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