DataOps-platform/migrations/versions/20260813_499_runtime_control_approval.py

"""Close WP09 recovery and canary approvals behind one-shot database facts."""

from alembic import op

revision = "20260813_499"
down_revision = "20260813_498"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute(r"""
    CREATE TABLE public.agent_runtime_control_events (
      event_id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
      tenant_id text NOT NULL,
      idempotency_key text NOT NULL,
      request_digest char(64) NOT NULL,
      payload_digest char(64) NOT NULL,
      operation text NOT NULL CHECK (operation IN ('recover','canary_promote','canary_rollback')),
      approval_task_uid uuid NOT NULL REFERENCES public.governance_tasks(uid),
      result jsonb NOT NULL,
      created_at timestamptz NOT NULL DEFAULT clock_timestamp(),
      UNIQUE(tenant_id,idempotency_key)
    );
    ALTER TABLE public.agent_runtime_control_events OWNER TO dataops_agent_runtime_owner;
    REVOKE ALL ON TABLE public.agent_runtime_control_events FROM PUBLIC,dataops_agent_runtime;
    GRANT SELECT,INSERT,UPDATE ON TABLE public.agent_runtime_control_events TO dataops_agent_runtime_owner;

    CREATE FUNCTION public.agent_runtime_control(p_payload jsonb)
    RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $control$
    DECLARE t record; s record; c record; event record; v_result jsonb; v_fence bigint; v_payload_digest text;
    BEGIN
      IF jsonb_typeof(p_payload)<>'object' OR NOT (p_payload ?& ARRAY['tenant_id','agent_uid','principal_id','business_domain_uid','environment','operation','approval_task_uid','idempotency_key','request_digest','expected_fence','incident_uid','from_state','to_state','generation','from_generation','to_generation','dataset_version','dataset_digest','metrics_digest','threshold_policy','issued_at']) OR p_payload-ARRAY['tenant_id','agent_uid','principal_id','business_domain_uid','environment','operation','approval_task_uid','idempotency_key','request_digest','expected_fence','incident_uid','from_state','to_state','generation','from_generation','to_generation','dataset_version','dataset_digest','metrics_digest','threshold_policy','issued_at']<>'{}'::jsonb THEN RAISE EXCEPTION 'runtime control payload is closed'; END IF;
      IF (p_payload->>'agent_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'principal_id') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'business_domain_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'approval_task_uid') !~ '^[0-9a-f-]{36}$' OR (p_payload->>'request_digest') !~ '^[0-9a-f]{64}$' OR (p_payload->>'expected_fence') !~ '^[0-9]+$' OR (p_payload->>'issued_at') !~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9:.+-]+$' OR p_payload->>'operation' NOT IN ('recover','canary_promote','canary_rollback') OR (p_payload->>'environment') !~ '^[A-Za-z0-9_-]{1,40}$' OR (p_payload->>'idempotency_key') !~ '^[A-Za-z0-9_.:-]{1,160}$' OR (p_payload->>'dataset_version') !~ '^[A-Za-z0-9._-]{1,120}$' OR (p_payload->>'dataset_digest') !~ '^[0-9a-f]{64}$' OR (p_payload->>'metrics_digest') !~ '^[0-9a-f]{64}$' OR (p_payload->>'threshold_policy') !~ '^[A-Za-z0-9._-]{1,120}$' THEN RAISE EXCEPTION 'runtime control value invalid'; END IF;
      v_payload_digest:=encode(sha256(convert_to(p_payload::text,'utf8')),'hex');
      SELECT * INTO event FROM public.agent_runtime_control_events WHERE tenant_id=p_payload->>'tenant_id' AND idempotency_key=p_payload->>'idempotency_key' FOR UPDATE;
      IF FOUND THEN IF event.request_digest<>p_payload->>'request_digest' OR event.operation<>p_payload->>'operation' OR event.payload_digest<>v_payload_digest THEN RAISE EXCEPTION 'control replay conflict'; END IF; RETURN event.result || jsonb_build_object('replay',true); END IF;
      SELECT * INTO t FROM public.governance_tasks WHERE uid=(p_payload->>'approval_task_uid')::uuid FOR UPDATE;
      IF NOT FOUND THEN RAISE EXCEPTION 'approval context denied'; END IF;
      IF t.status<>'approved' THEN RAISE EXCEPTION 'approval already consumed'; END IF;
      IF t.subject_uid<>(p_payload->>'agent_uid') OR t.business_domain_uid<>(p_payload->>'business_domain_uid') OR t.due_at<=clock_timestamp() OR t.context->>'tenant_id'<>p_payload->>'tenant_id' OR t.context->>'agent_uid'<>p_payload->>'agent_uid' OR t.context->>'principal_id'<>p_payload->>'principal_id' OR t.context->>'business_domain_uid'<>p_payload->>'business_domain_uid' OR t.context->>'environment'<>p_payload->>'environment' OR t.context->>'control_operation'<>p_payload->>'operation' OR t.context->>'request_digest'<>p_payload->>'request_digest' OR t.context->>'expected_fence'<>p_payload->>'expected_fence' OR t.context->>'issued_at'<>p_payload->>'issued_at' OR t.context->>'from_state'<>p_payload->>'from_state' OR t.context->>'to_state'<>p_payload->>'to_state' OR t.context->>'generation'<>p_payload->>'generation' OR t.context->>'from_generation'<>p_payload->>'from_generation' OR t.context->>'to_generation'<>p_payload->>'to_generation' OR t.context->>'dataset_version'<>p_payload->>'dataset_version' OR t.context->>'dataset_digest'<>p_payload->>'dataset_digest' OR t.context->>'metrics_digest'<>p_payload->>'metrics_digest' OR t.context->>'threshold_policy'<>p_payload->>'threshold_policy' OR COALESCE(t.context->>'incident_uid','')<>COALESCE(p_payload->>'incident_uid','') THEN RAISE EXCEPTION 'approval context denied'; END IF;
      IF NOT EXISTS(SELECT 1 FROM public.governance_task_reviews r WHERE r.task_uid=t.uid AND r.decision='approve' AND r.reviewer_uid<>(p_payload->>'principal_id')::uuid) THEN RAISE EXCEPTION 'approval reviewer denied'; END IF;
      IF p_payload->>'operation'='recover' THEN
        IF (p_payload->>'incident_uid') !~ '^[0-9a-f-]{36}$' OR p_payload->>'from_state'<>'paused' OR p_payload->>'to_state'<>'active' THEN RAISE EXCEPTION 'recover control invalid'; END IF;
        SELECT * INTO s FROM public.agent_runtime_states WHERE tenant_id=p_payload->>'tenant_id' FOR UPDATE;
        IF NOT FOUND OR s.state<>p_payload->>'from_state' OR s.lease_fence<>(p_payload->>'expected_fence')::bigint OR s.incident_ref<>(p_payload->>'incident_uid') THEN RAISE EXCEPTION 'runtime state fence conflict'; END IF;
        UPDATE public.agent_runtime_states SET state='active',approval_ref=t.uid::text,lease_fence=lease_fence+1,updated_at=clock_timestamp() WHERE tenant_id=s.tenant_id AND lease_fence=s.lease_fence RETURNING lease_fence INTO v_fence;
      ELSE
        IF COALESCE(p_payload->>'incident_uid','')<>'' OR p_payload->>'generation'<>p_payload->>'from_generation' OR p_payload->>'generation'<>p_payload->>'to_generation' OR p_payload->>'to_state'<>(CASE WHEN p_payload->>'operation'='canary_promote' THEN 'approved' ELSE 'rolled_back' END) THEN RAISE EXCEPTION 'canary control invalid'; END IF;
        SELECT * INTO c FROM public.agent_generation_canaries WHERE tenant_id=p_payload->>'tenant_id' AND generation=p_payload->>'generation' FOR UPDATE;
        IF NOT FOUND OR c.status<>p_payload->>'from_state' OR c.lease_fence<>(p_payload->>'expected_fence')::bigint THEN RAISE EXCEPTION 'generation fence conflict'; END IF;
        UPDATE public.agent_generation_canaries SET status=p_payload->>'to_state',approval_ref=t.uid::text,lease_fence=lease_fence+1 WHERE tenant_id=c.tenant_id AND generation=c.generation AND lease_fence=c.lease_fence RETURNING lease_fence INTO v_fence;
      END IF;
      UPDATE public.governance_tasks SET status='closed',closed_at=clock_timestamp(),updated_at=clock_timestamp() WHERE uid=t.uid AND status='approved';
      IF NOT FOUND THEN RAISE EXCEPTION 'approval already consumed'; END IF;
      v_result:=jsonb_build_object('operation',p_payload->>'operation','lease_fence',v_fence,'approval_task_uid',t.uid::text,'replay',false);
      INSERT INTO public.agent_runtime_control_events(tenant_id,idempotency_key,request_digest,payload_digest,operation,approval_task_uid,result) VALUES(p_payload->>'tenant_id',p_payload->>'idempotency_key',p_payload->>'request_digest',v_payload_digest,p_payload->>'operation',t.uid,v_result);
      RETURN v_result;
    END; $control$;
    ALTER FUNCTION public.agent_runtime_control(jsonb) OWNER TO dataops_agent_runtime_owner;
    REVOKE ALL ON FUNCTION public.agent_runtime_control(jsonb) FROM PUBLIC;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_control(jsonb) TO dataops_agent_runtime,dataops_agent_runtime_owner;
    CREATE OR REPLACE FUNCTION public.agent_runtime_transition(p_tenant text,p_state text,p_incident uuid,p_approval uuid,p_expected_fence bigint) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $deprecated_transition$ BEGIN RAISE EXCEPTION 'deprecated runtime transition; use agent_runtime_control'; END; $deprecated_transition$;
    CREATE OR REPLACE FUNCTION public.agent_runtime_canary_transition(p_tenant text,p_generation text,p_status text,p_approval uuid,p_expected_fence bigint,p_metrics jsonb,p_dataset_version text) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $deprecated_canary$ BEGIN RAISE EXCEPTION 'deprecated runtime canary transition; use agent_runtime_control'; END; $deprecated_canary$;
    REVOKE ALL ON FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint) FROM PUBLIC,dataops_agent_runtime;
    REVOKE ALL ON FUNCTION public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) FROM PUBLIC,dataops_agent_runtime;
    ALTER FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint) OWNER TO dataops_agent_runtime_owner;
    ALTER FUNCTION public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) OWNER TO dataops_agent_runtime_owner;
    """)


def downgrade() -> None:
    op.execute("""
    REVOKE ALL ON FUNCTION public.agent_runtime_control(jsonb) FROM PUBLIC;
    DROP FUNCTION IF EXISTS public.agent_runtime_control(jsonb);
    DROP TABLE IF EXISTS public.agent_runtime_control_events;
    CREATE OR REPLACE FUNCTION public.agent_runtime_transition(p_tenant text,p_state text,p_incident uuid,p_approval uuid,p_expected_fence bigint) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $old_transition$ DECLARE v_fence bigint; BEGIN IF p_state NOT IN ('degraded','paused','active') THEN RAISE EXCEPTION 'runtime state invalid'; END IF; IF p_state='active' AND (p_approval IS NULL OR NOT EXISTS(SELECT 1 FROM governance_tasks WHERE uid=p_approval AND status='approved' AND due_at>clock_timestamp())) THEN RAISE EXCEPTION 'runtime recovery approval required'; END IF; IF p_state IN ('degraded','paused') AND p_incident IS NULL THEN RAISE EXCEPTION 'runtime incident required'; END IF; UPDATE agent_runtime_states SET state=p_state,incident_ref=COALESCE(p_incident::text,incident_ref),approval_ref=p_approval::text,lease_fence=lease_fence+1,updated_at=clock_timestamp() WHERE tenant_id=p_tenant AND lease_fence=p_expected_fence RETURNING lease_fence INTO v_fence; IF NOT FOUND THEN RAISE EXCEPTION 'runtime state fence conflict'; END IF; RETURN v_fence; END; $old_transition$;
    CREATE OR REPLACE FUNCTION public.agent_runtime_canary_transition(p_tenant text,p_generation text,p_status text,p_approval uuid,p_expected_fence bigint,p_metrics jsonb,p_dataset_version text) RETURNS bigint LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $old_canary$ DECLARE v_fence bigint; BEGIN IF p_status NOT IN ('approved','rolled_back') OR p_dataset_version !~ '^[A-Za-z0-9._-]{1,120}$' OR jsonb_typeof(p_metrics)<>'object' THEN RAISE EXCEPTION 'canary evidence invalid'; END IF; IF p_approval IS NULL OR NOT EXISTS(SELECT 1 FROM governance_tasks WHERE uid=p_approval AND status='approved' AND due_at>clock_timestamp()) THEN RAISE EXCEPTION 'canary approval required'; END IF; UPDATE agent_generation_canaries SET status=p_status,approval_ref=p_approval::text,lease_fence=lease_fence+1 WHERE tenant_id=p_tenant AND generation=p_generation AND lease_fence=p_expected_fence RETURNING lease_fence INTO v_fence; IF NOT FOUND THEN RAISE EXCEPTION 'generation fence conflict'; END IF; RETURN v_fence; END; $old_canary$;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_transition(text,text,uuid,uuid,bigint),public.agent_runtime_canary_transition(text,text,text,uuid,bigint,jsonb,text) TO dataops_agent_runtime,dataops_agent_runtime_owner;
    """)