"""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()