DataOps-platform/tests/integration/test_mcp_gateway_persistence.py

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": "deploy_disabled_version",
                "decision": "allowed",
                "detail": {"candidate_id": candidate_id},
            }
        )
        with engine.connect() as connection:
            row = connection.execute(
                text(
                    "SELECT actor_subject, action, decision "
                    "FROM public.workflow_plan_audits "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            ).one()
        assert tuple(row) == (
            "agent-scheduler",
            "deploy_disabled_version",
            "allowed",
        )
    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()