DataOps-platform/app/commands/reconcile_datasource_credentials.py

"""Report and narrowly repair Neo4j/PostgreSQL credential references."""

import argparse
import json

from sqlalchemy import text


def inspect_reconciliation(graph_session, platform_session):
    graph_records = graph_session.run(
        """
        MATCH (n:DataSource)
        RETURN n.uid AS uid, n.credential_ref AS credential_ref,
               n.credential_version AS credential_version,
               n.username IS NOT NULL OR n.password IS NOT NULL
                   AS plaintext_remaining
        """
    )
    graph = {
        str(record["uid"]): dict(record)
        for record in graph_records
        if record.get("uid")
    }
    rows = platform_session.execute(
        text(
            """
            SELECT data_source_uid::text AS uid, credential_version, status
            FROM public.datasource_credentials
            ORDER BY data_source_uid, credential_version
            """
        )
    ).mappings()
    credentials = {}
    for row in rows:
        credentials.setdefault(str(row["uid"]), []).append(dict(row))

    report = []
    for uid in sorted(set(graph) | set(credentials)):
        issues = []
        node = graph.get(uid)
        versions = credentials.get(uid, [])
        active = [
            row for row in versions if row["status"] == "active"
        ]
        if node is None:
            issues.append("orphaned_credential")
        else:
            if node.get("plaintext_remaining"):
                issues.append("plaintext_remaining")
            referenced = node.get("credential_version")
            matching = [
                row
                for row in versions
                if row["credential_version"] == referenced
            ]
            if not matching:
                issues.append("missing_credential")
            elif matching[0]["status"] == "revoked":
                issues.append("revoked_credential_referenced")
            if active and referenced != active[-1]["credential_version"]:
                issues.append("stale_credential_version")
        if issues:
            report.append({"data_source_uid": uid, "issues": issues})
    return report


def repair_one(data_source_uid, graph_session, platform_session):
    row = (
        platform_session.execute(
            text(
                """
                SELECT credential_version
                FROM public.datasource_credentials
                WHERE data_source_uid = CAST(:uid AS uuid)
                  AND status = 'active'
                ORDER BY credential_version DESC
                LIMIT 1
                """
            ),
            {"uid": str(data_source_uid)},
        )
        .mappings()
        .one_or_none()
    )
    if row is None:
        raise ValueError("no active credential is available for repair")
    graph_session.run(
        """
        MATCH (n:DataSource {uid: $uid})
        SET n.credential_ref = $uid,
            n.credential_version = $credential_version
        """,
        {
            "uid": str(data_source_uid),
            "credential_version": int(row["credential_version"]),
        },
    )
    return {
        "data_source_uid": str(data_source_uid),
        "status": "reference_repaired",
    }


def main():
    parser = argparse.ArgumentParser(
        description="Reconcile DataSource credential references"
    )
    parser.add_argument("--repair", action="store_true")
    parser.add_argument("--uid")
    args = parser.parse_args()
    if args.repair and not args.uid:
        parser.error("--repair requires --uid")

    from app import create_app, db
    from app.services.neo4j_driver import neo4j_driver

    app = create_app()
    with app.app_context():
        with neo4j_driver.get_session() as graph_session:
            report = inspect_reconciliation(graph_session, db.session)
            for item in report:
                print(json.dumps(item, ensure_ascii=False))
            if args.repair:
                result = repair_one(args.uid, graph_session, db.session)
                print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()