132 lines
4.7 KiB
PL/PgSQL
132 lines
4.7 KiB
PL/PgSQL
-- ============================================================================
|
|
-- consume_quota: period-aware base/overage consumption
|
|
-- ============================================================================
|
|
-- The original consume_quota (20260412000001) compared only TODAY's
|
|
-- daily_usage row with the base limit. For weekly policies (free tier:
|
|
-- llm_haiku 250/week) that meant a user whose 7-day total had already reached
|
|
-- the limit kept taking the 'base' branch as long as today's own count stayed
|
|
-- under 250, so an overage credit (e.g. one rewarded ad) let every call through
|
|
-- without ever being deducted.
|
|
--
|
|
-- This version takes the policy period and sums the same window the other SQL
|
|
-- quota functions use (reserve_stt_quota, meeting-document generation):
|
|
-- weekly -> CURRENT_DATE - 6 .. CURRENT_DATE (7 calendar days)
|
|
-- daily -> CURRENT_DATE
|
|
-- and chooses base vs overage from that window total, under a per-user/feature
|
|
-- advisory lock plus the subscription row lock so concurrent calls cannot both
|
|
-- take the last base unit or the last credit.
|
|
--
|
|
-- Additional corrections:
|
|
-- * limit 0 ("not available") is denied instead of spending overage credits.
|
|
-- * the overage branch only succeeds when a credit is actually left
|
|
-- (`overage_credits > 0` in the UPDATE), so a race can never go negative.
|
|
-- * 'current' is the window total after this call (equal to today's count
|
|
-- for daily policies), matching what checkQuota reports.
|
|
--
|
|
-- p_period defaults to 'daily' so an edge function still calling the old
|
|
-- three-argument form keeps its previous (daily) behaviour until redeployed.
|
|
-- Apply this migration before deploying the quota.ts that passes p_period.
|
|
|
|
DROP FUNCTION IF EXISTS public.consume_quota(uuid, text, integer);
|
|
|
|
CREATE OR REPLACE FUNCTION public.consume_quota(
|
|
p_user_id uuid,
|
|
p_feature text,
|
|
p_base_limit integer,
|
|
p_period text DEFAULT 'daily'
|
|
) RETURNS jsonb
|
|
LANGUAGE plpgsql
|
|
SECURITY DEFINER
|
|
SET search_path = public, pg_temp
|
|
AS $$
|
|
DECLARE
|
|
v_window_start date;
|
|
v_current integer := 0;
|
|
v_overage integer := 0;
|
|
v_new_overage integer;
|
|
v_consumed_from text;
|
|
BEGIN
|
|
IF p_user_id IS NULL
|
|
OR p_feature IS NULL
|
|
OR p_base_limit IS NULL
|
|
OR p_base_limit < -1
|
|
OR p_period IS NULL
|
|
OR p_period NOT IN ('daily', 'weekly') THEN
|
|
RAISE EXCEPTION 'invalid_quota_consumption' USING ERRCODE = '22023';
|
|
END IF;
|
|
|
|
-- Serialise read-then-increment for this user/feature.
|
|
PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || p_feature, 20260928));
|
|
|
|
SELECT coalesce(overage_credits, 0) INTO v_overage
|
|
FROM public.subscriptions
|
|
WHERE user_id = p_user_id
|
|
FOR UPDATE;
|
|
v_overage := coalesce(v_overage, 0);
|
|
|
|
-- Not available: never spend credits on a feature the tier does not include.
|
|
IF p_base_limit = 0 THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'current', 0,
|
|
'limit', 0,
|
|
'overage_credits', v_overage,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
|
|
v_window_start := CASE WHEN p_period = 'weekly' THEN CURRENT_DATE - 6 ELSE CURRENT_DATE END;
|
|
|
|
SELECT coalesce(sum(count), 0)::integer INTO v_current
|
|
FROM public.daily_usage
|
|
WHERE user_id = p_user_id
|
|
AND feature = p_feature
|
|
AND date >= v_window_start
|
|
AND date <= CURRENT_DATE;
|
|
|
|
IF p_base_limit = -1 THEN
|
|
v_consumed_from := 'unlimited';
|
|
ELSIF v_current < p_base_limit THEN
|
|
v_consumed_from := 'base';
|
|
ELSE
|
|
UPDATE public.subscriptions
|
|
SET overage_credits = overage_credits - 1,
|
|
updated_at = now()
|
|
WHERE user_id = p_user_id
|
|
AND overage_credits > 0
|
|
RETURNING overage_credits INTO v_new_overage;
|
|
|
|
IF NOT FOUND THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'current', v_current,
|
|
'limit', p_base_limit,
|
|
'overage_credits', 0,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
|
|
v_overage := v_new_overage;
|
|
v_consumed_from := 'overage';
|
|
END IF;
|
|
|
|
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
|
VALUES (p_user_id, CURRENT_DATE, p_feature, 1)
|
|
ON CONFLICT (user_id, date, feature) DO UPDATE
|
|
SET count = public.daily_usage.count + 1;
|
|
|
|
RETURN jsonb_build_object(
|
|
'allowed', true,
|
|
'current', v_current + 1,
|
|
'limit', p_base_limit,
|
|
'overage_credits', v_overage,
|
|
'consumed_from', v_consumed_from
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
COMMENT ON FUNCTION public.consume_quota(uuid, text, integer, text) IS
|
|
'Atomic quota consumption over the policy window (daily = today, weekly = CURRENT_DATE-6..CURRENT_DATE): base first, then one overage credit. Used by llm-proxy and realtime-token.';
|
|
|
|
REVOKE ALL ON FUNCTION public.consume_quota(uuid, text, integer, text) FROM PUBLIC, anon, authenticated;
|
|
GRANT EXECUTE ON FUNCTION public.consume_quota(uuid, text, integer, text) TO service_role;
|