import os import pytest from sqlalchemy import create_engine, text from app.core.common.identifiers import new_governance_uid from app.core.mcp.persistence import ( PostgresAuditSink, PostgresContextRepository, PostgresSchedulingPlanStore, ) pytestmark = pytest.mark.skipif( os.getenv("DATAOPS_MCP_PERSISTENCE_INTEGRATION") != "1", reason="set DATAOPS_MCP_PERSISTENCE_INTEGRATION=1", ) def workflow_spec(dataflow_uid): return { "schema_version": "1.0", "dataflow_uid": dataflow_uid, "name": "V53 persistence acceptance", "nodes": [ { "id": "read_orders", "type": "sql.query", "data_source_uid": new_governance_uid(), "purpose": "read", "config": {"statement": "SELECT 1", "parameters": {}}, } ], "edges": [], "parameters": {}, } def schedule_plan(*, triggers=None): return { "schema_version": "1.0", "timezone": "Asia/Shanghai", "triggers": triggers or [{"type": "manual"}], "max_concurrency": 1, "conflict_policy": "skip", "timeout_seconds": 300, "retry": {"max_attempts": 1, "delay_seconds": 0}, "backfill": {"max_days": 2, "max_runs": 5}, } def test_postgres_store_persists_candidate_canary_idempotency_and_audit(): database_url = os.environ["DATAOPS_MCP_TEST_DATABASE_URL"] engine = create_engine(database_url, pool_pre_ping=True) store = PostgresSchedulingPlanStore(engine) audit = PostgresAuditSink(engine) dataflow_uid = new_governance_uid() correlation_id = new_governance_uid() candidate_id = None try: candidate = store.create_candidate( { "dataflow_uid": dataflow_uid, "business_domain": "sales", "environment": "test", "workflow_spec": workflow_spec(dataflow_uid), "schedule_plan": schedule_plan(), "workflow_hash": "a" * 64, "status": "candidate", "created_by": "agent-scheduler", "correlation_id": correlation_id, } ) candidate_id = candidate["candidate_id"] assert candidate["version_no"] >= 1 assert store.get_candidate(candidate_id)["business_domain"] == "sales" deployment = store.record_deployment( candidate_id, { "candidate_id": candidate_id, "namespace": "dataops.test", "flow_id": f"v53_{candidate_id.replace('-', '')}", "definition_hash": "b" * 64, "engine_revision": "revision-1", "status": "deployed_disabled", }, ) assert store.get_deployment(candidate_id) == deployment evidence = store.record_canary_execution( candidate_id, { "candidate_id": candidate_id, "execution_id": f"canary-{candidate_id}", "status": "started", "sample_runs": 0, }, ) assert evidence["status"] == "started" assert store.get_canary_evidence(candidate_id)["status"] == "started" verified = store.record_canary_verification( candidate_id, { **evidence, "status": "passed", "sample_runs": 1, "verified_by": "dataops-canary-verifier", "verification_summary": { "engine_state": "SUCCESS", "equivalent": True, "metrics": {"row_count_delta": 0}, }, }, ) assert verified["status"] == "passed" assert store.get_canary_evidence(candidate_id)["verified_by"] == ( "dataops-canary-verifier" ) promotion_key = f"promotion:v53:{candidate_id}" first = store.claim_promotion( promotion_key, candidate_id, actor_subject="agent-scheduler", correlation_id=correlation_id, ) second = store.claim_promotion( promotion_key, candidate_id, actor_subject="agent-scheduler", correlation_id=correlation_id, ) assert first[0] is True assert second[0] is False audit.record( { "subject": "agent-scheduler", "roles": ["scheduler"], "business_domain": "sales", "environment": "test", "correlation_id": correlation_id, "action": "ai_planning_closed_loop", "decision": "allowed", "model_provider": "deepseek", "model_name": "deepseek-chat", "prompt_version": "v54.1", "schema_version": "v54", "context_hash": "c" * 64, "detail": { "candidate_id": candidate_id, "workflow_hash": "a" * 64, "candidate_plan": {"schema_version": "1.0"}, }, } ) with engine.connect() as connection: row = connection.execute( text( "SELECT actor_subject, action, decision, model_provider, " "model_name, prompt_version, schema_version, context_hash, " "candidate_hash, decision_detail " "FROM public.workflow_plan_audits " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ).one() assert tuple(row[:9]) == ( "agent-scheduler", "ai_planning_closed_loop", "allowed", "deepseek", "deepseek-chat", "v54.1", "v54", "c" * 64, "a" * 64, ) assert row.decision_detail["candidate_plan"]["schema_version"] == "1.0" finally: if candidate_id is not None: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.workflow_gateway_operations " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) connection.execute( text( "DELETE FROM public.workflow_plan_audits " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) connection.execute( text( "DELETE FROM public.dataflow_workflow_versions " "WHERE id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) engine.dispose() def test_idempotency_key_cannot_replay_another_workflow_version(): database_url = os.environ["DATAOPS_MCP_TEST_DATABASE_URL"] engine = create_engine(database_url, pool_pre_ping=True) store = PostgresSchedulingPlanStore(engine) correlation_id = new_governance_uid() candidate_ids = [] key = f"promotion:v53:scope:{new_governance_uid()}" try: for suffix in ("1", "2"): dataflow_uid = new_governance_uid() candidate = store.create_candidate( { "dataflow_uid": dataflow_uid, "business_domain": "sales", "environment": "test", "workflow_spec": workflow_spec(dataflow_uid), "schedule_plan": schedule_plan(), "workflow_hash": suffix * 64, "status": "candidate", "created_by": "agent-scheduler", "correlation_id": correlation_id, } ) candidate_ids.append(candidate["candidate_id"]) store.record_deployment( candidate["candidate_id"], { "candidate_id": candidate["candidate_id"], "namespace": "dataops.test", "flow_id": f"idempotency-scope-{suffix}", "definition_hash": suffix * 64, "engine_revision": f"revision-{suffix}", "status": "deployed_disabled", }, ) store.claim_promotion( key, candidate_ids[0], actor_subject="agent-scheduler", correlation_id=correlation_id, ) with pytest.raises(ValueError, match="another workflow version"): store.claim_promotion( key, candidate_ids[1], actor_subject="agent-scheduler", correlation_id=correlation_id, ) finally: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.workflow_gateway_operations " "WHERE idempotency_key = :key" ), {"key": key}, ) if candidate_ids: connection.execute( text( "DELETE FROM public.dataflow_workflow_versions " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), {"ids": candidate_ids}, ) engine.dispose() def test_postgres_store_completes_promotion_pause_and_rollback_state(): database_url = os.environ["DATAOPS_MCP_TEST_DATABASE_URL"] engine = create_engine(database_url, pool_pre_ping=True) store = PostgresSchedulingPlanStore(engine) dataflow_uid = new_governance_uid() correlation_id = new_governance_uid() candidate_ids = [] try: for suffix in ("previous", "current"): candidate = store.create_candidate( { "dataflow_uid": dataflow_uid, "business_domain": "sales", "environment": "test", "workflow_spec": workflow_spec(dataflow_uid), "schedule_plan": schedule_plan(), "workflow_hash": ( "c" * 63 + ("1" if suffix == "previous" else "2") ), "status": "candidate", "created_by": "agent-scheduler", "correlation_id": correlation_id, } ) candidate_ids.append(candidate["candidate_id"]) store.record_deployment( candidate["candidate_id"], { "candidate_id": candidate["candidate_id"], "namespace": "dataops.test", "flow_id": f"flow-{suffix}-{candidate['version_no']}", "definition_hash": "d" * 64, "engine_revision": f"revision-{suffix}", "status": "deployed_disabled", }, ) previous_id, current_id = candidate_ids promotion_key = f"promotion:{previous_id}" claimed, promotion = store.claim_promotion( promotion_key, previous_id, actor_subject="agent-scheduler", correlation_id=correlation_id, ) assert claimed is True store.complete_promotion(promotion_key, previous_id, promotion) store.mark_paused(previous_id) with engine.connect() as connection: states = connection.execute( text( "SELECT v.status, b.role, b.status, s.status " "FROM public.dataflow_workflow_versions v " "JOIN public.workflow_engine_bindings b " "ON b.workflow_version_id = v.id " "JOIN public.workflow_schedules s " "ON s.workflow_version_id = v.id " "WHERE v.id = CAST(:id AS uuid)" ), {"id": previous_id}, ).one() assert tuple(states) == ("active", "primary", "disabled", "paused") previous = store.get_previous_deployment(current_id) assert previous["candidate_id"] == previous_id rollback_key = f"rollback:{current_id}" claimed, rollback = store.claim_rollback( rollback_key, current_id, actor_subject="agent-scheduler", correlation_id=correlation_id, ) assert claimed is True store.complete_rollback(rollback_key, current_id, rollback) with engine.connect() as connection: operation_status = connection.execute( text( "SELECT status FROM public.workflow_gateway_operations " "WHERE action = 'rollback_to_previous_version' " "AND idempotency_key = :key" ), {"key": rollback_key}, ).scalar_one() assert operation_status == "succeeded" finally: if candidate_ids: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.workflow_gateway_operations " "WHERE workflow_version_id = ANY(CAST(:ids AS uuid[]))" ), {"ids": candidate_ids}, ) connection.execute( text( "DELETE FROM public.workflow_plan_audits " "WHERE workflow_version_id = ANY(CAST(:ids AS uuid[]))" ), {"ids": candidate_ids}, ) connection.execute( text( "DELETE FROM public.dataflow_workflow_versions " "WHERE id = ANY(CAST(:ids AS uuid[]))" ), {"ids": candidate_ids}, ) engine.dispose() def test_postgres_store_tracks_retries_and_bounded_backfill_runs(): database_url = os.environ["DATAOPS_MCP_TEST_DATABASE_URL"] engine = create_engine(database_url, pool_pre_ping=True) store = PostgresSchedulingPlanStore(engine) dataflow_uid = new_governance_uid() correlation_id = new_governance_uid() candidate_id = None original_run_id = new_governance_uid() try: candidate = store.create_candidate( { "dataflow_uid": dataflow_uid, "business_domain": "sales", "environment": "test", "workflow_spec": workflow_spec(dataflow_uid), "schedule_plan": schedule_plan( triggers=[{"type": "cron", "expression": "0 0 * * *"}] ), "workflow_hash": "e" * 64, "status": "candidate", "created_by": "agent-scheduler", "correlation_id": correlation_id, } ) candidate_id = candidate["candidate_id"] store.record_deployment( candidate_id, { "candidate_id": candidate_id, "namespace": "dataops.test", "flow_id": f"retry-backfill-{candidate['version_no']}", "definition_hash": "f" * 64, "engine_revision": "revision-retry", "status": "deployed_disabled", }, ) with engine.begin() as connection: connection.execute( text( "INSERT INTO public.workflow_runs " "(id, workflow_version_id, engine_type, " "engine_execution_id, trigger_type, status, " "correlation_id, attempt, run_metadata) " "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), " "'kestra', 'failed-execution-1', 'manual', 'failed', " "CAST(:correlation_id AS uuid), 1, '{}'::jsonb)" ), { "id": original_run_id, "version_id": candidate_id, "correlation_id": correlation_id, }, ) record = store.get_execution_record("failed-execution-1") assert record["retry_count"] == 0 assert record["business_domain"] == "sales" store.record_retry( "failed-execution-1", {"id": "replay-execution-1"}, ) assert store.get_execution_record("failed-execution-1")["retry_count"] == 1 estimated = store.estimate_backfill_runs( candidate_id, "2026-07-01T00:00:00Z", "2026-07-03T00:00:00Z", ) assert estimated == 2 store.record_backfill( candidate_id, {"id": "backfill-execution-1", "estimated_runs": estimated}, ) with engine.connect() as connection: trigger_type = connection.execute( text( "SELECT trigger_type FROM public.workflow_runs " "WHERE engine_execution_id = 'backfill-execution-1'" ) ).scalar_one() assert trigger_type == "backfill" finally: if candidate_id is not None: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.workflow_runs " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) connection.execute( text( "DELETE FROM public.workflow_gateway_operations " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) connection.execute( text( "DELETE FROM public.dataflow_workflow_versions " "WHERE id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) engine.dispose() def test_postgres_context_read_model_exposes_only_bounded_operational_data(): database_url = os.environ["DATAOPS_MCP_TEST_DATABASE_URL"] engine = create_engine(database_url, pool_pre_ping=True) store = PostgresSchedulingPlanStore(engine) repository = PostgresContextRepository(engine, max_concurrency=5) dataflow_uid = new_governance_uid() correlation_id = new_governance_uid() candidate_id = None execution_id = f"context-{new_governance_uid()}" try: candidate = store.create_candidate( { "dataflow_uid": dataflow_uid, "business_domain": "sales", "environment": "test", "workflow_spec": workflow_spec(dataflow_uid), "schedule_plan": schedule_plan( triggers=[{"type": "cron", "expression": "0 0 * * *"}] ), "workflow_hash": "9" * 64, "status": "candidate", "created_by": "agent-context", "correlation_id": correlation_id, } ) candidate_id = candidate["candidate_id"] with engine.begin() as connection: connection.execute( text( "INSERT INTO public.workflow_runs " "(id, workflow_version_id, engine_type, " "engine_execution_id, trigger_type, status, " "correlation_id, attempt, run_metadata) " "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), " "'kestra', :execution_id, 'manual', 'success', " "CAST(:correlation_id AS uuid), 1, '{}'::jsonb)" ), { "id": new_governance_uid(), "version_id": candidate_id, "execution_id": execution_id, "correlation_id": correlation_id, }, ) assert repository.describe_dataflow(dataflow_uid)["business_domain"] == "sales" assert any(row["uid"] == dataflow_uid for row in repository.list_dataflows()) assert repository.get_sla_constraints(dataflow_uid) == { "timezone": "Asia/Shanghai", "max_duration_seconds": 300, } history = repository.get_execution_history(dataflow_uid, limit=10, days=1) assert history[0]["status"] == "success" capacity = repository.estimate_schedule_capacity("sales", schedule_plan()) assert capacity["available_slots"] <= 5 simulation = repository.simulate_schedule( "sales", schedule_plan(triggers=[{"type": "cron", "expression": "0 0 * * *"}]), ) assert len(simulation["next_runs"]) == 5 finally: if candidate_id is not None: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.workflow_runs " "WHERE workflow_version_id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) connection.execute( text( "DELETE FROM public.dataflow_workflow_versions " "WHERE id = CAST(:id AS uuid)" ), {"id": candidate_id}, ) engine.dispose()