DataOps-platform/app/runner/bootstrap.py

"""Environment parsing and dependency wiring for the standalone Runner."""

from __future__ import annotations

import os
from dataclasses import dataclass, field

import click
from minio import Minio

from app.core.data_source.runtime import (
    DataSourceRuntimeConfig,
    build_standalone_data_source_runtime,
)
from app.runner.api import create_runner_app
from app.runner.artifacts import ArtifactStore, PostgresArtifactResolver
from app.runner.auth import TaskTokenVerifier
from app.runner.ledger import PostgresTaskLedger
from app.runner.nodes import (
    GovernedHttpExecutor,
    NodeRegistry,
    RestrictedPythonExecutor,
    SqlExecuteExecutor,
    SqlQueryExecutor,
)
from app.runner.rule_polars import PolarsRulePlanAdapter
from app.runner.rule_sql import (
    SqlGlotQualityPlanAdapter,
    SqlGlotRulePlanAdapter,
)
from app.runner.rules import (
    PostgresRulePlanRepository,
    RulePlanExecutor,
)


def _required(name):
    value = str(os.environ.get(name, "")).strip()
    if not value:
        raise ValueError(f"{name} is required")
    return value


def _integer(name, default, minimum, maximum):
    try:
        value = int(os.environ.get(name, str(default)))
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{name} must be an integer") from exc
    if value < minimum or value > maximum:
        raise ValueError(f"{name} must be between {minimum} and {maximum}")
    return value


def _boolean(name, default=False):
    value = str(os.environ.get(name, str(default))).strip().lower()
    if value in {"1", "true", "yes", "on"}:
        return True
    if value in {"0", "false", "no", "off"}:
        return False
    raise ValueError(f"{name} must be a boolean")


@dataclass(frozen=True)
class RunnerSettings:
    runtime: DataSourceRuntimeConfig = field(repr=False)
    task_token_secret: str = field(repr=False)
    artifact_host: str = field(repr=False)
    artifact_user: str = field(repr=False)
    artifact_password: str = field(repr=False)
    artifact_bucket: str = "dataops-rules"
    artifact_secure: bool = False
    artifact_max_bytes: int = 32 * 1024 * 1024
    artifact_max_rows: int = 100_000
    artifact_memory_limit_bytes: int = 256 * 1024 * 1024
    artifact_ttl_seconds: int = 3600
    allowed_http_hosts: frozenset = field(default_factory=frozenset)
    task_token_ttl_seconds: int = 60
    max_query_rows: int = 1000


def register_artifact_cleanup_cli(app, artifact_resolver):
    """Register the bounded production maintenance entrypoint."""

    @app.cli.command("reconcile-rule-artifacts")
    @click.option("--limit", type=click.IntRange(1, 1_000), default=100)
    @click.option(
        "--grace-seconds",
        type=click.IntRange(30, 86_400),
        default=300,
    )
    def reconcile_rule_artifacts(limit, grace_seconds):
        result = artifact_resolver.reconcile(
            limit=limit,
            grace_seconds=grace_seconds,
        )
        click.echo(
            " ".join(
                f"{key}={result[key]}"
                for key in (
                    "pending_finalized",
                    "pending_deleted",
                    "ready_failed",
                    "orphans_deleted",
                )
            )
        )

    @app.cli.command("cleanup-rule-artifacts")
    @click.option("--limit", type=click.IntRange(1, 1_000), default=100)
    def cleanup_rule_artifacts(limit):
        removed = artifact_resolver.cleanup_expired(limit=limit)
        click.echo(f"removed {removed} expired rule artifacts")


def runner_settings_from_env():
    task_token_secret = _required("RUNNER_TASK_TOKEN_SECRET")
    worker_count = _integer("RUNNER_WORKERS", 2, 1, 8)
    runtime = DataSourceRuntimeConfig.from_mapping(
        {
            "platform_database_url": _required("DATABASE_URL"),
            "neo4j_uri": _required("NEO4J_URI"),
            "neo4j_user": _required("NEO4J_USER"),
            "neo4j_password": _required("NEO4J_PASSWORD"),
            "credential_master_key": _required(
                "DATASOURCE_CREDENTIAL_MASTER_KEY"
            ),
            "credential_key_version": _required(
                "DATASOURCE_CREDENTIAL_KEY_VERSION"
            ),
            "certificate_dir": os.environ.get(
                "DATASOURCE_CERT_DIR",
                "/etc/dataops-platform/datasource-certs",
            ),
            "pool_size": _integer("RUNNER_DATASOURCE_POOL_SIZE", 1, 1, 3),
            "max_overflow": _integer(
                "RUNNER_DATASOURCE_MAX_OVERFLOW", 1, 0, 3
            ),
            "pool_timeout": _integer(
                "RUNNER_DATASOURCE_POOL_TIMEOUT", 10, 1, 60
            ),
            "pool_recycle": _integer(
                "RUNNER_DATASOURCE_POOL_RECYCLE", 1800, 60, 86400
            ),
            "idle_ttl": _integer(
                "RUNNER_DATASOURCE_POOL_IDLE_TTL", 900, 60, 86400
            ),
            "max_idle_pools": _integer(
                "RUNNER_DATASOURCE_MAX_IDLE_POOLS", 4, 1, 20
            ),
            "query_timeout": _integer(
                "RUNNER_DATASOURCE_QUERY_TIMEOUT", 30, 1, 300
            ),
            "worker_count": worker_count,
            "connection_budget": _integer(
                "RUNNER_DATASOURCE_CONNECTION_BUDGET", 32, 1, 200
            ),
        }
    )
    allowed_hosts = frozenset(
        host.strip().lower()
        for host in os.environ.get("RUNNER_HTTP_ALLOWED_HOSTS", "").split(",")
        if host.strip()
    )
    return RunnerSettings(
        runtime=runtime,
        task_token_secret=task_token_secret,
        artifact_host=_required("RUNNER_MINIO_HOST"),
        artifact_user=_required("RUNNER_MINIO_USER"),
        artifact_password=_required("RUNNER_MINIO_PASSWORD"),
        artifact_bucket=_required("RUNNER_MINIO_BUCKET"),
        artifact_secure=_boolean("RUNNER_MINIO_SECURE", False),
        artifact_max_bytes=_integer(
            "RUNNER_ARTIFACT_MAX_BYTES",
            32 * 1024 * 1024,
            1024,
            2 * 1024 * 1024 * 1024,
        ),
        artifact_max_rows=_integer(
            "RUNNER_ARTIFACT_MAX_ROWS", 100_000, 1, 10_000_000
        ),
        artifact_memory_limit_bytes=_integer(
            "RUNNER_ARTIFACT_MEMORY_LIMIT_BYTES",
            256 * 1024 * 1024,
            16 * 1024 * 1024,
            16 * 1024 * 1024 * 1024,
        ),
        artifact_ttl_seconds=_integer(
            "RUNNER_ARTIFACT_TTL_SECONDS", 3600, 1, 86400
        ),
        allowed_http_hosts=allowed_hosts,
        task_token_ttl_seconds=_integer(
            "RUNNER_TASK_TOKEN_TTL_SECONDS", 60, 1, 300
        ),
        max_query_rows=_integer("RUNNER_MAX_QUERY_ROWS", 1000, 1, 10000),
    )


def build_runner_application(settings=None):
    settings = settings or runner_settings_from_env()
    runtime = build_standalone_data_source_runtime(settings.runtime)
    query_executor = SqlQueryExecutor(
        runtime.manager,
        max_rows=settings.max_query_rows,
    )
    write_executor = SqlExecuteExecutor(runtime.manager)
    sql_rule_adapter = SqlGlotRulePlanAdapter(runtime.manager)
    artifact_store = ArtifactStore(
        Minio(
            settings.artifact_host,
            access_key=settings.artifact_user,
            secret_key=settings.artifact_password,
            secure=settings.artifact_secure,
        ),
        bucket=settings.artifact_bucket,
        max_artifact_bytes=settings.artifact_max_bytes,
        max_rows=settings.artifact_max_rows,
        memory_limit_bytes=settings.artifact_memory_limit_bytes,
        max_ttl_seconds=settings.artifact_ttl_seconds,
    )
    artifact_resolver = PostgresArtifactResolver(
        runtime.platform_engine, artifact_store
    )
    polars_rule_adapter = PolarsRulePlanAdapter(
        artifact_store=artifact_store,
        artifact_resolver=artifact_resolver,
        masking_policies={
            "customer_mobile_last4": "preserve_last_4",
            "redact": "redact",
        },
        artifact_ttl_seconds=settings.artifact_ttl_seconds,
    )
    rule_executor = RulePlanExecutor(
        PostgresRulePlanRepository(runtime.platform_engine),
        adapters={
            "sql_pushdown": sql_rule_adapter,
            "polars_batch": polars_rule_adapter,
            "quality_check": SqlGlotQualityPlanAdapter(),
        },
    )
    registry = NodeRegistry(
        {
            "sql.query": query_executor,
            "sql.execute": write_executor,
            "python": RestrictedPythonExecutor({}),
            "http": GovernedHttpExecutor(
                allowed_hosts=settings.allowed_http_hosts
            ),
            "rule.apply": rule_executor,
            "quality.check": rule_executor,
        }
    )
    application = create_runner_app(
        verifier=TaskTokenVerifier(settings.task_token_secret),
        ledger=PostgresTaskLedger(runtime.platform_engine),
        registry=registry,
    )
    register_artifact_cleanup_cli(application, artifact_resolver)
    application.extensions["dataops_runner_runtime"] = runtime
    application.extensions["dataops_runner_settings"] = settings
    return application