DataOps-platform/migrations/versions/20260818_547_metering_allocation_budget_gateway.py

"""Persist WP12 allocation and budget controls behind the control gateway."""

from alembic import op

revision = "20260818_547"
down_revision = "20260818_546"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.get_bind().exec_driver_sql(r'''
    CREATE FUNCTION public.metering_showback_control_write(p_action text,p_principal text,p_payload jsonb)
    RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $control$
    DECLARE t text; d text; r text; total integer; item_count integer; valid_count integer;
      used bigint; created boolean:=false; existing_digest text;
    BEGIN
      IF NOT pg_has_role(session_user,'dataops_bi_ai_catalog_issuer','MEMBER') OR p_action NOT IN ('allocation','budget') OR jsonb_typeof(p_payload)<>'object' OR p_principal !~ '^[A-Za-z0-9_.:-]{1,120}$' THEN RAISE EXCEPTION 'metering_control_denied'; END IF;
      SELECT tenant_ref,domain_ref,role_name INTO t,d,r FROM public.metering_scope_grants WHERE principal_ref=p_principal AND active AND revoked_at IS NULL ORDER BY CASE role_name WHEN 'admin' THEN 1 WHEN 'operator' THEN 2 ELSE 3 END LIMIT 1;
      IF NOT FOUND OR r NOT IN ('operator','admin') THEN RAISE EXCEPTION 'metering_control_denied'; END IF;
      IF p_action='allocation' THEN
        IF p_payload-ARRAY['rule_uid','rule_version','effective_start','effective_end','mapping','allocations']<>'{}'::jsonb OR jsonb_typeof(p_payload->'mapping')<>'object' OR jsonb_typeof(p_payload->'allocations')<>'array' OR jsonb_array_length(p_payload->'allocations') NOT BETWEEN 1 AND 32 OR p_payload->>'rule_uid' !~ '^[A-Za-z0-9][A-Za-z0-9._:-]{0,119}$' OR p_payload->>'rule_version' !~ '^[1-9][0-9]{0,6}$' OR (p_payload->'mapping')-ARRAY['department','business_domain','project','cost_center']<>'{}'::jsonb OR p_payload->'mapping'->>'business_domain'<>d OR p_payload->'mapping'->>'department' !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->'mapping'->>'project' !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->'mapping'->>'cost_center' !~ '^[a-z][a-z0-9-]{0,62}$' OR (p_payload->>'effective_start')::timestamptz >= (p_payload->>'effective_end')::timestamptz THEN RAISE EXCEPTION 'metering_allocation_closed'; END IF;
        SELECT count(*), count(*) FILTER (WHERE jsonb_typeof(value)='object' AND value-ARRAY['target','weight_micros']='{}'::jsonb AND value->>'target' ~ '^[a-z][a-z0-9-]{0,62}$' AND value->>'weight_micros' ~ '^[0-9]{1,7}$' AND (value->>'weight_micros')::integer<=1000000), COALESCE(sum(CASE WHEN jsonb_typeof(value)='object' AND value-ARRAY['target','weight_micros']='{}'::jsonb AND value->>'target' ~ '^[a-z][a-z0-9-]{0,62}$' AND value->>'weight_micros' ~ '^[0-9]{1,7}$' AND (value->>'weight_micros')::integer<=1000000 THEN (value->>'weight_micros')::integer ELSE 0 END),0) INTO item_count,valid_count,total FROM jsonb_array_elements(p_payload->'allocations');
        IF item_count<>valid_count OR total<>1000000 THEN RAISE EXCEPTION 'metering_allocation_weight_invalid'; END IF;
        SELECT rule_digest INTO existing_digest FROM public.metering_allocation_rules WHERE rule_uid=p_payload->>'rule_uid' AND rule_version=(p_payload->>'rule_version')::integer;
        IF existing_digest IS NOT NULL AND existing_digest<>encode(sha256(convert_to(p_payload::text,'utf8')),'hex') THEN RAISE EXCEPTION 'metering_allocation_conflict'; END IF;
        INSERT INTO public.metering_allocation_rules(rule_uid,rule_version,tenant_ref,domain_ref,department_ref,project_ref,cost_center_ref,effective_start,effective_end,rule_digest) VALUES(p_payload->>'rule_uid',(p_payload->>'rule_version')::integer,t,d,p_payload->'mapping'->>'department',p_payload->'mapping'->>'project',p_payload->'mapping'->>'cost_center',(p_payload->>'effective_start')::timestamptz,(p_payload->>'effective_end')::timestamptz,encode(sha256(convert_to(p_payload::text,'utf8')),'hex')) ON CONFLICT(rule_uid,rule_version) DO NOTHING;
        INSERT INTO public.metering_allocations(rule_uid,rule_version,target_ref,weight_micros) SELECT p_payload->>'rule_uid',(p_payload->>'rule_version')::integer,value->>'target',(value->>'weight_micros')::integer FROM jsonb_array_elements(p_payload->'allocations') ON CONFLICT DO NOTHING;
        RETURN jsonb_build_object('rule_uid',p_payload->>'rule_uid','rule_version',(p_payload->>'rule_version')::integer,'persisted_before_ack',true);
      END IF;
      IF p_payload-ARRAY['budget_uid','window','mapping','limit_micros','threshold_micros']<>'{}'::jsonb OR jsonb_typeof(p_payload->'mapping')<>'object' OR (p_payload->'mapping')-ARRAY['department','business_domain','project','cost_center']<>'{}'::jsonb OR p_payload->'mapping'->>'business_domain'<>d OR p_payload->'mapping'->>'department' !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->'mapping'->>'project' !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->'mapping'->>'cost_center' !~ '^[a-z][a-z0-9-]{0,62}$' OR p_payload->>'budget_uid' !~ '^[A-Za-z0-9][A-Za-z0-9._:-]{0,119}$' OR p_payload->>'window' !~ '^[0-9]{4}-(0[1-9]|1[0-2])$' OR p_payload->>'limit_micros' !~ '^[0-9]{1,18}$' OR p_payload->>'threshold_micros' !~ '^[0-9]{1,18}$' OR (p_payload->>'threshold_micros')::bigint>(p_payload->>'limit_micros')::bigint THEN RAISE EXCEPTION 'metering_budget_closed'; END IF;
      IF EXISTS(SELECT 1 FROM public.metering_budgets WHERE budget_uid=p_payload->>'budget_uid' AND (tenant_ref,domain_ref,department_ref,project_ref,cost_center_ref,window_key,limit_micros,threshold_micros) IS DISTINCT FROM (t,d,p_payload->'mapping'->>'department',p_payload->'mapping'->>'project',p_payload->'mapping'->>'cost_center',p_payload->>'window',(p_payload->>'limit_micros')::bigint,(p_payload->>'threshold_micros')::bigint)) THEN RAISE EXCEPTION 'metering_budget_conflict'; END IF;
      INSERT INTO public.metering_budgets(budget_uid,tenant_ref,domain_ref,department_ref,project_ref,cost_center_ref,window_key,limit_micros,threshold_micros) VALUES(p_payload->>'budget_uid',t,d,p_payload->'mapping'->>'department',p_payload->'mapping'->>'project',p_payload->'mapping'->>'cost_center',p_payload->>'window',(p_payload->>'limit_micros')::bigint,(p_payload->>'threshold_micros')::bigint) ON CONFLICT(budget_uid) DO NOTHING;
      SELECT COALESCE(sum(quantity_micros),0) INTO used FROM public.metering_events WHERE tenant_ref=t AND domain_ref=d AND department_ref=p_payload->'mapping'->>'department' AND project_ref=p_payload->'mapping'->>'project' AND cost_center_ref=p_payload->'mapping'->>'cost_center' AND to_char(window_start AT TIME ZONE 'UTC','YYYY-MM')=p_payload->>'window';
      IF used >= (p_payload->>'threshold_micros')::bigint THEN WITH written AS (INSERT INTO public.metering_alert_outbox(budget_uid,window_key,threshold_micros,payload_digest) VALUES(p_payload->>'budget_uid',p_payload->>'window',(p_payload->>'threshold_micros')::bigint,encode(sha256(convert_to(p_payload::text,'utf8')),'hex')) ON CONFLICT DO NOTHING RETURNING 1) SELECT EXISTS(SELECT 1 FROM written) INTO created; END IF;
      RETURN jsonb_build_object('budget_uid',p_payload->>'budget_uid','used_micros',used,'alert_created',COALESCE(created,false),'provider','disabled');
    END; $control$;
    ALTER FUNCTION public.metering_showback_control_write(text,text,jsonb) OWNER TO dataops_tenant_foundation_owner;
    REVOKE ALL ON FUNCTION public.metering_showback_control_write(text,text,jsonb) FROM PUBLIC,dataops_app,dataops_app_runtime;
    GRANT EXECUTE ON FUNCTION public.metering_showback_control_write(text,text,jsonb) TO dataops_bi_ai_catalog_control;
    ''')


def downgrade() -> None:
    if op.get_bind().exec_driver_sql("SELECT EXISTS(SELECT 1 FROM public.metering_allocation_rules LIMIT 1) OR EXISTS(SELECT 1 FROM public.metering_budgets LIMIT 1) OR EXISTS(SELECT 1 FROM public.metering_alert_outbox LIMIT 1)").scalar():
        raise RuntimeError("downgrade refused: WP12 allocation or budget facts are nonempty")
    op.execute("REVOKE EXECUTE ON FUNCTION public.metering_showback_control_write(text,text,jsonb) FROM dataops_bi_ai_catalog_control; DROP FUNCTION public.metering_showback_control_write(text,text,jsonb);")