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

"""Permit bounded decimal quota reservations through a corrected closed gateway."""

from alembic import op

revision = "20260817_503"
down_revision = "20260817_502"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute(
        r"""
        CREATE FUNCTION public.tenant_quota_reserve(p_payload jsonb)
        RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $tenant_decimal$
        DECLARE v_tenant text; v_digest text; v_amount numeric(30,6); v_quota record; v_reservation record; v_fence bigint;
        BEGIN
          IF jsonb_typeof(p_payload)<>'object' OR NOT (p_payload ?& ARRAY['tenant_id','quota_name','amount','idempotency_key'])
             OR p_payload-ARRAY['tenant_id','quota_name','amount','idempotency_key']<>'{}'::jsonb THEN RAISE EXCEPTION 'tenant_payload_closed'; END IF;
          v_tenant:=p_payload->>'tenant_id';
          IF v_tenant !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->>'quota_name' NOT IN ('records','storage_objects','background_tasks','events','models','plugins','connectors','api_concurrency')
             OR p_payload->>'idempotency_key' !~ '^[A-Za-z0-9._:-]{1,120}$' OR p_payload->>'amount' !~ '^[0-9]+([.][0-9]{1,6})?$' THEN RAISE EXCEPTION 'tenant_payload_invalid'; END IF;
          v_amount:=(p_payload->>'amount')::numeric;
          IF v_amount<=0 OR v_amount>1000000000000000000000000::numeric THEN RAISE EXCEPTION 'tenant_payload_invalid'; END IF;
          v_digest:=encode(sha256(convert_to(p_payload::text,'utf8')),'hex');
          PERFORM set_config('dataops.tenant_id',v_tenant,true);
          PERFORM pg_advisory_xact_lock(hashtext(v_tenant || '-quota-' || (p_payload->>'quota_name')));
          SELECT * INTO v_reservation FROM public.tenant_quota_reservations WHERE tenant_id=v_tenant AND quota_name=p_payload->>'quota_name' AND idempotency_key=p_payload->>'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_payload->>'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_payload->>'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_payload->>'quota_name',p_payload->>'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; $tenant_decimal$;
        ALTER FUNCTION public.tenant_quota_reserve(jsonb) OWNER TO dataops_tenant_foundation_owner;
        REVOKE ALL ON FUNCTION public.tenant_quota_reserve(jsonb) FROM PUBLIC;
        GRANT EXECUTE ON FUNCTION public.tenant_quota_reserve(jsonb) TO dataops_app_runtime;
        """
    )


def downgrade() -> None:
    op.execute("REVOKE EXECUTE ON FUNCTION public.tenant_quota_reserve(jsonb) FROM dataops_app_runtime; DROP FUNCTION IF EXISTS public.tenant_quota_reserve(jsonb);")