DataOps-platform/app/commands/reconcile_governance_uids.py

from __future__ import annotations

import argparse
import json
from typing import Any

from app.core.common.identifiers import new_governance_uid


GOVERNANCE_LABELS = ("BusinessDomain", "DataMeta", "DataFlow", "DataSource")


def ensure_uid_constraints(session: Any) -> None:
    constraints = {
        "BusinessDomain": "business_domain_uid",
        "DataMeta": "data_meta_uid",
        "DataFlow": "data_flow_uid",
        "DataSource": "datasource_uid",
    }
    for label, constraint in constraints.items():
        session.run(
            f"CREATE CONSTRAINT {constraint} IF NOT EXISTS "
            f"FOR (n:{label}) REQUIRE n.uid IS UNIQUE"
        )


def report_missing_uids(session: Any, limit: int = 100) -> dict[str, list[dict[str, Any]]]:
    report: dict[str, list[dict[str, Any]]] = {}
    for label in GOVERNANCE_LABELS:
        result = session.run(
            f"MATCH (n:{label}) WHERE n.uid IS NULL "
            "RETURN elementId(n) AS legacy_id, n.name_zh AS name_zh "
            "ORDER BY legacy_id LIMIT $limit",
            {"limit": int(limit)},
        )
        report[label] = [dict(record) for record in result]
    return report


def backfill_missing_uids(session: Any, limit: int = 100) -> dict[str, int]:
    report = report_missing_uids(session, limit=limit)
    counts: dict[str, int] = {}
    for label, records in report.items():
        counts[label] = 0
        for record in records:
            session.run(
                f"MATCH (n:{label}) WHERE elementId(n) = $legacy_id AND n.uid IS NULL "
                "SET n.uid = $uid",
                {
                    "legacy_id": str(record["legacy_id"]),
                    "uid": new_governance_uid(),
                },
            )
            counts[label] += 1
    return counts


def main() -> None:
    parser = argparse.ArgumentParser(description="Audit stable governance UIDs")
    parser.add_argument("--limit", type=int, default=100)
    parser.add_argument("--backfill", action="store_true")
    parser.add_argument("--ensure-constraints", action="store_true")
    args = parser.parse_args()

    from app.services.neo4j_driver import neo4j_driver

    with neo4j_driver.get_session() as session:
        report = report_missing_uids(session, limit=args.limit)
        print(json.dumps(report, ensure_ascii=False, indent=2))
        if args.backfill:
            counts = backfill_missing_uids(session, limit=args.limit)
            print(json.dumps({"backfilled": counts}, ensure_ascii=False))
        if args.ensure_constraints:
            ensure_uid_constraints(session)


if __name__ == "__main__":
    main()