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

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