DataOps-platform/app/commands/migrate_datasource_credentials.py

"""Explicitly encrypt one legacy DataSource credential and remove plaintext."""

import argparse
import json
from dataclasses import dataclass

from app.core.data_source.models import DataSourceCredential


@dataclass(frozen=True)
class MigrationResult:
    data_source_uid: str
    credential_version: int
    status: str


def migrate_one(
    data_source_uid,
    graph_session,
    credential_repository,
    platform_session,
    *,
    actor_uid,
):
    record = graph_session.run(
        """
        MATCH (n:DataSource {uid: $uid})
        RETURN properties(n) AS properties
        """,
        {"uid": str(data_source_uid)},
    ).single()
    if record is None:
        raise ValueError("data source was not found")
    properties = dict(record["properties"])
    username = properties.get("username")
    password = properties.get("password")
    if not username or not password:
        raise ValueError("legacy plaintext credential is incomplete")

    credential = DataSourceCredential(str(username), str(password))
    try:
        sealed = credential_repository.create_version(
            platform_session,
            data_source_uid=str(data_source_uid),
            credential=credential,
            actor_uid=actor_uid,
        )
        verified = credential_repository.codec.decrypt(sealed)
        if verified != credential:
            raise ValueError("credential round-trip verification failed")
        platform_session.commit()
    except Exception:
        platform_session.rollback()
        raise

    try:
        graph_session.run(
            """
            MATCH (n:DataSource {uid: $uid})
            SET n.credential_ref = $credential_ref,
                n.credential_version = $credential_version
            """,
            {
                "uid": str(data_source_uid),
                "credential_ref": str(data_source_uid),
                "credential_version": sealed.credential_version,
            },
        )
        graph_session.run(
            """
            MATCH (n:DataSource {uid: $uid})
            REMOVE n.username, n.password, n.conn_str,
                   n.connection_string, n.connection_url
            """,
            {"uid": str(data_source_uid)},
        )
    except Exception as graph_error:
        try:
            credential_repository.compensate_failed_activation(
                platform_session,
                data_source_uid=str(data_source_uid),
                failed_version=sealed.credential_version,
                restore_version=None,
                actor_uid=actor_uid,
            )
            platform_session.commit()
        except Exception:
            platform_session.rollback()
        raise ValueError(
            "credential migration requires reconciliation"
        ) from graph_error

    return MigrationResult(
        data_source_uid=str(data_source_uid),
        credential_version=sealed.credential_version,
        status="migrated",
    )


def main():
    parser = argparse.ArgumentParser(
        description="Encrypt one legacy DataSource credential"
    )
    parser.add_argument("--uid", required=True)
    parser.add_argument(
        "--confirm-encrypt-and-remove-plaintext",
        action="store_true",
    )
    args = parser.parse_args()
    if not args.confirm_encrypt_and_remove_plaintext:
        parser.error(
            "--confirm-encrypt-and-remove-plaintext is required"
        )

    from app import create_app, db
    from app.core.data_source.credentials import (
        CredentialCodec,
        DataSourceCredentialRepository,
    )
    from app.services.neo4j_driver import neo4j_driver

    app = create_app()
    with app.app_context():
        codec = CredentialCodec.from_base64(
            app.config.get("DATASOURCE_CREDENTIAL_MASTER_KEY", ""),
            app.config.get("DATASOURCE_CREDENTIAL_KEY_VERSION", "v1"),
        )
        repository = DataSourceCredentialRepository(codec)
        with neo4j_driver.get_session() as graph_session:
            result = migrate_one(
                args.uid,
                graph_session,
                repository,
                db.session,
                actor_uid=None,
            )
        print(
            json.dumps(
                {
                    "data_source_uid": result.data_source_uid,
                    "credential_version": result.credential_version,
                    "status": result.status,
                },
                ensure_ascii=False,
            )
        )


if __name__ == "__main__":
    main()