Skip to content

Commit 73e1d92

Browse files
ProgramcodeProgramcode
authored
fix: separate source import quota usage (#50)
* fix: separate source import quota usage * test: use portable quota result parsing
1 parent b2ba178 commit 73e1d92

6 files changed

Lines changed: 406 additions & 3 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
"test": "vitest run",
2020
"test:watch": "vitest",
2121
"test:edge": "deno test --allow-read --allow-env=DEMO_CODE_PEPPER,EXA_API_KEY,SUPABASE_URL,SUPABASE_SERVICE_ROLE_KEY,SPARRING_INTEGRITY_SECRET,OPENROUTER_API_KEY --allow-net=127.0.0.1 supabase/functions/",
22-
"test:db": "supabase test db",
22+
"test:db": "supabase test db && bash scripts/test-daily-quota.sh",
2323
"test:e2e": "playwright test --project=desktop-chromium --project=mobile-chromium",
2424
"test:e2e:webkit": "playwright test --project=webkit",
2525
"test:e2e:staging": "playwright test --project=staging",

‎scripts/test-daily-quota.sh‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
5+
database_url="${SUPABASE_DB_URL:-postgresql://postgres:postgres@127.0.0.1:54322/postgres}"
6+
database_password="${SUPABASE_DB_PASSWORD:-postgres}"
7+
test_user_id="66666666-6666-4666-8666-666666666666"
8+
result_dir="$(mktemp -d)"
9+
first_pid=""
10+
second_pid=""
11+
12+
cleanup() {
13+
touch "$result_dir/release-first"
14+
if [ -n "$first_pid" ]; then
15+
wait "$first_pid" 2>/dev/null || true
16+
fi
17+
if [ -n "$second_pid" ]; then
18+
wait "$second_pid" 2>/dev/null || true
19+
fi
20+
PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 \
21+
--command "DELETE FROM auth.users WHERE id = '$test_user_id'" >/dev/null
22+
rm -rf "$result_dir"
23+
}
24+
trap cleanup EXIT
25+
26+
(
27+
cd "$repo_root"
28+
PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 <<'SQL'
29+
BEGIN;
30+
INSERT INTO auth.users (id)
31+
VALUES ('77777777-7777-4777-8777-777777777777');
32+
DROP FUNCTION public.consume_daily_quota(uuid, integer, text);
33+
CREATE FUNCTION public.consume_daily_quota(uuid, integer)
34+
RETURNS TABLE (allowed boolean, current_count integer)
35+
LANGUAGE sql
36+
AS $$ SELECT false, 0 $$;
37+
ALTER TABLE public.daily_usage_counters DROP CONSTRAINT daily_usage_counters_pkey;
38+
ALTER TABLE public.daily_usage_counters DROP COLUMN quota_key;
39+
ALTER TABLE public.daily_usage_counters
40+
ADD CONSTRAINT daily_usage_counters_pkey PRIMARY KEY (user_id, usage_date);
41+
INSERT INTO public.daily_usage_counters (user_id, usage_date, request_count)
42+
VALUES ('77777777-7777-4777-8777-777777777777', CURRENT_DATE - 1, 7);
43+
\i supabase/migrations/20260920125555_scope_daily_usage_counters.sql
44+
DO $$
45+
BEGIN
46+
IF NOT EXISTS (
47+
SELECT 1
48+
FROM public.daily_usage_counters
49+
WHERE user_id = '77777777-7777-4777-8777-777777777777'
50+
AND usage_date = CURRENT_DATE - 1
51+
AND quota_key = 'ai'
52+
AND request_count = 7
53+
) THEN
54+
RAISE EXCEPTION 'legacy daily usage row was not preserved under the ai quota';
55+
END IF;
56+
END;
57+
$$;
58+
ROLLBACK;
59+
SQL
60+
)
61+
echo "Legacy quota row migration: count 7 preserved under ai"
62+
63+
PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 \
64+
--command "DELETE FROM auth.users WHERE id = '$test_user_id'; INSERT INTO auth.users (id) VALUES ('$test_user_id')" \
65+
>/dev/null
66+
67+
reserve_first() {
68+
PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 <<SQL
69+
BEGIN;
70+
SET LOCAL ROLE service_role;
71+
SELECT set_config('request.jwt.claim.role', 'service_role', true);
72+
SELECT allowed, current_count
73+
FROM public.consume_daily_quota('$test_user_id', 1, 'concurrent');
74+
\! touch "$result_dir/first-holds-lock"
75+
\! while [ ! -e "$result_dir/release-first" ]; do sleep 0.05; done
76+
COMMIT;
77+
SQL
78+
}
79+
80+
reserve_second() {
81+
PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 <<SQL
82+
BEGIN;
83+
SET LOCAL ROLE service_role;
84+
SELECT set_config('request.jwt.claim.role', 'service_role', true);
85+
SELECT allowed, current_count
86+
FROM public.consume_daily_quota('$test_user_id', 1, 'concurrent') /* daily_quota_concurrency_second */;
87+
COMMIT;
88+
SQL
89+
}
90+
91+
reserve_first >"$result_dir/first" &
92+
first_pid=$!
93+
94+
for _ in {1..100}; do
95+
if [ -e "$result_dir/first-holds-lock" ]; then
96+
break
97+
fi
98+
if ! kill -0 "$first_pid" 2>/dev/null; then
99+
wait "$first_pid"
100+
fi
101+
sleep 0.05
102+
done
103+
if [ ! -e "$result_dir/first-holds-lock" ]; then
104+
echo "first quota reservation did not acquire its row lock"
105+
exit 1
106+
fi
107+
108+
reserve_second >"$result_dir/second" &
109+
second_pid=$!
110+
111+
second_is_waiting=false
112+
for _ in {1..100}; do
113+
if [ "$(PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 --command \
114+
"SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid <> pg_backend_pid() AND wait_event_type = 'Lock' AND query LIKE '%daily_quota_concurrency_second%')")" = "t" ]; then
115+
second_is_waiting=true
116+
break
117+
fi
118+
if ! kill -0 "$second_pid" 2>/dev/null; then
119+
break
120+
fi
121+
sleep 0.05
122+
done
123+
if [ "$second_is_waiting" != true ]; then
124+
echo "second quota reservation did not wait on the first row lock"
125+
exit 1
126+
fi
127+
128+
touch "$result_dir/release-first"
129+
130+
first_result=0
131+
second_result=0
132+
wait "$first_pid" || first_result=$?
133+
wait "$second_pid" || second_result=$?
134+
if [ "$first_result" -ne 0 ] || [ "$second_result" -ne 0 ]; then
135+
exit 1
136+
fi
137+
138+
actual_results="$(grep -hE '^[tf]\|[0-9]+$' "$result_dir/first" "$result_dir/second" | sort)"
139+
expected_results=$'f|1\nt|1'
140+
if [ "$actual_results" != "$expected_results" ]; then
141+
echo "unexpected concurrent quota results:"
142+
echo "$actual_results"
143+
exit 1
144+
fi
145+
146+
stored_count="$(PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 --command \
147+
"SELECT request_count FROM public.daily_usage_counters WHERE user_id = '$test_user_id' AND usage_date = CURRENT_DATE AND quota_key = 'concurrent'")"
148+
if [ "$stored_count" != "1" ]; then
149+
echo "expected one stored concurrent reservation, got: $stored_count"
150+
exit 1
151+
fi
152+
153+
echo "Concurrent quota reservations: 1 allowed, 1 rejected, stored count 1"

‎supabase/functions/import-source/index.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ async function consumeDailyImportQuota(userId: string): Promise<{ allowed: boole
4949
const { data, error } = await supabase.rpc("consume_daily_quota", {
5050
p_user_id: userId,
5151
p_limit: DAILY_IMPORT_LIMIT,
52+
p_quota_key: "import_source",
5253
});
5354
if (error) throw error;
5455
const row = Array.isArray(data) ? data[0] : data;
@@ -238,7 +239,7 @@ function normalizeText(text: string) {
238239
.slice(0, MAX_CONTENT_CHARS);
239240
}
240241

241-
Deno.serve(async (req) => {
242+
export async function handleRequest(req: Request): Promise<Response> {
242243
const corsHeaders = getCorsHeaders(req.headers.get("origin"));
243244

244245
if (req.method === "OPTIONS") {
@@ -331,4 +332,8 @@ Deno.serve(async (req) => {
331332
await logImport(user.id, "error", "exception", Date.now() - startedAt);
332333
return jsonResponse({ error: message }, 400, corsHeaders);
333334
}
334-
});
335+
}
336+
337+
if (import.meta.main) {
338+
Deno.serve(handleRequest);
339+
}

‎supabase/functions/import-source/index_test.ts‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
const environment = {
2+
SUPABASE_URL: "https://supabase.test",
3+
SUPABASE_SERVICE_ROLE_KEY: "test-service-role-key",
4+
};
5+
6+
for (const [key, value] of Object.entries(environment)) {
7+
Deno.env.set(key, value);
8+
}
9+
10+
const { handleRequest } = await import("./index.ts");
11+
12+
function jsonResponse(body: unknown, status = 200): Response {
13+
return Response.json(body, { status });
14+
}
15+
16+
Deno.test("import-source reserves only import_source quota", async () => {
17+
const originalFetch = globalThis.fetch;
18+
let quotaRequest: Record<string, unknown> | undefined;
19+
20+
globalThis.fetch =
21+
(async (input: Request | URL | string, init?: RequestInit) => {
22+
const url = new URL(input instanceof Request ? input.url : String(input));
23+
24+
if (url.origin !== "https://supabase.test") {
25+
throw new Error(`Unexpected request: ${url}`);
26+
}
27+
if (url.pathname === "/auth/v1/user") {
28+
return jsonResponse({ id: "user-1", aud: "authenticated" });
29+
}
30+
if (url.pathname.includes("/rest/v1/user_roles")) {
31+
return jsonResponse({ role: "user" });
32+
}
33+
if (url.pathname.includes("/rest/v1/rpc/consume_daily_quota")) {
34+
quotaRequest = JSON.parse(String(init?.body)) as Record<
35+
string,
36+
unknown
37+
>;
38+
return jsonResponse([{ allowed: false, current_count: 50 }]);
39+
}
40+
throw new Error(`Unexpected Supabase request: ${url}`);
41+
}) as typeof fetch;
42+
43+
try {
44+
const response = await handleRequest(
45+
new Request("http://127.0.0.1/functions/v1/import-source", {
46+
method: "POST",
47+
headers: {
48+
Authorization: "Bearer test-user-token",
49+
"Content-Type": "application/json",
50+
},
51+
body: JSON.stringify({ url: "https://example.com/article" }),
52+
}),
53+
);
54+
55+
if (response.status !== 429) {
56+
throw new Error(
57+
`import-source returned ${response.status}: ${await response.text()}`,
58+
);
59+
}
60+
if (quotaRequest?.p_quota_key !== "import_source") {
61+
throw new Error(
62+
`import-source reserved the wrong quota: ${
63+
JSON.stringify(quotaRequest)
64+
}`,
65+
);
66+
}
67+
} finally {
68+
globalThis.fetch = originalFetch;
69+
}
70+
});

‎supabase/migrations/20260920125555_scope_daily_usage_counters.sql‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
-- Keep independently configured daily quotas in separate atomic counters.
2+
ALTER TABLE public.daily_usage_counters
3+
ADD COLUMN quota_key text NOT NULL DEFAULT 'ai';
4+
5+
ALTER TABLE public.daily_usage_counters
6+
DROP CONSTRAINT daily_usage_counters_pkey;
7+
8+
ALTER TABLE public.daily_usage_counters
9+
ADD CONSTRAINT daily_usage_counters_pkey
10+
PRIMARY KEY (user_id, usage_date, quota_key);
11+
12+
-- The new defaulted parameter changes the PostgreSQL function signature, so
13+
-- remove the old overload before creating its replacement.
14+
DROP FUNCTION public.consume_daily_quota(uuid, integer);
15+
16+
CREATE FUNCTION public.consume_daily_quota(
17+
p_user_id uuid,
18+
p_limit integer,
19+
p_quota_key text DEFAULT 'ai'
20+
)
21+
RETURNS TABLE (allowed boolean, current_count integer)
22+
LANGUAGE plpgsql
23+
SECURITY DEFINER
24+
SET search_path = public
25+
AS $$
26+
DECLARE
27+
v_quota_key text := nullif(trim(p_quota_key), '');
28+
v_current_count integer;
29+
BEGIN
30+
IF auth.role() <> 'service_role' OR p_user_id IS NULL OR p_limit < 1 OR v_quota_key IS NULL THEN
31+
RAISE EXCEPTION 'quota reservation is restricted to trusted server code';
32+
END IF;
33+
34+
INSERT INTO public.daily_usage_counters (user_id, usage_date, quota_key, request_count)
35+
VALUES (p_user_id, CURRENT_DATE, v_quota_key, 1)
36+
ON CONFLICT (user_id, usage_date, quota_key) DO UPDATE
37+
SET request_count = public.daily_usage_counters.request_count + 1
38+
WHERE public.daily_usage_counters.request_count < p_limit
39+
RETURNING request_count INTO v_current_count;
40+
41+
IF FOUND THEN
42+
RETURN QUERY SELECT true, v_current_count;
43+
RETURN;
44+
END IF;
45+
46+
SELECT request_count INTO v_current_count
47+
FROM public.daily_usage_counters
48+
WHERE user_id = p_user_id
49+
AND usage_date = CURRENT_DATE
50+
AND quota_key = v_quota_key;
51+
52+
RETURN QUERY SELECT false, COALESCE(v_current_count, 0);
53+
END;
54+
$$;
55+
56+
REVOKE ALL ON FUNCTION public.consume_daily_quota(uuid, integer, text)
57+
FROM PUBLIC, anon, authenticated;
58+
GRANT EXECUTE ON FUNCTION public.consume_daily_quota(uuid, integer, text)
59+
TO service_role;

0 commit comments

Comments
 (0)