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

"""Public context-managed access to external business data sources."""

import time
from contextlib import contextmanager

from sqlalchemy.exc import DBAPIError, OperationalError, TimeoutError

from app.core.data_source.adapters.base import PURPOSES
from app.core.data_source.errors import (
    DataSourceConfigurationInvalid,
    DataSourceConnectionFailed,
    DataSourceNotFound,
    DataSourcePoolTimeout,
    DataSourceQueryTimeout,
    DataSourceReadOnlyViolation,
)
from app.core.data_source.models import PoolKey


def _driver_error_code(error):
    original = getattr(error, "orig", error)
    code = getattr(original, "pgcode", None)
    if code:
        return str(code)
    args = getattr(original, "args", ())
    return str(args[0]) if args else ""


def _translate_query_error(error):
    code = _driver_error_code(error)
    if code in {"57014", "3024", "1317"}:
        return DataSourceQueryTimeout()
    if code in {"25006", "1792"}:
        return DataSourceReadOnlyViolation()
    return error


class ManagedConnection:
    def __init__(self, connection, *, key, registry, clock=None):
        self._connection = connection
        self._key = key
        self._registry = registry
        self._clock = clock or time.monotonic

    def execute(self, statement, *args, **kwargs):
        started = self._clock()
        try:
            return self._connection.execute(statement, *args, **kwargs)
        except DBAPIError as exc:
            translated = _translate_query_error(exc)
            if translated is exc:
                raise
            raise translated from exc
        finally:
            self._registry.record_query_duration(
                self._key,
                (self._clock() - started) * 1000,
            )

    def __getattr__(self, name):
        return getattr(self._connection, name)


class DataSourceConnectionManager:
    def __init__(
        self,
        *,
        definitions,
        credentials,
        platform_session,
        adapter_resolver,
        registry,
        settings_resolver,
        clock=None,
    ):
        self.definitions = definitions
        self.credentials = credentials
        self.platform_session = platform_session
        self.adapter_resolver = adapter_resolver
        self.registry = registry
        self.settings_resolver = settings_resolver
        self._clock = clock or time.monotonic

    @contextmanager
    def connect(self, data_source_uid, purpose):
        if purpose not in PURPOSES or purpose == "connection_test":
            raise DataSourceConfigurationInvalid(
                "unsupported pooled connection purpose"
            )
        definition = self.definitions.get(data_source_uid)
        if not definition:
            raise DataSourceNotFound()
        credential = self.credentials.get_active(
            self.platform_session(),
            data_source_uid,
            definition.credential_version,
        )
        adapter = self.adapter_resolver(definition.database_type)
        settings = self.settings_resolver(definition.pool_overrides())
        key = PoolKey(
            data_source_uid=str(data_source_uid),
            credential_version=int(definition.credential_version),
            config_fingerprint=definition.connection_fingerprint(),
        )

        with self.registry.lease(
            key,
            lambda: adapter.create_pooled_engine(
                definition,
                credential,
                settings,
            ),
        ) as engine:
            breaker = self.registry.breaker_for(key)
            try:
                with breaker.probe_permission():
                    started = self._clock()
                    connection = engine.connect()
                    self.registry.add_checkout_wait(
                        key,
                        (self._clock() - started) * 1000,
                    )
                    self.registry.record_connection_success(key)
            except TimeoutError as exc:
                self.registry.record_pool_timeout(key)
                self.registry.record_connection_failure(key)
                raise DataSourcePoolTimeout() from exc
            except (OperationalError, DBAPIError) as exc:
                self.registry.record_connection_failure(key)
                raise DataSourceConnectionFailed() from exc

            with connection:
                transaction = connection.begin()
                try:
                    adapter.configure_transaction(connection, purpose)
                    yield ManagedConnection(
                        connection,
                        key=key,
                        registry=self.registry,
                        clock=self._clock,
                    )
                    if purpose == "dataflow_write":
                        transaction.commit()
                    else:
                        transaction.rollback()
                except Exception:
                    try:
                        transaction.rollback()
                    except Exception:
                        pass
                    raise

    def invalidate(self, data_source_uid, reason):
        self.registry.invalidate(data_source_uid, reason)

    def snapshot(self):
        return self.registry.snapshot()

    def close(self):
        self.registry.close_all()