"""One lazily constructed data-source manager per backend Worker process.""" import threading _lock = threading.Lock() _manager = None _factory = None class _Neo4jDefinitionGateway: def _call(self, method, *args, **kwargs): from app.core.data_source.definitions import ( DataSourceDefinitionRepository, ) from app.services.neo4j_driver import neo4j_driver with neo4j_driver.get_session() 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, ) 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()