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

from app.core.data_source.runtime import (
    DataSourceRuntimeConfig,
    build_standalone_data_source_runtime,
)
from app.runner.api import create_runner_app
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_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


@dataclass(frozen=True)
class RunnerSettings:
    runtime: DataSourceRuntimeConfig = field(repr=False)
    task_token_secret: str = field(repr=False)
    allowed_http_hosts: frozenset = field(default_factory=frozenset)
    task_token_ttl_seconds: int = 60
    max_query_rows: int = 1000


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,
        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)
    rule_executor = RulePlanExecutor(
        PostgresRulePlanRepository(runtime.platform_engine),
        adapters={
            "sql_pushdown": sql_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,
    )
    application.extensions["dataops_runner_runtime"] = runtime
    application.extensions["dataops_runner_settings"] = settings
    return application