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

"""Secret-free Neo4j repository for external data-source definitions."""

import json
from dataclasses import dataclass
from typing import Optional

from app.core.common.identifiers import ensure_governance_uid
from app.core.data_source.models import DataSourceDefinition


NEO4J_SECRET_KEYS = {
    "username",
    "password",
    "passwd",
    "credential",
    "credentials",
    "encrypted_payload",
    "nonce",
    "api_key",
    "token",
    "authorization",
    "conn_str",
    "connection_string",
    "connection_url",
}


@dataclass(frozen=True)
class DataSourceFilters:
    uid: Optional[str] = None
    name_en: Optional[str] = None
    name_zh: Optional[str] = None
    database_type: Optional[str] = None
    status: Optional[bool] = None


class DataSourceDefinitionRepository:
    """Read and write DataSource nodes using stable UIDs only."""

    def __init__(self, session):
        self.session = session

    @staticmethod
    def _assert_secret_free(properties: dict) -> None:
        def visit(value):
            if isinstance(value, dict):
                for key, item in value.items():
                    if str(key).strip().lower() in NEO4J_SECRET_KEYS:
                        raise ValueError(
                            "secret properties cannot be stored in DataSource"
                        )
                    visit(item)
            elif isinstance(value, (list, tuple)):
                for item in value:
                    visit(item)

        visit(properties)

    @classmethod
    def _to_properties(cls, definition: DataSourceDefinition) -> dict:
        properties = {
            "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,
            "credential_ref": definition.credential_ref,
            "credential_version": definition.credential_version,
            "pool_size": definition.pool_size,
            "max_overflow": definition.max_overflow,
            "tls_options_json": json.dumps(
                dict(definition.tls_options),
                sort_keys=True,
                separators=(",", ":"),
            ),
            "status": definition.status,
            "desc": definition.description,
        }
        properties.update(dict(definition.extra_properties))
        properties = {
            key: value for key, value in properties.items() if value is not None
        }
        cls._assert_secret_free(properties)
        return properties

    @staticmethod
    def _from_properties(properties: dict) -> DataSourceDefinition:
        values = dict(properties)
        tls_raw = values.pop("tls_options_json", "{}") or "{}"
        try:
            tls_options = json.loads(tls_raw)
        except (TypeError, ValueError):
            tls_options = {}
        known = {
            "uid",
            "name_en",
            "name_zh",
            "type",
            "host",
            "port",
            "database",
            "schema",
            "credential_ref",
            "credential_version",
            "pool_size",
            "max_overflow",
            "status",
            "desc",
        }
        return DataSourceDefinition(
            uid=values.get("uid"),
            name_en=values.get("name_en", ""),
            name_zh=values.get("name_zh"),
            database_type=values.get("type", ""),
            host=values.get("host", ""),
            port=values.get("port", 0),
            database=values.get("database", ""),
            schema=values.get("schema"),
            credential_ref=values.get("credential_ref"),
            credential_version=values.get("credential_version"),
            pool_size=values.get("pool_size"),
            max_overflow=values.get("max_overflow"),
            tls_options=tls_options if isinstance(tls_options, dict) else {},
            status=bool(values.get("status", True)),
            description=values.get("desc"),
            extra_properties={
                key: value
                for key, value in values.items()
                if key not in known and key != "_id"
            },
        )

    def get(self, uid: str) -> Optional[DataSourceDefinition]:
        record = self.session.run(
            """
            MATCH (n:DataSource {uid: $uid})
            RETURN properties(n) AS properties
            """,
            {"uid": str(uid)},
        ).single()
        if record is None:
            return None
        return self._from_properties(dict(record["properties"]))

    def list(self, filters: DataSourceFilters = None):
        filters = filters or DataSourceFilters()
        clauses = []
        parameters = {}
        mapping = {
            "uid": "n.uid",
            "name_en": "n.name_en",
            "name_zh": "n.name_zh",
            "database_type": "n.type",
            "status": "n.status",
        }
        for field_name, property_name in mapping.items():
            value = getattr(filters, field_name)
            if value is not None:
                clauses.append(f"{property_name} = ${field_name}")
                parameters[field_name] = value
        where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        records = self.session.run(
            f"""
            MATCH (n:DataSource)
            {where}
            RETURN properties(n) AS properties
            ORDER BY n.name_en
            """,
            parameters,
        )
        return [
            self._from_properties(dict(record["properties"]))
            for record in records
        ]

    def save(self, definition: DataSourceDefinition) -> DataSourceDefinition:
        uid_holder = {"uid": definition.uid} if definition.uid else {}
        uid = ensure_governance_uid(uid_holder)
        saved = definition.with_uid(uid)
        properties = self._to_properties(saved)
        self.session.run(
            """
            MERGE (n:DataSource {uid: $uid})
            SET n = $properties
            RETURN properties(n) AS properties
            """,
            {"uid": uid, "properties": properties},
        )
        return saved

    def delete(self, uid: str) -> bool:
        record = self.session.run(
            """
            MATCH (n:DataSource {uid: $uid})
            WITH n, count(n) AS found
            DETACH DELETE n
            RETURN found AS deleted_count
            """,
            {"uid": str(uid)},
        ).single()
        return bool(record and record["deleted_count"])

    def credential_references(self):
        records = self.session.run(
            """
            MATCH (n:DataSource)
            WHERE n.uid IS NOT NULL
            RETURN n.uid AS uid, n.credential_ref AS credential_ref,
                   n.credential_version AS credential_version
            """
        )
        return {
            str(record["uid"]): (
                record["credential_ref"],
                int(record["credential_version"]),
            )
            for record in records
            if record.get("credential_ref") is not None
            and record.get("credential_version") is not None
        }

    def remove_legacy_secret_fields(self, uid: str) -> None:
        self.session.run(
            """
            MATCH (n:DataSource {uid: $uid})
            REMOVE n.username, n.password, n.conn_str,
                   n.connection_string, n.connection_url
            """,
            {"uid": str(uid)},
        )