DataOps-platform/migrations/versions/20260817_506_tenant_controlled_quota.py

"""Bind decimal quota reserve, settle, and release to consumed control claims."""

from alembic import op

revision = "20260817_506"
down_revision = "20260817_505"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute(
        r"""
        ALTER TABLE public.tenant_control_claims DROP CONSTRAINT tenant_control_claims_action_check;
        ALTER TABLE public.tenant_control_claims ADD CONSTRAINT tenant_control_claims_action_check
          CHECK (action IN ('quota_reserve','quota_settle','quota_release','lifecycle_transition'));

        CREATE OR REPLACE FUNCTION public.tenant_control_issue_claim(p_payload jsonb)
        RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $tenant_claim_issue$
        DECLARE v_principal uuid; v_host text; v_action text; v_tenant text; v_claim uuid; v_nonce uuid;
        DECLARE v_request jsonb; v_digest text; v_expires timestamptz;
        BEGIN
          IF NOT pg_has_role(session_user,'dataops_tenant_control','MEMBER') THEN RAISE EXCEPTION 'tenant_control_identity_required'; END IF;
          IF jsonb_typeof(p_payload)<>'object' OR NOT (p_payload ?& ARRAY['claim_uid','nonce','principal_id','host','action','request'])
             OR p_payload-ARRAY['claim_uid','nonce','principal_id','host','action','request']<>'{}'::jsonb THEN RAISE EXCEPTION 'tenant_payload_closed'; END IF;
          IF coalesce(p_payload->>'claim_uid','') !~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
             OR coalesce(p_payload->>'nonce','') !~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
             OR coalesce(p_payload->>'principal_id','') !~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
             OR p_payload->>'host' !~ '^[a-z0-9][a-z0-9.-]{0,252}$' OR p_payload->>'host' <> lower(p_payload->>'host')
             OR p_payload->>'action' NOT IN ('quota_reserve','quota_settle','quota_release','lifecycle_transition')
             OR jsonb_typeof(p_payload->'request')<>'object' THEN RAISE EXCEPTION 'tenant_payload_invalid'; END IF;
          v_claim := (p_payload->>'claim_uid')::uuid; v_nonce := (p_payload->>'nonce')::uuid;
          v_principal := (p_payload->>'principal_id')::uuid; v_host := p_payload->>'host';
          v_action := p_payload->>'action'; v_request := p_payload->'request';
          SELECT m.tenant_id INTO v_tenant FROM public.tenant_control_memberships m JOIN public.tenants t ON t.tenant_id=m.tenant_id
           WHERE m.principal_id=v_principal AND m.host=v_host AND t.state='active';
          IF NOT FOUND THEN RAISE EXCEPTION 'tenant_membership_denied'; END IF;
          v_digest := encode(sha256(convert_to(v_request::text,'utf8')),'hex'); v_expires := clock_timestamp()+interval '120 seconds';
          INSERT INTO public.tenant_control_claims(claim_uid,nonce,tenant_id,principal_id,host,action,request,request_digest,expires_at)
            VALUES(v_claim,v_nonce,v_tenant,v_principal,v_host,v_action,v_request,v_digest,v_expires);
          PERFORM set_config('dataops.tenant_id',v_tenant,true);
          INSERT INTO public.tenant_audit_events(tenant_id,event_type,payload_digest,lease_fence) VALUES(v_tenant,'tenant.claim.issued',v_digest,0);
          RETURN jsonb_build_object('claim_uid',v_claim,'tenant_id',v_tenant,'action',v_action,'request_digest',v_digest,'expires_at',v_expires);
        END; $tenant_claim_issue$;

        CREATE OR REPLACE FUNCTION public.tenant_control_consume_claim(p_claim_uid uuid,p_nonce uuid,p_action text,p_request jsonb)
        RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $tenant_claim_consume$
        DECLARE v_claim public.tenant_control_claims%ROWTYPE;
        BEGIN
          IF NOT pg_has_role(session_user,'dataops_tenant_control','MEMBER') THEN RAISE EXCEPTION 'tenant_control_identity_required'; END IF;
          IF p_action NOT IN ('quota_reserve','quota_settle','quota_release','lifecycle_transition') OR jsonb_typeof(p_request)<>'object' THEN RAISE EXCEPTION 'tenant_payload_closed'; END IF;
          SELECT * INTO v_claim FROM public.tenant_control_claims WHERE claim_uid=p_claim_uid FOR UPDATE;
          IF NOT FOUND OR v_claim.nonce<>p_nonce OR v_claim.action<>p_action OR v_claim.request<>p_request THEN RAISE EXCEPTION 'tenant_claim_denied'; END IF;
          IF v_claim.status<>'issued' OR v_claim.expires_at<=clock_timestamp() THEN RAISE EXCEPTION 'tenant_claim_replayed_or_expired'; END IF;
          UPDATE public.tenant_control_claims SET status='consumed',consumed_at=clock_timestamp() WHERE claim_uid=v_claim.claim_uid AND status='issued';
          IF NOT FOUND THEN RAISE EXCEPTION 'tenant_claim_replayed_or_expired'; END IF;
          PERFORM set_config('dataops.tenant_id',v_claim.tenant_id,true);
          INSERT INTO public.tenant_audit_events(tenant_id,event_type,payload_digest,lease_fence) VALUES(v_claim.tenant_id,'tenant.claim.consumed',v_claim.request_digest,0);
          RETURN jsonb_build_object('claim_uid',v_claim.claim_uid,'tenant_id',v_claim.tenant_id,'principal_id',v_claim.principal_id,'action',v_claim.action,'request_digest',v_claim.request_digest);
        END; $tenant_claim_consume$;

        CREATE FUNCTION public.tenant_control_quota_mutate(p_claim_uid uuid,p_nonce uuid,p_action text,p_request jsonb)
        RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $tenant_quota_control$
        DECLARE v_claim jsonb; v_tenant text; v_quota record; v_reservation record; v_amount numeric(30,6); v_fence bigint; v_digest text;
        BEGIN
          IF p_action NOT IN ('quota_reserve','quota_settle','quota_release') OR jsonb_typeof(p_request)<>'object' THEN RAISE EXCEPTION 'tenant_payload_closed'; END IF;
          IF (p_action='quota_reserve' AND (NOT (p_request ?& ARRAY['quota_name','amount','idempotency_key']) OR p_request-ARRAY['quota_name','amount','idempotency_key']<>'{}'::jsonb OR jsonb_typeof(p_request->'amount')<>'string'))
             OR (p_action IN ('quota_settle','quota_release') AND (NOT (p_request ?& ARRAY['quota_name','idempotency_key','lease_fence']) OR p_request-ARRAY['quota_name','idempotency_key','lease_fence']<>'{}'::jsonb OR jsonb_typeof(p_request->'lease_fence')<>'string')) THEN RAISE EXCEPTION 'tenant_payload_closed'; END IF;
          IF p_request->>'quota_name' NOT IN ('records','storage_objects','background_tasks','events','models','plugins','connectors','api_concurrency')
             OR p_request->>'idempotency_key' !~ '^[A-Za-z0-9._:-]{1,120}$'
             OR (p_action='quota_reserve' AND (p_request->>'amount' !~ '^[0-9]+([.][0-9]{1,6})?$' OR (p_request->>'amount')::numeric<=0 OR (p_request->>'amount')::numeric>1000000000000000000000000::numeric))
             OR (p_action IN ('quota_settle','quota_release') AND p_request->>'lease_fence' !~ '^[1-9][0-9]{0,18}$') THEN RAISE EXCEPTION 'tenant_payload_invalid'; END IF;
          v_claim := public.tenant_control_consume_claim(p_claim_uid,p_nonce,p_action,p_request);
          v_tenant := v_claim->>'tenant_id'; v_digest := v_claim->>'request_digest';
          PERFORM set_config('dataops.tenant_id',v_tenant,true);
          PERFORM pg_advisory_xact_lock(hashtext(v_tenant || '-quota-' || (p_request->>'quota_name')));
          IF p_action='quota_reserve' THEN
            v_amount := (p_request->>'amount')::numeric;
            SELECT * INTO v_reservation FROM public.tenant_quota_reservations WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' AND idempotency_key=p_request->>'idempotency_key' FOR UPDATE;
            IF FOUND THEN
              IF v_reservation.request_digest<>v_digest THEN RAISE EXCEPTION 'tenant_replay_conflict'; END IF;
              RETURN jsonb_build_object('tenant_id',v_tenant,'lease_fence',v_reservation.lease_fence,'status',v_reservation.status);
            END IF;
            SELECT * INTO v_quota FROM public.tenant_quotas WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' FOR UPDATE;
            IF NOT FOUND OR NOT EXISTS(SELECT 1 FROM public.tenants WHERE tenant_id=v_tenant AND state='active') THEN RAISE EXCEPTION 'tenant_not_active'; END IF;
            IF v_quota.limit_units-v_quota.reserved_units<v_amount THEN RAISE EXCEPTION 'tenant_quota_exhausted'; END IF;
            UPDATE public.tenant_quotas SET reserved_units=reserved_units+v_amount,lease_fence=lease_fence+1,updated_at=clock_timestamp() WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' AND lease_fence=v_quota.lease_fence RETURNING lease_fence INTO v_fence;
            IF NOT FOUND THEN RAISE EXCEPTION 'tenant_fence_conflict'; END IF;
            INSERT INTO public.tenant_quota_reservations(tenant_id,quota_name,idempotency_key,request_digest,amount,lease_fence,status) VALUES(v_tenant,p_request->>'quota_name',p_request->>'idempotency_key',v_digest,v_amount,v_fence,'reserved');
            INSERT INTO public.tenant_audit_events(tenant_id,event_type,payload_digest,lease_fence) VALUES(v_tenant,'quota.reserve',v_digest,v_fence);
            RETURN jsonb_build_object('tenant_id',v_tenant,'lease_fence',v_fence,'status','reserved');
          END IF;
          SELECT * INTO v_reservation FROM public.tenant_quota_reservations WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' AND idempotency_key=p_request->>'idempotency_key' FOR UPDATE;
          IF NOT FOUND OR v_reservation.lease_fence<>(p_request->>'lease_fence')::bigint THEN RAISE EXCEPTION 'tenant_fence_conflict'; END IF;
          IF v_reservation.status<> 'reserved' THEN
            IF (p_action='quota_settle' AND v_reservation.status='settled') OR (p_action='quota_release' AND v_reservation.status='released') THEN RETURN jsonb_build_object('tenant_id',v_tenant,'lease_fence',v_reservation.lease_fence,'status',v_reservation.status); END IF;
            RAISE EXCEPTION 'tenant_quota_terminal';
          END IF;
          IF p_action='quota_settle' THEN
            UPDATE public.tenant_quota_reservations SET status='settled',updated_at=clock_timestamp() WHERE reservation_uid=v_reservation.reservation_uid;
            INSERT INTO public.tenant_audit_events(tenant_id,event_type,payload_digest,lease_fence) VALUES(v_tenant,'quota.settle',v_digest,v_reservation.lease_fence);
            RETURN jsonb_build_object('tenant_id',v_tenant,'lease_fence',v_reservation.lease_fence,'status','settled');
          END IF;
          SELECT * INTO v_quota FROM public.tenant_quotas WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' FOR UPDATE;
          UPDATE public.tenant_quotas SET reserved_units=reserved_units-v_reservation.amount,lease_fence=lease_fence+1,updated_at=clock_timestamp() WHERE tenant_id=v_tenant AND quota_name=p_request->>'quota_name' AND lease_fence=v_quota.lease_fence RETURNING lease_fence INTO v_fence;
          IF NOT FOUND THEN RAISE EXCEPTION 'tenant_fence_conflict'; END IF;
          UPDATE public.tenant_quota_reservations SET status='released',updated_at=clock_timestamp() WHERE reservation_uid=v_reservation.reservation_uid;
          INSERT INTO public.tenant_audit_events(tenant_id,event_type,payload_digest,lease_fence) VALUES(v_tenant,'quota.release',v_digest,v_fence);
          RETURN jsonb_build_object('tenant_id',v_tenant,'lease_fence',v_fence,'status','released');
        END; $tenant_quota_control$;
        ALTER FUNCTION public.tenant_control_issue_claim(jsonb) OWNER TO dataops_tenant_foundation_owner;
        ALTER FUNCTION public.tenant_control_consume_claim(uuid,uuid,text,jsonb) OWNER TO dataops_tenant_foundation_owner;
        ALTER FUNCTION public.tenant_control_quota_mutate(uuid,uuid,text,jsonb) OWNER TO dataops_tenant_foundation_owner;
        REVOKE ALL ON FUNCTION public.tenant_control_quota_mutate(uuid,uuid,text,jsonb) FROM PUBLIC,dataops_app_runtime;
        GRANT EXECUTE ON FUNCTION public.tenant_control_quota_mutate(uuid,uuid,text,jsonb) TO dataops_tenant_control;
        """
    )


def downgrade() -> None:
    op.execute(
        """
        DO $$ BEGIN
          IF EXISTS (SELECT 1 FROM public.tenant_control_claims WHERE action IN ('quota_settle','quota_release')) THEN
            RAISE EXCEPTION 'tenant controlled quota downgrade refused while newer claims exist';
          END IF;
        END $$;
        REVOKE EXECUTE ON FUNCTION public.tenant_control_quota_mutate(uuid,uuid,text,jsonb) FROM dataops_tenant_control;
        DROP FUNCTION public.tenant_control_quota_mutate(uuid,uuid,text,jsonb);
        ALTER TABLE public.tenant_control_claims DROP CONSTRAINT tenant_control_claims_action_check;
        ALTER TABLE public.tenant_control_claims ADD CONSTRAINT tenant_control_claims_action_check
          CHECK (action IN ('quota_reserve','lifecycle_transition'));
        """
    )