175_default_openai_long_context_billing.sql 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. -- Keep mixed-version writers consistent before backfilling rows that already exist.
  2. CREATE OR REPLACE FUNCTION public.enforce_openai_long_context_billing_extra()
  3. RETURNS TRIGGER
  4. LANGUAGE plpgsql
  5. AS $$
  6. DECLARE
  7. parent_effective_value JSONB;
  8. BEGIN
  9. IF NEW.platform IS DISTINCT FROM 'openai' THEN
  10. RETURN NEW;
  11. END IF;
  12. NEW.extra := COALESCE(NEW.extra, '{}'::jsonb);
  13. IF NEW.parent_account_id IS NOT NULL AND NEW.quota_dimension = 'spark' THEN
  14. SELECT CASE
  15. WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb
  16. WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb
  17. WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean'
  18. THEN parent.extra->'openai_long_context_billing_enabled'
  19. ELSE 'false'::jsonb
  20. END
  21. INTO parent_effective_value
  22. FROM accounts AS parent
  23. WHERE parent.id = NEW.parent_account_id;
  24. NEW.extra := jsonb_set(
  25. NEW.extra,
  26. '{openai_long_context_billing_enabled}',
  27. COALESCE(parent_effective_value, 'false'::jsonb),
  28. true
  29. );
  30. ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled')
  31. AND TG_OP = 'UPDATE'
  32. AND OLD.platform = 'openai'
  33. AND jsonb_typeof(OLD.extra->'openai_long_context_billing_enabled') = 'boolean' THEN
  34. NEW.extra := jsonb_set(
  35. NEW.extra,
  36. '{openai_long_context_billing_enabled}',
  37. OLD.extra->'openai_long_context_billing_enabled',
  38. true
  39. );
  40. ELSIF NOT (NEW.extra ? 'openai_long_context_billing_enabled') THEN
  41. NEW.extra := jsonb_set(
  42. NEW.extra,
  43. '{openai_long_context_billing_enabled}',
  44. 'false'::jsonb,
  45. true
  46. );
  47. END IF;
  48. IF jsonb_typeof(NEW.extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean' THEN
  49. RAISE EXCEPTION 'openai_long_context_billing_enabled must be a boolean'
  50. USING ERRCODE = '22023';
  51. END IF;
  52. RETURN NEW;
  53. END;
  54. $$;
  55. CREATE OR REPLACE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows()
  56. RETURNS TRIGGER
  57. LANGUAGE plpgsql
  58. AS $$
  59. BEGIN
  60. WITH updated_shadows AS (
  61. UPDATE accounts AS shadow
  62. SET extra = jsonb_set(
  63. COALESCE(shadow.extra, '{}'::jsonb),
  64. '{openai_long_context_billing_enabled}',
  65. NEW.extra->'openai_long_context_billing_enabled',
  66. true
  67. )
  68. WHERE shadow.parent_account_id = NEW.id
  69. AND shadow.platform = 'openai'
  70. AND shadow.quota_dimension = 'spark'
  71. AND shadow.extra->'openai_long_context_billing_enabled'
  72. IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled'
  73. RETURNING shadow.id
  74. )
  75. INSERT INTO scheduler_outbox (event_type, account_id)
  76. SELECT 'account_changed', id
  77. FROM updated_shadows;
  78. RETURN NULL;
  79. END;
  80. $$;
  81. DROP TRIGGER IF EXISTS accounts_enforce_openai_long_context_billing_extra ON accounts;
  82. CREATE TRIGGER accounts_enforce_openai_long_context_billing_extra
  83. BEFORE INSERT OR UPDATE OF platform, extra, parent_account_id, quota_dimension
  84. ON accounts
  85. FOR EACH ROW
  86. EXECUTE FUNCTION public.enforce_openai_long_context_billing_extra();
  87. DROP TRIGGER IF EXISTS accounts_propagate_openai_long_context_billing_extra ON accounts;
  88. CREATE TRIGGER accounts_propagate_openai_long_context_billing_extra
  89. AFTER UPDATE OF platform, extra
  90. ON accounts
  91. FOR EACH ROW
  92. WHEN (
  93. NEW.platform = 'openai'
  94. AND NEW.parent_account_id IS NULL
  95. AND (
  96. OLD.platform IS DISTINCT FROM NEW.platform
  97. OR OLD.extra->'openai_long_context_billing_enabled'
  98. IS DISTINCT FROM NEW.extra->'openai_long_context_billing_enabled'
  99. )
  100. )
  101. EXECUTE FUNCTION public.propagate_openai_long_context_billing_extra_to_shadows();
  102. UPDATE accounts
  103. SET extra = jsonb_set(
  104. COALESCE(extra, '{}'::jsonb),
  105. '{openai_long_context_billing_enabled}',
  106. 'false'::jsonb,
  107. true
  108. )
  109. WHERE platform = 'openai'
  110. AND COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled'
  111. AND jsonb_typeof(extra->'openai_long_context_billing_enabled') IS DISTINCT FROM 'boolean';
  112. UPDATE accounts
  113. SET extra = jsonb_set(
  114. COALESCE(extra, '{}'::jsonb),
  115. '{openai_long_context_billing_enabled}',
  116. 'false'::jsonb,
  117. true
  118. )
  119. WHERE platform = 'openai'
  120. AND parent_account_id IS NULL
  121. AND NOT (COALESCE(extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled');
  122. WITH shadow_values AS (
  123. SELECT
  124. shadow.id,
  125. CASE
  126. WHEN parent.platform IS DISTINCT FROM 'openai' THEN 'false'::jsonb
  127. WHEN NOT (COALESCE(parent.extra, '{}'::jsonb) ? 'openai_long_context_billing_enabled') THEN 'false'::jsonb
  128. WHEN jsonb_typeof(parent.extra->'openai_long_context_billing_enabled') = 'boolean'
  129. THEN parent.extra->'openai_long_context_billing_enabled'
  130. ELSE 'false'::jsonb
  131. END AS effective_value
  132. FROM accounts AS shadow
  133. JOIN accounts AS parent ON parent.id = shadow.parent_account_id
  134. WHERE shadow.platform = 'openai'
  135. AND shadow.quota_dimension = 'spark'
  136. ),
  137. updated_shadows AS (
  138. UPDATE accounts AS shadow
  139. SET extra = jsonb_set(
  140. COALESCE(shadow.extra, '{}'::jsonb),
  141. '{openai_long_context_billing_enabled}',
  142. shadow_values.effective_value,
  143. true
  144. )
  145. FROM shadow_values
  146. WHERE shadow.id = shadow_values.id
  147. AND shadow.extra->'openai_long_context_billing_enabled'
  148. IS DISTINCT FROM shadow_values.effective_value
  149. RETURNING shadow.id
  150. )
  151. INSERT INTO scheduler_outbox (event_type, account_id)
  152. SELECT 'account_changed', id
  153. FROM updated_shadows;