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

"""Secure cross-store lifecycle orchestration for external data sources."""

from dataclasses import replace

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.definitions import DataSourceFilters
from app.core.data_source.errors import (
    DataSourceConfigurationInvalid,
    DataSourceError,
    DataSourceNotFound,
)
from app.core.data_source.models import (
    DataSourceCredential,
    DataSourceDefinition,
)


TYPE_ALIASES = {
    "postgres": "postgresql",
    "postgresql": "postgresql",
    "mysql": "mysql",
}
POOL_INVALIDATION_REASONS = {
    "admin_reset",
    "configuration_changed",
    "credential_rotated",
}


class DataSourceService:
    def __init__(
        self,
        *,
        definitions,
        credentials,
        platform_session,
        adapter_resolver,
        connection_manager,
        settings_resolver,
        outbox_enqueuer,
    ):
        self.definitions = definitions
        self.credentials = credentials
        self.platform_session = platform_session
        self.adapter_resolver = adapter_resolver
        self.connection_manager = connection_manager
        self.settings_resolver = settings_resolver
        self.outbox_enqueuer = outbox_enqueuer

    @staticmethod
    def _database_type(payload):
        requested = str(payload.get("type") or "").strip().lower()
        database_type = TYPE_ALIASES.get(requested)
        if database_type is None:
            raise DataSourceConfigurationInvalid(
                "only PostgreSQL and MySQL data sources are supported"
            )
        return database_type

    @staticmethod
    def _required_text(payload, name):
        value = str(payload.get(name) or "").strip()
        if not value:
            raise DataSourceConfigurationInvalid(f"{name} is required")
        return value

    @classmethod
    def _definition(
        cls,
        payload,
        *,
        uid,
        credential_version,
        existing=None,
    ):
        database_type = cls._database_type(payload)
        try:
            port = int(payload.get("port"))
        except (TypeError, ValueError) as exc:
            raise DataSourceConfigurationInvalid("port is invalid") from exc
        if port < 1 or port > 65535:
            raise DataSourceConfigurationInvalid("port is invalid")

        tls_options = payload.get("tls_options") or {}
        if not isinstance(tls_options, dict):
            raise DataSourceConfigurationInvalid("tls_options is invalid")
        return DataSourceDefinition(
            uid=uid,
            name_en=cls._required_text(payload, "name_en"),
            name_zh=payload.get("name_zh"),
            database_type=database_type,
            host=cls._required_text(payload, "host"),
            port=port,
            database=cls._required_text(payload, "database"),
            schema=payload.get("schema"),
            credential_ref=uid,
            credential_version=credential_version,
            pool_size=payload.get("pool_size"),
            max_overflow=payload.get("max_overflow"),
            tls_options=tls_options,
            status=bool(payload.get("status", True)),
            description=payload.get("desc"),
            extra_properties=(
                dict(existing.extra_properties) if existing else {}
            ),
        )

    def _effective_credential(self, payload, existing):
        username = payload.get("username")
        password = payload.get("password")
        current = None
        if existing is not None and (not username or not password):
            current = self.credentials.get_active(
                self.platform_session(),
                existing.uid,
                existing.credential_version,
            )
        username = str(username or (current.username if current else "")).strip()
        password = str(password or (current.password if current else ""))
        if not username or not password:
            raise DataSourceConfigurationInvalid(
                "username and password are required"
            )
        options = payload.get("credential_options") or (
            dict(current.options) if current else {}
        )
        if not isinstance(options, dict):
            raise DataSourceConfigurationInvalid(
                "credential options are invalid"
            )
        return DataSourceCredential(username, password, options)

    def save(self, payload, *, actor_uid):
        if not isinstance(payload, dict):
            raise DataSourceConfigurationInvalid()
        requested_uid = str(payload.get("uid") or "").strip() or None
        existing = (
            self.definitions.get(requested_uid)
            if requested_uid is not None
            else None
        )
        if requested_uid is not None and existing is None:
            raise DataSourceNotFound()
        uid = requested_uid or new_governance_uid()
        credential = self._effective_credential(payload, existing)

        candidate = self._definition(
            payload,
            uid=uid,
            credential_version=(
                existing.credential_version if existing else 1
            ),
            existing=existing,
        )
        adapter = self.adapter_resolver(candidate.database_type)
        adapter.validate_options(candidate.tls_options)
        settings = self.settings_resolver(candidate.pool_overrides())
        adapter.test_connection(
            candidate,
            credential,
            query_timeout=settings["query_timeout"],
        )

        session = self.platform_session()
        previous_version = (
            existing.credential_version if existing is not None else None
        )
        try:
            sealed = self.credentials.create_version(
                session,
                data_source_uid=uid,
                credential=credential,
                actor_uid=actor_uid,
            )
            self.outbox_enqueuer(
                session,
                aggregate_type="datasource",
                aggregate_id=uid,
                event_type="datasource.credential_version_created",
                payload={
                    "uid": uid,
                    "database_type": candidate.database_type,
                    "credential_version": sealed.credential_version,
                },
            )
            session.commit()
        except Exception:
            session.rollback()
            raise

        candidate = replace(
            candidate,
            credential_version=sealed.credential_version,
        )
        try:
            saved = self.definitions.save(candidate)
        except Exception as graph_error:
            try:
                self.credentials.compensate_failed_activation(
                    session,
                    data_source_uid=uid,
                    failed_version=sealed.credential_version,
                    restore_version=previous_version,
                    actor_uid=actor_uid,
                )
                session.commit()
            except Exception as compensation_error:
                session.rollback()
                raise DataSourceError(
                    "data source save requires reconciliation"
                ) from compensation_error
            raise DataSourceError(
                "data source definition could not be saved"
            ) from graph_error

        if existing is not None:
            self.connection_manager.invalidate(
                uid,
                "configuration_changed",
            )
        return saved, existing is None

    def list(self, payload):
        payload = payload if isinstance(payload, dict) else {}
        filters = DataSourceFilters(
            uid=payload.get("uid"),
            name_en=payload.get("name_en"),
            name_zh=payload.get("name_zh"),
            database_type=payload.get("type"),
            status=payload.get("status"),
        )
        return self.definitions.list(filters)

    def test_connection(self, payload):
        if not isinstance(payload, dict):
            raise DataSourceConfigurationInvalid()
        uid = str(payload.get("uid") or "").strip() or None
        existing = self.definitions.get(uid) if uid else None
        if uid and existing is None:
            raise DataSourceNotFound()
        effective_uid = uid or new_governance_uid()
        credential = self._effective_credential(payload, existing)
        candidate = self._definition(
            payload,
            uid=effective_uid,
            credential_version=(
                existing.credential_version if existing else 1
            ),
            existing=existing,
        )
        adapter = self.adapter_resolver(candidate.database_type)
        settings = self.settings_resolver(candidate.pool_overrides())
        adapter.test_connection(
            candidate,
            credential,
            query_timeout=settings["query_timeout"],
        )
        return {"connected": True, "message": "连接测试成功"}

    def delete(self, uid, *, actor_uid):
        uid = str(uid or "").strip()
        existing = self.definitions.get(uid)
        if existing is None:
            raise DataSourceNotFound()
        self.connection_manager.invalidate(uid, "datasource_deleted")
        if not self.definitions.delete(uid):
            raise DataSourceNotFound()

        session = self.platform_session()
        try:
            revoked = self.credentials.revoke_all(
                session,
                data_source_uid=uid,
                actor_uid=actor_uid,
            )
            self.outbox_enqueuer(
                session,
                aggregate_type="datasource",
                aggregate_id=uid,
                event_type="datasource.deleted",
                payload={"uid": uid},
            )
            session.commit()
        except Exception as exc:
            session.rollback()
            raise DataSourceError(
                "data source deletion requires reconciliation"
            ) from exc
        return {"uid": uid, "revoked_credential_versions": revoked}

    def pool_statuses(self):
        return self.connection_manager.snapshot()

    def pool_status(self, uid):
        uid = str(uid or "").strip()
        for status in self.pool_statuses():
            if status.data_source_uid == uid:
                return status
        raise DataSourceNotFound()

    def invalidate_pool(self, uid, *, reason, actor_uid):
        uid = str(uid or "").strip()
        if reason not in POOL_INVALIDATION_REASONS:
            raise DataSourceConfigurationInvalid(
                "pool invalidation reason is invalid"
            )
        if self.definitions.get(uid) is None:
            raise DataSourceNotFound()
        self.connection_manager.invalidate(uid, reason)
        session = self.platform_session()
        try:
            self.credentials.record_pool_event(
                session,
                data_source_uid=uid,
                actor_uid=actor_uid,
                event_type="pool_invalidated",
                safe_detail=f"pool invalidated: {reason}",
            )
            session.commit()
        except Exception:
            session.rollback()
            raise
        return {"uid": uid, "invalidated": True, "reason": reason}

    @staticmethod
    def serialize(definition):
        return {
            "uid": definition.uid,
            "name_en": definition.name_en,
            "name_zh": definition.name_zh,
            "type": definition.database_type,
            "host": definition.host,
            "port": definition.port,
            "database": definition.database,
            "schema": definition.schema,
            "pool_size": definition.pool_size,
            "max_overflow": definition.max_overflow,
            "status": definition.status,
            "desc": definition.description,
            "credential_configured": bool(
                definition.credential_version
            ),
        }

    @staticmethod
    def serialize_pool_status(status):
        return {
            "data_source_uid": status.data_source_uid,
            "credential_version": status.credential_version,
            "pool_state": status.pool_state,
            "pool_size": status.pool_size,
            "checked_out": status.checked_out,
            "checked_out_peak": status.checked_out_peak,
            "checked_in": status.checked_in,
            "overflow": status.overflow,
            "leases": status.leases,
            "checkout_wait_ms": status.checkout_wait_ms,
            "last_used_at": status.last_used_at,
            "connection_created_total": status.connection_created_total,
            "connection_failed_total": status.connection_failed_total,
            "pool_timeout_total": status.pool_timeout_total,
            "invalidated_total": status.invalidated_total,
            "query_total": status.query_total,
            "last_query_duration_ms": status.last_query_duration_ms,
            "consecutive_failures": status.consecutive_failures,
            "circuit_open_until": status.circuit_open_until,
        }


def build_data_source_service(manager):
    from app.core.events.outbox import enqueue_outbox

    return DataSourceService(
        definitions=manager.definitions,
        credentials=manager.credentials,
        platform_session=manager.platform_session,
        adapter_resolver=manager.adapter_resolver,
        connection_manager=manager,
        settings_resolver=manager.settings_resolver,
        outbox_enqueuer=enqueue_outbox,
    )