Files

163 lines
5.6 KiB
PL/PgSQL
Raw Permalink Normal View History

-- Keep mixed-version writers consistent before backfilling rows that already exist.
CREATE OR REPLACE FUNCTION public.enforce_openai_long_context_billing_extra()
RETURNS TRIGGER
LANGUAGE plpgsql
AS $$
DECLARE
parent_effective_value JSONB;
BEGIN
IF NEW.platform IS DISTINCT FROM 'openai' THEN
RETURN NEW;
END IF;
NEW.extra := COALESCE(NEW.extra, '{}'::jsonb);
IF NEW.parent_account_id IS NOT NULL AND NEW.quota_dimension = 'spark' THEN
SELECT CASE
WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb
WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb
WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean'
THEN parent.extra->'openai_long_context_billing_enabled'
ELSE 'false'::jsonb
END
INTO parent_effective_value
FROM accounts AS parent
WHERE parent.id = NEW.parent_account_id;
NEW.extra := jsonb_set(
NEW.extra,
'{openai_long_context_billing_enabled}',
COALESCE(parent_effective_value, 'false'::jsonb),
true
);
ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled')
AND TG_OP = 'UPDATE'
AND OLD.platform = 'openai'
AND jsonb_typeof(OLD.extra->'openai_long_context_billing_enabled') = 'boolean' THEN
NEW.extra := jsonb_set(
NEW.extra,
'{openai_long_context_billing_enabled}',
OLD.extra->'openai_long_context_billing_enabled',
true
);
ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled') THEN
NEW.extra := jsonb_set(
NEW.extra,
'{openai_long_context_billing_enabled}',
'false'::jsonb,
true
);
END IF;
IF jsonb_typeof(NEW.extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean' THEN
RAISE EXCEPTION 'openai_long_context_billing_enabled must be a boolean'
USING ERRCODE = '22023';
END IF;
RETURN NEW;
END;
$$;
CREATE OR REPLACE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows()
RETURNS TRIGGER
LANGUAGE plpgsql
AS $$
BEGIN
WITH updated_shadows AS (
UPDATE accounts AS shadow
SET extra = jsonb_set(
COALESCE(shadow.extra, '{}'::jsonb),
'{openai_long_context_billing_enabled}',
NEW.extra->'openai_long_context_billing_enabled',
true
)
WHERE shadow.parent_account_id = NEW.id
AND shadow.platform = 'openai'
AND shadow.quota_dimension = 'spark'
AND shadow.extra->'openai_long_context_billing_enabled'
IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled'
RETURNING shadow.id
)
INSERT INTO scheduler_outbox (event_type, account_id)
SELECT 'account_changed', id
FROM updated_shadows;
RETURN NULL;
END;
$$;
DROP TRIGGER IF EXISTS accounts_enforce_openai_long_context_billing_extra ON accounts;
CREATE TRIGGER accounts_enforce_openai_long_context_billing_extra
BEFORE INSERT OR UPDATE OF platform, extra, parent_account_id, quota_dimension
ON accounts
FOR EACH ROW
EXECUTE FUNCTION public.enforce_openai_long_context_billing_extra();
DROP TRIGGER IF EXISTS accounts_propagate_openai_long_context_billing_extra ON accounts;
CREATE TRIGGER accounts_propagate_openai_long_context_billing_extra
AFTER UPDATE OF platform, extra
ON accounts
FOR EACH ROW
WHEN (
NEW.platform = 'openai'
AND NEW.parent_account_id IS NULL
AND (
OLD.platform IS DISTINCT FROM NEW.platform
OR OLD.extra->'openai_long_context_billing_enabled'
IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled'
)
)
EXECUTE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows();
UPDATE accounts
SET extra = jsonb_set(
COALESCE(extra, '{}'::jsonb),
'{openai_long_context_billing_enabled}',
'false'::jsonb,
true
)
WHERE platform = 'openai'
AND COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled'
AND jsonb_typeof(extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean';
UPDATE accounts
SET extra = jsonb_set(
COALESCE(extra, '{}'::jsonb),
'{openai_long_context_billing_enabled}',
'false'::jsonb,
true
)
WHERE platform = 'openai'
AND parent_account_id IS NULL
AND NOT (COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled');
WITH shadow_values AS (
SELECT
shadow.id,
CASE
WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb
WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb
WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean'
THEN parent.extra->'openai_long_context_billing_enabled'
ELSE 'false'::jsonb
END AS effective_value
FROM accounts AS shadow
JOIN accounts AS parent ON parent.id = shadow.parent_account_id
WHERE shadow.platform = 'openai'
AND shadow.quota_dimension = 'spark'
),
updated_shadows AS (
UPDATE accounts AS shadow
SET extra = jsonb_set(
COALESCE(shadow.extra, '{}'::jsonb),
'{openai_long_context_billing_enabled}',
shadow_values.effective_value,
true
)
FROM shadow_values
WHERE shadow.id = shadow_values.id
AND shadow.extra->'openai_long_context_billing_enabled'
IS DISTINCT FROM shadow_values.effective_value
RETURNING shadow.id
)
INSERT INTO scheduler_outbox (event_type, account_id)
SELECT 'account_changed', id
FROM updated_shadows;