DataOps-platform/tests/integration/test_data_factory_postgres_lifecycle.py

from __future__ import annotations

import hashlib
import json

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

from app.core.common.identifiers import new_governance_uid
from app.core.data_rules.repository import DataRuleRepository

pytestmark = pytest.mark.integration
DATABASE_URL = "postgresql://dataops:dataops-test-password@127.0.0.1:15432/dataops"


def _insert_deployment(
    connection,
    *,
    deployment_id: str,
    version_id: str,
    dataflow_uid: str,
    actor_uid: str,
    status: str,
    idempotency_key: str,
) -> None:
    hashes = {
        "package_hash": "a" * 64,
        "binding_hash": "b" * 64,
        "schema_snapshot_hash": "c" * 64,
        "workflow_spec_hash": "d" * 64,
        "schedule_hash": "e" * 64,
        "engine_definition_hash": "f" * 64,
    }
    connection.execute(
        text(
            """
            INSERT INTO public.dataflow_deployments
            (id, dataflow_version_id, dataflow_uid, environment, status,
             package, package_hash, binding_snapshot, binding_hash,
             schema_snapshots, schema_snapshot_hash, physical_plan_hashes,
             workflow_spec, workflow_spec_hash, schedule_snapshot,
             schedule_hash, engine_namespace, engine_definition_id,
             engine_definition_hash, created_by, create_reason,
             create_idempotency_key, correlation_id)
            VALUES
            (CAST(:id AS uuid), CAST(:version_id AS uuid),
             CAST(:dataflow_uid AS uuid), 'production', :status,
             CAST(:package AS jsonb), :package_hash,
             CAST(:binding AS jsonb), :binding_hash,
             CAST(:schemas AS jsonb), :schema_snapshot_hash,
             CAST(:plans AS jsonb), CAST(:workflow AS jsonb),
             :workflow_spec_hash, CAST(:schedule AS jsonb), :schedule_hash,
             'dataops.factory', :definition_id, :engine_definition_hash,
             CAST(:actor_uid AS uuid), 'integration acceptance', :key,
             CAST(:correlation_id AS uuid))
            """
        ),
        {
            "id": deployment_id,
            "version_id": version_id,
            "dataflow_uid": dataflow_uid,
            "status": status,
            "package": json.dumps({"schema_version": "1.0"}),
            "binding": json.dumps(
                {"input": {"binding_id": "in"}, "output": {"binding_id": "out"}}
            ),
            "schemas": json.dumps({"input": {}, "output": {}}),
            "plans": json.dumps(["0" * 64]),
            "workflow": json.dumps(
                {
                    "schema_version": "1.0",
                    "dataflow_uid": dataflow_uid,
                    "name": "factory postgres acceptance",
                    "nodes": [],
                    "edges": [],
                    "parameters": {},
                }
            ),
            "schedule": json.dumps(
                {
                    "schema_version": "1.0",
                    "timezone": "Asia/Shanghai",
                    "triggers": [{"type": "manual"}],
                    "max_concurrency": 1,
                    "timeout_seconds": 60,
                    "retry": {"max_attempts": 1, "delay_seconds": 0},
                    "conflict_policy": "queue",
                    "backfill": {"max_days": 1, "max_runs": 1},
                }
            ),
            "definition_id": f"flow-{version_id}",
            "actor_uid": actor_uid,
            "key": idempotency_key,
            "correlation_id": new_governance_uid(),
            **hashes,
        },
    )


def test_postgres_activation_and_rollback_are_atomic_and_idempotent():
    engine = create_engine(DATABASE_URL, pool_pre_ping=True)
    actor_uid = new_governance_uid()
    dataflow_uid = new_governance_uid()
    version_ids = [new_governance_uid(), new_governance_uid()]
    deployment_ids = [new_governance_uid(), new_governance_uid()]
    evidence_id = new_governance_uid()
    operation_ids: list[str] = []
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    "INSERT INTO public.users "
                    "(id, username, password_hash) "
                    "VALUES (CAST(:id AS uuid), :username, 'not-used')"
                ),
                {"id": actor_uid, "username": f"factory-{actor_uid}"},
            )
            for index, version_id in enumerate(version_ids, start=1):
                connection.execute(
                    text(
                        """
                        INSERT INTO public.dataflow_versions
                        (id, dataflow_uid, version_no, name, dataflow_spec,
                         input_schema_hashes, output_schema_hash, status,
                         created_by)
                        VALUES
                        (CAST(:id AS uuid), CAST(:uid AS uuid), :version,
                         'factory postgres acceptance', '{}'::jsonb,
                         '{}'::jsonb, :hash, 'released',
                         CAST(:actor AS uuid))
                        """
                    ),
                    {
                        "id": version_id,
                        "uid": dataflow_uid,
                        "version": index,
                        "hash": "9" * 64,
                        "actor": actor_uid,
                    },
                )
            _insert_deployment(
                connection,
                deployment_id=deployment_ids[0],
                version_id=version_ids[0],
                dataflow_uid=dataflow_uid,
                actor_uid=actor_uid,
                status="active",
                idempotency_key="factory-pg-create-v1",
            )
            _insert_deployment(
                connection,
                deployment_id=deployment_ids[1],
                version_id=version_ids[1],
                dataflow_uid=dataflow_uid,
                actor_uid=actor_uid,
                status="canary",
                idempotency_key="factory-pg-create-v2",
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_canary_evidence
                    (id, deployment_id, execution_id, status, package_hash,
                     binding_hash, schema_snapshot_hash, physical_plan_hashes,
                     workflow_spec_hash, schedule_hash, engine_definition_hash,
                     verified_by, expires_at)
                    VALUES
                    (CAST(:id AS uuid), CAST(:deployment AS uuid), 'exec-pg',
                     'passed', :package_hash, :binding_hash, :schema_hash,
                     CAST(:plans AS jsonb), :workflow_hash, :schedule_hash,
                     :engine_hash, CAST(:actor AS uuid),
                     CURRENT_TIMESTAMP + INTERVAL '15 minutes')
                    """
                ),
                {
                    "id": evidence_id,
                    "deployment": deployment_ids[1],
                    "actor": actor_uid,
                    "package_hash": "a" * 64,
                    "binding_hash": "b" * 64,
                    "schema_hash": "c" * 64,
                    "plans": json.dumps(["0" * 64]),
                    "workflow_hash": "d" * 64,
                    "schedule_hash": "e" * 64,
                    "engine_hash": "f" * 64,
                },
            )
            connection.execute(
                text(
                    "UPDATE public.dataflow_deployments "
                    "SET canary_evidence_id = CAST(:evidence AS uuid) "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"evidence": evidence_id, "id": deployment_ids[1]},
            )

        with Session(engine) as session:
            repository = DataRuleRepository(session)
            claimed, activation = repository.claim_operation(
                deployment_ids[1],
                "activate",
                "factory-pg-activate",
                actor_uid,
                new_governance_uid(),
                "activate acceptance candidate",
                {"evidence_id": evidence_id},
            )
            assert claimed is True
            operation_ids.append(activation["id"])
            session.commit()

            # A different idempotency key/action in the same
            # (dataflow_uid, environment) scope observes the current owner and
            # cannot create a second in-flight external mutation.
            with Session(engine) as competing_session:
                competing = DataRuleRepository(competing_session)
                second_claimed, blocker = competing.claim_operation(
                    deployment_ids[1],
                    "rollback",
                    "factory-pg-competing-operation",
                    actor_uid,
                    new_governance_uid(),
                    "must be fenced by activation",
                    {"candidate_deployment_id": deployment_ids[1]},
                )
                assert second_claimed is False
                assert blocker["id"] == activation["id"]
                assert blocker["fencing_epoch"] == activation["fencing_epoch"]
                competing_session.rollback()

            active = repository.activate_atomic(
                deployment_ids[1],
                0,
                evidence_id,
                actor_uid,
                activation,
            )
            session.commit()
            assert active["status"] == "active"
            assert active["previous_active_deployment_id"] == deployment_ids[0]
            assert (
                repository.get_deployment(deployment_ids[0])["status"] == "superseded"
            )

            claimed, rollback = repository.claim_operation(
                deployment_ids[1],
                "rollback",
                "factory-pg-rollback",
                actor_uid,
                new_governance_uid(),
                "rollback acceptance candidate",
                {"candidate_deployment_id": deployment_ids[1]},
            )
            assert claimed is True
            operation_ids.append(rollback["id"])
            restored = repository.rollback_atomic(
                deployment_ids[1],
                active["lock_version"],
                actor_uid,
                rollback,
            )
            session.commit()
            assert restored["rolled_back"]["status"] == "rolled_back"
            assert restored["active"]["id"] == deployment_ids[0]
            assert restored["active"]["status"] == "active"

            replayed, prior = repository.claim_operation(
                deployment_ids[1],
                "rollback",
                "factory-pg-rollback",
                actor_uid,
                rollback["correlation_id"],
                "rollback acceptance candidate",
                {"candidate_deployment_id": deployment_ids[1]},
            )
            assert replayed is False
            assert prior["status"] == "completed"
            assert prior["result"]["active"]["id"] == deployment_ids[0]

            claimed, crashed = repository.claim_operation(
                deployment_ids[0],
                "deploy_disabled",
                "factory-pg-crashed-claim",
                actor_uid,
                new_governance_uid(),
                "simulate process exit after durable claim",
                {},
            )
            assert claimed is True
            operation_ids.append(crashed["id"])
            first_epoch = crashed["fencing_epoch"]
            session.commit()
            repository.mark_operation_unknown(
                crashed, "acceptance_unknown_before_external_outcome"
            )
            session.commit()

            # Once any operation in the dataflow/environment scope is unknown,
            # a different deployment, action and idempotency key must return
            # the original blocker and must not create another claim.
            blocked, original_unknown = repository.claim_operation(
                deployment_ids[1],
                "run_canary",
                "factory-pg-must-not-bypass-unknown",
                actor_uid,
                new_governance_uid(),
                "must reconcile original unknown first",
                {"inputs_hash": "7" * 64},
            )
            assert blocked is False
            assert original_unknown["id"] == crashed["id"]
            assert original_unknown["deployment_id"] == deployment_ids[0]
            assert original_unknown["action"] == "deploy_disabled"
            assert original_unknown["status"] == "unknown"
            assert original_unknown["idempotency_key"] == (
                "factory-pg-crashed-claim"
            )
            assert (
                session.execute(
                    text(
                        "SELECT count(*) FROM "
                        "public.dataflow_deployment_operations "
                        "WHERE idempotency_key = "
                        "'factory-pg-must-not-bypass-unknown'"
                    )
                ).scalar_one()
                == 0
            )

            newer_unknown_id = new_governance_uid()
            operation_ids.append(newer_unknown_id)
            empty_request = "{}"
            session.execute(
                text(
                    "INSERT INTO public.dataflow_deployment_operations "
                    "(id, deployment_id, action, idempotency_key, status, "
                    "actor_uid, correlation_id, reason, request, request_hash, "
                    "error_code, owner_token, lease_expires_at, attempt_epoch, "
                    "fencing_epoch, started_at) VALUES "
                    "(CAST(:id AS uuid), CAST(:deployment_id AS uuid), "
                    "'deploy_disabled', :key, 'unknown', "
                    "CAST(:actor_uid AS uuid), CAST(:correlation_id AS uuid), "
                    "'newer legacy unknown', CAST(:request AS jsonb), "
                    ":request_hash, 'legacy_claim_requires_reconciliation', "
                    "CAST(:owner_token AS uuid), "
                    "CURRENT_TIMESTAMP - interval '1 second', 1, 999999, "
                    "CURRENT_TIMESTAMP + interval '1 second')"
                ),
                {
                    "id": newer_unknown_id,
                    "deployment_id": deployment_ids[1],
                    "key": "factory-pg-newer-legacy-unknown",
                    "actor_uid": actor_uid,
                    "correlation_id": new_governance_uid(),
                    "request": empty_request,
                    "request_hash": hashlib.sha256(
                        empty_request.encode("utf-8")
                    ).hexdigest(),
                    "owner_token": new_governance_uid(),
                },
            )
            session.commit()
            with pytest.raises(
                ValueError,
                match="older deployment operation requires reconciliation",
            ):
                repository.recover_operation(
                    deployment_ids[1],
                    "deploy_disabled",
                    "factory-pg-newer-legacy-unknown",
                    actor_uid,
                )
            session.rollback()

            session.execute(
                text(
                    "UPDATE public.dataflow_deployment_operations "
                    "SET lease_expires_at = CURRENT_TIMESTAMP - interval '1 second' "
                    "WHERE id = CAST(:id AS uuid); "
                    "UPDATE public.dataflow_deployment_operation_leases "
                    "SET lease_expires_at = CURRENT_TIMESTAMP - interval '1 second' "
                    "WHERE operation_id = CAST(:id AS uuid)"
                ),
                {"id": crashed["id"]},
            )
            session.commit()
            recovered = repository.recover_operation(
                deployment_ids[0],
                "deploy_disabled",
                "factory-pg-crashed-claim",
                actor_uid,
            )
            assert recovered["attempt_epoch"] == 2
            assert recovered["fencing_epoch"] > first_epoch
            assert repository.owns_operation_lease(recovered) is True
            repository.fail_operation(recovered, "acceptance_cleanup")
            session.commit()
    finally:
        with engine.begin() as connection:
            connection.execute(
                text(
                    "UPDATE public.dataflow_deployments "
                    "SET canary_evidence_id = NULL, "
                    "previous_active_deployment_id = NULL "
                    "WHERE id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": deployment_ids},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_deployment_transitions "
                    "WHERE deployment_id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": deployment_ids},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_deployment_operations "
                    "WHERE deployment_id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": deployment_ids},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_canary_evidence "
                    "WHERE deployment_id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": deployment_ids},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_deployments "
                    "WHERE id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": deployment_ids},
            )
            connection.execute(
                text(
                    "DELETE FROM public.dataflow_versions "
                    "WHERE id = ANY(CAST(:ids AS uuid[]))"
                ),
                {"ids": version_ids},
            )
            connection.execute(
                text("DELETE FROM public.users WHERE id = CAST(:id AS uuid)"),
                {"id": actor_uid},
            )
        engine.dispose()