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()