283 lines
9 KiB
PL/PgSQL
283 lines
9 KiB
PL/PgSQL
-- Reserve STT quota before provider work so concurrent requests cannot spend
|
|
-- provider capacity beyond the user's allowance. Failed/expired work refunds
|
|
-- the exact base or overage unit it reserved.
|
|
|
|
CREATE TABLE IF NOT EXISTS public.stt_quota_reservations (
|
|
id uuid PRIMARY KEY,
|
|
user_id uuid NOT NULL REFERENCES auth.users(id) ON DELETE CASCADE,
|
|
feature text NOT NULL DEFAULT 'stt_transcribe' CHECK (feature = 'stt_transcribe'),
|
|
usage_date date NOT NULL DEFAULT CURRENT_DATE,
|
|
consumed_from text NOT NULL CHECK (consumed_from IN ('base', 'overage', 'unlimited')),
|
|
status text NOT NULL DEFAULT 'reserved' CHECK (status IN ('reserved', 'completed', 'released')),
|
|
tier text NOT NULL,
|
|
quota_period text NOT NULL CHECK (quota_period IN ('daily', 'weekly')),
|
|
quota_limit integer NOT NULL,
|
|
current_count integer NOT NULL,
|
|
overage_after integer NOT NULL,
|
|
lease_expires_at timestamptz NOT NULL,
|
|
created_at timestamptz NOT NULL DEFAULT now(),
|
|
finalized_at timestamptz,
|
|
release_reason text
|
|
);
|
|
|
|
ALTER TABLE public.stt_quota_reservations ENABLE ROW LEVEL SECURITY;
|
|
REVOKE ALL ON TABLE public.stt_quota_reservations FROM PUBLIC, anon, authenticated;
|
|
GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE public.stt_quota_reservations TO service_role;
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_stt_quota_reservations_reclaim
|
|
ON public.stt_quota_reservations(user_id, status, lease_expires_at)
|
|
WHERE status = 'reserved';
|
|
|
|
CREATE OR REPLACE FUNCTION public.reserve_stt_quota(
|
|
p_user_id uuid,
|
|
p_reservation_id uuid
|
|
) RETURNS jsonb
|
|
LANGUAGE plpgsql
|
|
SECURITY DEFINER
|
|
SET search_path = public, pg_temp
|
|
AS $$
|
|
DECLARE
|
|
existing public.stt_quota_reservations%ROWTYPE;
|
|
expired public.stt_quota_reservations%ROWTYPE;
|
|
subscription_tier text := 'free';
|
|
overage integer := 0;
|
|
quota_period text;
|
|
quota_limit integer;
|
|
current_count integer := 0;
|
|
new_count integer := 0;
|
|
new_overage integer := 0;
|
|
consumed_from text;
|
|
BEGIN
|
|
IF p_user_id IS NULL OR p_reservation_id IS NULL THEN
|
|
RAISE EXCEPTION 'invalid_stt_quota_reservation' USING ERRCODE = '22023';
|
|
END IF;
|
|
|
|
PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':stt_transcribe', 26026));
|
|
|
|
SELECT * INTO existing
|
|
FROM public.stt_quota_reservations
|
|
WHERE id = p_reservation_id
|
|
FOR UPDATE;
|
|
|
|
IF FOUND THEN
|
|
IF existing.user_id <> p_user_id OR existing.feature <> 'stt_transcribe' THEN
|
|
RAISE EXCEPTION 'stt_quota_reservation_conflict' USING ERRCODE = 'PT409';
|
|
END IF;
|
|
RETURN jsonb_build_object(
|
|
'allowed', existing.status IN ('reserved', 'completed'),
|
|
'reservation_id', existing.id,
|
|
'status', existing.status,
|
|
'current', existing.current_count,
|
|
'limit', existing.quota_limit,
|
|
'period', existing.quota_period,
|
|
'tier', existing.tier,
|
|
'overage_credits', existing.overage_after,
|
|
'consumed_from', existing.consumed_from
|
|
);
|
|
END IF;
|
|
|
|
-- Reclaim crashed requests before calculating the next allowance.
|
|
FOR expired IN
|
|
SELECT *
|
|
FROM public.stt_quota_reservations
|
|
WHERE user_id = p_user_id
|
|
AND feature = 'stt_transcribe'
|
|
AND status = 'reserved'
|
|
AND lease_expires_at <= now()
|
|
FOR UPDATE
|
|
LOOP
|
|
UPDATE public.daily_usage
|
|
SET count = greatest(count - 1, 0)
|
|
WHERE user_id = expired.user_id
|
|
AND date = expired.usage_date
|
|
AND feature = expired.feature;
|
|
|
|
IF expired.consumed_from = 'overage' THEN
|
|
UPDATE public.subscriptions
|
|
SET overage_credits = overage_credits + 1,
|
|
updated_at = now()
|
|
WHERE user_id = expired.user_id;
|
|
END IF;
|
|
|
|
UPDATE public.stt_quota_reservations
|
|
SET status = 'released', finalized_at = now(), release_reason = 'lease_expired'
|
|
WHERE id = expired.id;
|
|
END LOOP;
|
|
|
|
SELECT coalesce(tier, 'free'), coalesce(overage_credits, 0)
|
|
INTO subscription_tier, overage
|
|
FROM public.subscriptions
|
|
WHERE user_id = p_user_id
|
|
FOR UPDATE;
|
|
|
|
IF NOT FOUND THEN
|
|
subscription_tier := 'free';
|
|
overage := 0;
|
|
END IF;
|
|
|
|
IF subscription_tier = 'free' THEN
|
|
quota_period := 'weekly';
|
|
quota_limit := 250;
|
|
ELSIF subscription_tier IN ('pro', 'pro_plus', 'team', 'enterprise') THEN
|
|
quota_period := 'daily';
|
|
quota_limit := -1;
|
|
ELSE
|
|
quota_period := 'daily';
|
|
quota_limit := 0;
|
|
END IF;
|
|
|
|
IF quota_period = 'weekly' THEN
|
|
SELECT coalesce(sum(count), 0)::integer INTO current_count
|
|
FROM public.daily_usage
|
|
WHERE user_id = p_user_id
|
|
AND feature = 'stt_transcribe'
|
|
AND date >= CURRENT_DATE - 6
|
|
AND date <= CURRENT_DATE;
|
|
ELSE
|
|
SELECT coalesce(count, 0) INTO current_count
|
|
FROM public.daily_usage
|
|
WHERE user_id = p_user_id
|
|
AND feature = 'stt_transcribe'
|
|
AND date = CURRENT_DATE;
|
|
current_count := coalesce(current_count, 0);
|
|
END IF;
|
|
|
|
IF quota_limit = 0 OR (quota_limit > -1 AND current_count >= quota_limit AND overage <= 0) THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'reservation_id', NULL,
|
|
'status', 'denied',
|
|
'current', current_count,
|
|
'limit', quota_limit,
|
|
'period', quota_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', overage,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
|
|
IF quota_limit = -1 THEN
|
|
consumed_from := 'unlimited';
|
|
ELSIF current_count < quota_limit THEN
|
|
consumed_from := 'base';
|
|
ELSE
|
|
consumed_from := 'overage';
|
|
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 new_overage;
|
|
IF NOT FOUND THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'reservation_id', NULL,
|
|
'status', 'denied',
|
|
'current', current_count,
|
|
'limit', quota_limit,
|
|
'period', quota_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', 0,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
overage := new_overage;
|
|
END IF;
|
|
|
|
INSERT INTO public.daily_usage(user_id, date, feature, count)
|
|
VALUES (p_user_id, CURRENT_DATE, 'stt_transcribe', 1)
|
|
ON CONFLICT (user_id, date, feature)
|
|
DO UPDATE SET count = public.daily_usage.count + 1
|
|
RETURNING count INTO new_count;
|
|
|
|
current_count := current_count + 1;
|
|
INSERT INTO public.stt_quota_reservations(
|
|
id, user_id, consumed_from, tier, quota_period, quota_limit,
|
|
current_count, overage_after, lease_expires_at
|
|
) VALUES (
|
|
p_reservation_id, p_user_id, consumed_from, subscription_tier, quota_period, quota_limit,
|
|
current_count, overage, now() + interval '10 minutes'
|
|
);
|
|
|
|
RETURN jsonb_build_object(
|
|
'allowed', true,
|
|
'reservation_id', p_reservation_id,
|
|
'status', 'reserved',
|
|
'current', current_count,
|
|
'limit', quota_limit,
|
|
'period', quota_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', overage,
|
|
'consumed_from', consumed_from
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
CREATE OR REPLACE FUNCTION public.finalize_stt_quota(
|
|
p_reservation_id uuid,
|
|
p_succeeded boolean
|
|
) RETURNS jsonb
|
|
LANGUAGE plpgsql
|
|
SECURITY DEFINER
|
|
SET search_path = public, pg_temp
|
|
AS $$
|
|
DECLARE
|
|
reservation public.stt_quota_reservations%ROWTYPE;
|
|
final_status text;
|
|
BEGIN
|
|
IF p_reservation_id IS NULL OR p_succeeded IS NULL THEN
|
|
RAISE EXCEPTION 'invalid_stt_quota_finalize' USING ERRCODE = '22023';
|
|
END IF;
|
|
|
|
SELECT * INTO reservation
|
|
FROM public.stt_quota_reservations
|
|
WHERE id = p_reservation_id;
|
|
IF NOT FOUND THEN
|
|
RAISE EXCEPTION 'stt_quota_reservation_not_found' USING ERRCODE = 'P0002';
|
|
END IF;
|
|
|
|
PERFORM pg_advisory_xact_lock(hashtextextended(reservation.user_id::text || ':stt_transcribe', 26026));
|
|
SELECT * INTO reservation
|
|
FROM public.stt_quota_reservations
|
|
WHERE id = p_reservation_id
|
|
FOR UPDATE;
|
|
|
|
IF reservation.status <> 'reserved' THEN
|
|
RETURN jsonb_build_object('reservation_id', reservation.id, 'status', reservation.status);
|
|
END IF;
|
|
|
|
IF p_succeeded THEN
|
|
final_status := 'completed';
|
|
ELSE
|
|
UPDATE public.daily_usage
|
|
SET count = greatest(count - 1, 0)
|
|
WHERE user_id = reservation.user_id
|
|
AND date = reservation.usage_date
|
|
AND feature = reservation.feature;
|
|
|
|
IF reservation.consumed_from = 'overage' THEN
|
|
UPDATE public.subscriptions
|
|
SET overage_credits = overage_credits + 1,
|
|
updated_at = now()
|
|
WHERE user_id = reservation.user_id;
|
|
END IF;
|
|
final_status := 'released';
|
|
END IF;
|
|
|
|
UPDATE public.stt_quota_reservations
|
|
SET status = final_status,
|
|
finalized_at = now(),
|
|
release_reason = CASE WHEN p_succeeded THEN NULL ELSE 'provider_failed' END
|
|
WHERE id = reservation.id;
|
|
|
|
RETURN jsonb_build_object('reservation_id', reservation.id, 'status', final_status);
|
|
END;
|
|
$$;
|
|
|
|
REVOKE ALL ON FUNCTION public.reserve_stt_quota(uuid, uuid) FROM PUBLIC, anon, authenticated;
|
|
REVOKE ALL ON FUNCTION public.finalize_stt_quota(uuid, boolean) FROM PUBLIC, anon, authenticated;
|
|
GRANT EXECUTE ON FUNCTION public.reserve_stt_quota(uuid, uuid) TO service_role;
|
|
GRANT EXECUTE ON FUNCTION public.finalize_stt_quota(uuid, boolean) TO service_role;
|
|
|
|
COMMENT ON TABLE public.stt_quota_reservations IS
|
|
'Service-only leases that reserve STT quota before provider work and refund failed or expired work.';
|