DataOps-platform/app/core/data_research/sources.py

from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Any, Callable

from app.core.common.timezone_utils import now_china_naive
from app.core.data_research.errors import IngestionSourceInvalid


@dataclass(frozen=True)
class IngestionSourceRecord:
    uid: str
    source_type: str
    name: str
    config: dict[str, Any]
    permission_scope: dict[str, Any]
    status: str
    created_by: str | None
    created_at: Any = None
    updated_at: Any = None


class DatabaseSourceRegistrationService:
    SUPPORTED_DATABASES = frozenset({"postgresql", "mysql"})

    def __init__(
        self,
        repository,
        *,
        definition_resolver: Callable[[str], Any],
        clock: Callable[[], Any] = now_china_naive,
        commit: Callable[[], Any] = lambda: None,
        rollback: Callable[[], Any] = lambda: None,
    ):
        self.repository = repository
        self.definition_resolver = definition_resolver
        self.clock = clock
        self.commit = commit
        self.rollback = rollback

    def ensure(self, data_source_uid, actor_uid):
        uid = str(data_source_uid or "").strip()
        definition = self.definition_resolver(uid) if uid else None
        if definition is None:
            raise IngestionSourceInvalid("data source was not found")
        if not bool(definition.status):
            raise IngestionSourceInvalid("data source is disabled")
        database_type = str(definition.database_type or "").strip().lower()
        if database_type not in self.SUPPORTED_DATABASES:
            raise IngestionSourceInvalid(
                f"database type {database_type or 'unknown'} is not supported"
            )

        now = self.clock()
        existing = self.repository.get(uid)
        if existing is not None and existing.source_type != "database":
            raise IngestionSourceInvalid(
                "data source UID is already used by another source type"
            )
        config = {
            "database_type": database_type,
            "database": str(definition.database or ""),
            "schema": str(definition.schema or "") or None,
        }
        name = str(
            definition.name_zh or definition.name_en or uid
        ).strip()
        if existing is None:
            record = IngestionSourceRecord(
                uid=uid,
                source_type="database",
                name=name,
                config=config,
                permission_scope={},
                status="active",
                created_by=actor_uid,
                created_at=now,
                updated_at=now,
            )
            created = True
        else:
            record = replace(
                existing,
                name=name,
                config=config,
                status="active",
                updated_at=now,
            )
            created = False
        try:
            saved = self.repository.save(record)
            self.commit()
            return saved, created
        except Exception:
            self.rollback()
            raise