DataOps-platform/tests/integration/test_data_factory_postgres_lifecycle.py

from __future__ import annotations

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"])
            active = repository.activate_atomic(
                deployment_ids[1],
                0,
                evidence_id,
                actor_uid,
                activation["id"],
            )
            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["id"],
            )
            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]
            session.rollback()
    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()