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

"""Immutable value objects used by external data-source services."""

import hashlib
import json
from dataclasses import dataclass, field, replace
from types import MappingProxyType
from typing import Mapping, Optional


@dataclass(frozen=True)
class DataSourceCredential:
    username: str
    password: str = field(repr=False)
    options: Mapping[str, str] = field(default_factory=dict)

    def __post_init__(self):
        object.__setattr__(
            self,
            "options",
            MappingProxyType(dict(self.options)),
        )


@dataclass(frozen=True)
class SealedCredential:
    id: str
    data_source_uid: str
    credential_version: int
    encrypted_payload: bytes = field(repr=False)
    nonce: bytes = field(repr=False)
    key_version: str
    status: str = "active"


@dataclass(frozen=True)
class DataSourceDefinition:
    uid: Optional[str]
    name_en: str
    database_type: str
    host: str
    port: int
    database: str
    name_zh: Optional[str] = None
    schema: Optional[str] = None
    credential_ref: Optional[str] = None
    credential_version: Optional[int] = None
    pool_size: Optional[int] = None
    max_overflow: Optional[int] = None
    tls_options: Mapping[str, str] = field(default_factory=dict)
    status: bool = True
    description: Optional[str] = None
    extra_properties: Mapping[str, object] = field(default_factory=dict)

    def __post_init__(self):
        object.__setattr__(
            self,
            "database_type",
            str(self.database_type).strip().lower(),
        )
        object.__setattr__(self, "port", int(self.port))
        object.__setattr__(
            self,
            "tls_options",
            MappingProxyType(dict(self.tls_options)),
        )
        object.__setattr__(
            self,
            "extra_properties",
            MappingProxyType(dict(self.extra_properties)),
        )

    def with_uid(self, uid: str):
        return replace(self, uid=uid)

    def pool_overrides(self) -> dict:
        return {
            key: value
            for key, value in {
                "pool_size": self.pool_size,
                "max_overflow": self.max_overflow,
            }.items()
            if value is not None
        }

    def connection_fingerprint(self) -> str:
        payload = {
            "database_type": self.database_type,
            "host": self.host,
            "port": self.port,
            "database": self.database,
            "schema": self.schema,
            "pool_overrides": self.pool_overrides(),
            "tls_options": dict(self.tls_options),
        }
        canonical = json.dumps(
            payload,
            sort_keys=True,
            separators=(",", ":"),
        )
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


@dataclass(frozen=True)
class PoolKey:
    data_source_uid: str
    credential_version: int
    config_fingerprint: str


@dataclass(frozen=True)
class PoolStatus:
    data_source_uid: str
    credential_version: int
    config_fingerprint: str
    pool_state: str
    leases: int
    draining: bool
    created_at: float
    last_used_at: float
    pool_size: int = 0
    checked_out: int = 0
    checked_out_peak: int = 0
    checked_in: int = 0
    overflow: int = 0
    checkout_wait_ms: float = 0.0
    connection_created_total: int = 0
    connection_failed_total: int = 0
    pool_timeout_total: int = 0
    invalidated_total: int = 0
    query_total: int = 0
    last_query_duration_ms: float = 0.0
    consecutive_failures: int = 0
    circuit_open_until: float = 0.0