DataOps-platform/app/core/data_source/runtime.py

"""Explicit and lazily constructed external data-source runtimes."""

import threading
from dataclasses import dataclass, field
from typing import Mapping

from neo4j import GraphDatabase
from sqlalchemy import create_engine


_lock = threading.Lock()
_manager = None
_factory = None


@dataclass(frozen=True)
class DataSourceRuntimeConfig:
    platform_database_url: str = field(repr=False)
    neo4j_uri: str
    neo4j_user: str
    neo4j_password: str = field(repr=False)
    credential_master_key: str = field(repr=False)
    credential_key_version: str
    certificate_dir: str = "/etc/dataops-platform/datasource-certs"
    pool_size: int = 1
    max_overflow: int = 1
    pool_timeout: int = 10
    pool_recycle: int = 1800
    idle_ttl: int = 900
    max_idle_pools: int = 4
    query_timeout: int = 30
    worker_count: int = 2
    connection_budget: int = 32

    @classmethod
    def from_mapping(cls, values: Mapping[str, object]):
        config = cls(**dict(values))
        if config.pool_size < 1 or config.max_overflow < 0:
            raise ValueError("data source pool size is invalid")
        if config.max_idle_pools < 1 or config.worker_count < 1:
            raise ValueError("data source worker configuration is invalid")
        if config.max_business_connections > config.connection_budget:
            raise ValueError("runner data source connection budget is exceeded")
        return config

    @property
    def max_business_connections(self):
        return (
            self.worker_count
            * self.max_idle_pools
            * (self.pool_size + self.max_overflow)
        )

    def pool_settings(self):
        return {
            "pool_size": self.pool_size,
            "max_overflow": self.max_overflow,
            "pool_timeout": self.pool_timeout,
            "pool_recycle": self.pool_recycle,
            "idle_ttl": self.idle_ttl,
            "max_idle_pools": self.max_idle_pools,
            "query_timeout": self.query_timeout,
            "drain_timeout": 30,
        }


class _Neo4jDefinitionGateway:
    def __init__(self, driver=None):
        self._driver = driver

    def _call(self, method, *args, **kwargs):
        from app.core.data_source.definitions import (
            DataSourceDefinitionRepository,
        )

        if self._driver is None:
            from app.services.neo4j_driver import neo4j_driver

            session_context = neo4j_driver.get_session()
        else:
            session_context = self._driver.session()
        with session_context as session:
            repository = DataSourceDefinitionRepository(session)
            return getattr(repository, method)(*args, **kwargs)

    def get(self, *args, **kwargs):
        return self._call("get", *args, **kwargs)

    def list(self, *args, **kwargs):
        return self._call("list", *args, **kwargs)

    def save(self, *args, **kwargs):
        return self._call("save", *args, **kwargs)

    def delete(self, *args, **kwargs):
        return self._call("delete", *args, **kwargs)

    def credential_references(self, *args, **kwargs):
        return self._call("credential_references", *args, **kwargs)

    def remove_legacy_secret_fields(self, *args, **kwargs):
        return self._call(
            "remove_legacy_secret_fields",
            *args,
            **kwargs,
        )


class _StandaloneCredentialGateway:
    def __init__(self, engine, repository):
        self._engine = engine
        self._repository = repository

    def get_active(self, _session, uid, version):
        with self._engine.connect() as connection:
            return self._repository.get_active(connection, uid, version)


@dataclass
class StandaloneDataSourceRuntime:
    manager: object
    platform_engine: object
    graph_driver: object

    def close(self):
        self.manager.close()
        self.graph_driver.close()
        self.platform_engine.dispose()


def build_standalone_data_source_runtime(config):
    """Build a Runner-owned runtime without Flask globals or request state."""
    from app.core.data_source.adapters import adapter_for
    from app.core.data_source.credentials import (
        CredentialCodec,
        DataSourceCredentialRepository,
    )
    from app.core.data_source.manager import DataSourceConnectionManager
    from app.core.data_source.pool_registry import PoolRegistry

    platform_engine = create_engine(
        config.platform_database_url,
        pool_size=1,
        max_overflow=1,
        pool_pre_ping=True,
        pool_recycle=300,
    )
    graph_driver = GraphDatabase.driver(
        config.neo4j_uri,
        auth=(config.neo4j_user, config.neo4j_password),
    )
    codec = CredentialCodec.from_base64(
        config.credential_master_key,
        config.credential_key_version,
    )
    base_settings = config.pool_settings()

    def settings_resolver(overrides):
        values = dict(base_settings)
        for key in ("pool_size", "max_overflow"):
            if key in overrides:
                values[key] = int(overrides[key])
        if values["pool_size"] + values["max_overflow"] > (
            config.pool_size + config.max_overflow
        ):
            raise ValueError("data source override exceeds runner connection budget")
        return values

    manager = DataSourceConnectionManager(
        definitions=_Neo4jDefinitionGateway(graph_driver),
        credentials=_StandaloneCredentialGateway(
            platform_engine,
            DataSourceCredentialRepository(codec),
        ),
        platform_session=lambda: None,
        adapter_resolver=lambda database_type: adapter_for(
            database_type,
            certificate_dir=config.certificate_dir,
        ),
        registry=PoolRegistry(base_settings),
        settings_resolver=settings_resolver,
    )
    return StandaloneDataSourceRuntime(
        manager=manager,
        platform_engine=platform_engine,
        graph_driver=graph_driver,
    )


def _build_default_manager():
    from flask import current_app

    from app import db
    from app.config.config import datasource_pool_settings
    from app.core.data_source.adapters import adapter_for
    from app.core.data_source.credentials import (
        CredentialCodec,
        DataSourceCredentialRepository,
    )
    from app.core.data_source.manager import DataSourceConnectionManager
    from app.core.data_source.pool_registry import PoolRegistry

    codec = CredentialCodec.from_base64(
        current_app.config.get("DATASOURCE_CREDENTIAL_MASTER_KEY", ""),
        current_app.config.get("DATASOURCE_CREDENTIAL_KEY_VERSION", "v1"),
    )
    base_settings = datasource_pool_settings()
    registry_settings = {
        **base_settings,
        "drain_timeout": 30,
    }
    certificate_dir = current_app.config.get("DATASOURCE_CERT_DIR")
    return DataSourceConnectionManager(
        definitions=_Neo4jDefinitionGateway(),
        credentials=DataSourceCredentialRepository(codec),
        platform_session=lambda: db.session,
        adapter_resolver=lambda database_type: adapter_for(
            database_type,
            certificate_dir=certificate_dir,
        ),
        registry=PoolRegistry(registry_settings),
        settings_resolver=datasource_pool_settings,
    )


def configure_data_source_runtime(factory):
    """Replace the process-local factory and clear any current manager."""
    global _factory, _manager
    with _lock:
        if _manager is not None:
            _manager.close()
        _manager = None
        _factory = factory


def get_data_source_manager():
    global _manager
    if _manager is not None:
        return _manager
    with _lock:
        if _manager is None:
            _manager = (_factory or _build_default_manager)()
        return _manager


def peek_data_source_manager():
    """Return the current Worker manager without constructing one."""
    return _manager


def close_data_source_runtime():
    global _manager
    with _lock:
        manager = _manager
        _manager = None
    if manager is not None:
        manager.close()