163 lines
5.6 KiB
PL/PgSQL
163 lines
5.6 KiB
PL/PgSQL
-- 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;
|