From 26d36c4995950adaadabf4a801c0e959ad7b6287 Mon Sep 17 00:00:00 2001 From: johnyeo Date: Fri, 22 May 2026 14:13:29 +0100 Subject: [PATCH] fix: optimize product counts --- server/experiments/explainGetProductCounts.ts | 236 ++++++++++++++++++ server/experiments/explainGetRewards.ts | 151 +++++++++++ .../out/product-counts-explain-1.txt | 6 + .../out/product-counts-explain-2.txt | 9 + server/experiments/out/rewards-explain.txt | 15 ++ .../src/internal/admin/handleListAdminOrgs.ts | 14 +- .../internal/admin/handleListAdminUsers.ts | 3 +- .../cusProducts/CusProdReadService.ts | 47 ++-- .../rate-limits/rate-limit-overrides.test.ts | 126 ++++++++++ .../rate-limits/rate-limit-overrides.test.ts | 126 ++++++++++ .../productModels/entModels/entTable.ts | 6 + 11 files changed, 717 insertions(+), 22 deletions(-) create mode 100644 server/experiments/explainGetProductCounts.ts create mode 100644 server/experiments/explainGetRewards.ts create mode 100644 server/experiments/out/product-counts-explain-1.txt create mode 100644 server/experiments/out/product-counts-explain-2.txt create mode 100644 server/experiments/out/rewards-explain.txt create mode 100644 server/tests/integration/others/rate-limits/rate-limit-overrides.test.ts create mode 100644 server/tests/unit/rate-limits/rate-limit-overrides.test.ts diff --git a/server/experiments/explainGetProductCounts.ts b/server/experiments/explainGetProductCounts.ts new file mode 100644 index 000000000..bdf0b1701 --- /dev/null +++ b/server/experiments/explainGetProductCounts.ts @@ -0,0 +1,236 @@ +import { writeFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { AppEnv, schemas } from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import { drizzle } from "drizzle-orm/node-postgres"; +import pg from "pg"; +import { initDrizzle, prodTestOrgId } from "./experimentEnv"; + +const { CusProdReadService } = await import( + "../src/internal/customers/cusProducts/CusProdReadService" +); +const { ProductService } = await import( + "../src/internal/products/ProductService" +); + +// Run with: +// PRODUCT_COUNTS_ORG_ID=org_2x3YWDWucn3OSul12pIV6XrcXyo bun run experiments/explainGetProductCounts.ts + +const main = async () => { + const orgId = process.env.PRODUCT_COUNTS_ORG_ID || prodTestOrgId; + const env = AppEnv.Live; + + const { db } = initDrizzle(); + + console.log( + `--- Running GET /products/product_counts for org=${orgId} env=${env} ---\n`, + ); + + // Mirror the handler exactly + const startTotal = performance.now(); + + const startListFull = performance.now(); + const products = await ProductService.listFull({ db, orgId, env }); + const elapsedListFull = performance.now() - startListFull; + + console.log(`ProductService.listFull: ${elapsedListFull.toFixed(2)}ms (${products.length} products)\n`); + + const startCounts = performance.now(); + const perProductTimings: { id: string; ms: number }[] = []; + const counts = await Promise.all( + products.map(async (product) => { + const s = performance.now(); + const result = await CusProdReadService.getCountsForAllVersions({ + db, + productId: product.id, + orgId, + env, + }); + perProductTimings.push({ id: product.id, ms: performance.now() - s }); + return result; + }), + ); + const elapsedCounts = performance.now() - startCounts; + + const elapsedTotal = performance.now() - startTotal; + + console.log(`Promise.all over getCountsForAllVersions: ${elapsedCounts.toFixed(2)}ms`); + console.log(`Total handler logic: ${elapsedTotal.toFixed(2)}ms`); + console.log(`Counts returned: ${counts.length}\n`); + + console.log("--- Per-product timings (top 10 slowest) ---\n"); + perProductTimings.sort((a, b) => b.ms - a.ms); + for (const t of perProductTimings.slice(0, 10)) { + console.log(` ${t.id}: ${t.ms.toFixed(2)}ms`); + } + console.log(); + + const sumPerProduct = perProductTimings.reduce((acc, t) => acc + t.ms, 0); + const avgPerProduct = sumPerProduct / perProductTimings.length; + console.log(`Avg per-product: ${avgPerProduct.toFixed(2)}ms`); + console.log(`Sum (serial equivalent): ${sumPerProduct.toFixed(2)}ms\n`); + + // Capture Drizzle SQL for ONE per-product query and EXPLAIN it + const capturedQueries: { query: string; params: unknown[] }[] = []; + const pool = new pg.Pool({ + connectionString: process.env.DATABASE_URL, + max: 2, + }); + const loggedDb = drizzle(pool, { + schema: schemas, + logger: { + logQuery: (query, params) => capturedQueries.push({ query, params }), + }, + }); + + if (products[0]) { + console.log(`--- Capturing Drizzle SQL for product ${products[0].id} ---\n`); + await CusProdReadService.getCountsForAllVersions({ + db: loggedDb as never, + productId: products[0].id, + orgId, + env, + }); + + // Two queries are emitted: products lookup + the aggregate + for (const [i, q] of capturedQueries.entries()) { + const isAggregate = q.query.includes("count(distinct"); + console.log(`Query ${i + 1} ${isAggregate ? "(aggregate)" : "(product lookup)"}:`); + console.log(q.query); + console.log("Params:", q.params, "\n"); + + console.log(`--- EXPLAIN (ANALYZE, BUFFERS) ---\n`); + const plan = await pool.query( + `EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${q.query}`, + q.params as unknown[], + ); + const planLines: string[] = []; + for (const row of plan.rows) { + const line = (row as Record)["QUERY PLAN"]; + if (typeof line === "string") planLines.push(line); + } + console.log(planLines.join("\n")); + console.log(); + + const outDir = resolve(import.meta.dir, "out"); + try { + const { mkdirSync } = await import("node:fs"); + mkdirSync(outDir, { recursive: true }); + } catch { + // ignore + } + writeFileSync( + resolve(outDir, `product-counts-explain-${i + 1}.txt`), + planLines.join("\n"), + ); + } + } + + // Batched alternative — apples-to-apples with the new code shape: + // status filter pushed into the JOIN predicate so it can hit the partial + // covering index (idx_customer_products_active_counts). + console.log("--- Alternative: single batched query (all products, status filter pushed down) ---\n"); + const startBatched = performance.now(); + const batchedResult = await db.execute(sql` + SELECT + p.id AS product_id, + count(DISTINCT cp.internal_customer_id) AS active, + count(DISTINCT CASE WHEN cp.canceled_at IS NOT NULL THEN cp.internal_customer_id END) AS canceled, + count(DISTINCT CASE WHEN cp.is_custom = true THEN cp.internal_customer_id END) AS custom, + count(DISTINCT CASE WHEN cp.trial_ends_at IS NOT NULL AND cp.trial_ends_at > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint THEN cp.internal_customer_id END) AS trialing, + count(DISTINCT cp.internal_customer_id) AS all_count + FROM products p + LEFT JOIN customer_products cp + ON cp.internal_product_id = p.internal_id + AND cp.status IN ('active', 'past_due') + WHERE p.org_id = ${orgId} AND p.env = ${env} + GROUP BY p.id + `); + const elapsedBatched = performance.now() - startBatched; + console.log(`Batched query: ${elapsedBatched.toFixed(2)}ms (${batchedResult.length} rows)\n`); + + console.log(`--- EXPLAIN (ANALYZE, BUFFERS) on the batched query ---\n`); + const batchedPlan = await db.execute(sql` + EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) + SELECT + p.id AS product_id, + count(DISTINCT cp.internal_customer_id) AS active, + count(DISTINCT CASE WHEN cp.canceled_at IS NOT NULL THEN cp.internal_customer_id END) AS canceled, + count(DISTINCT CASE WHEN cp.is_custom = true THEN cp.internal_customer_id END) AS custom, + count(DISTINCT CASE WHEN cp.trial_ends_at IS NOT NULL AND cp.trial_ends_at > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint THEN cp.internal_customer_id END) AS trialing, + count(DISTINCT cp.internal_customer_id) AS all_count + FROM products p + LEFT JOIN customer_products cp + ON cp.internal_product_id = p.internal_id + AND cp.status IN ('active', 'past_due') + WHERE p.org_id = ${orgId} AND p.env = ${env} + GROUP BY p.id + `); + const batchedPlanLines: string[] = []; + for (const row of batchedPlan) { + const line = (row as Record)["QUERY PLAN"]; + if (typeof line === "string") batchedPlanLines.push(line); + } + console.log(batchedPlanLines.join("\n")); + console.log(); + + console.log("--- Table sizes ---\n"); + const sizes = await db.execute(sql` + SELECT + 'products' AS table_name, + (SELECT count(*) FROM products) AS total_rows, + (SELECT count(*) FROM products WHERE org_id = ${orgId} AND env = ${env}) AS org_rows, + pg_size_pretty(pg_relation_size('products')) AS table_size + UNION ALL + SELECT + 'customer_products', + (SELECT count(*) FROM customer_products), + (SELECT count(*) FROM customer_products cp + WHERE cp.internal_product_id IN + (SELECT internal_id FROM products WHERE org_id = ${orgId} AND env = ${env})), + pg_size_pretty(pg_relation_size('customer_products')) + `); + for (const row of sizes) { + console.log(row); + } + console.log(); + + console.log("--- Indexes on customer_products / products ---\n"); + const indexes = await db.execute(sql` + SELECT tablename, indexname, indexdef, + pg_size_pretty(pg_relation_size(quote_ident(indexname)::regclass)) AS size + FROM pg_indexes + WHERE tablename IN ('customer_products', 'products') + ORDER BY tablename, indexname + `); + for (const row of indexes) { + console.log(row); + } + + // Verdict — did the new partial covering index get picked up? + console.log("\n--- Verdict ---"); + const aggregateQuery = capturedQueries.find((q) => q.query.includes("count(distinct")); + if (aggregateQuery) { + const plan = await pool.query( + `EXPLAIN (FORMAT TEXT) ${aggregateQuery.query}`, + aggregateQuery.params as unknown[], + ); + const planText = plan.rows + .map((r) => (r as Record)["QUERY PLAN"]) + .join("\n"); + if (planText.includes("idx_customer_products_active_counts")) { + console.log("✅ Per-product query uses idx_customer_products_active_counts"); + } else if (planText.includes("Index Only Scan")) { + console.log("✅ Per-product query does an Index Only Scan (no heap)"); + } else if (planText.includes("Seq Scan on customer_products")) { + console.log("❌ Per-product query falls back to Seq Scan"); + } else { + console.log("⚠️ Per-product query plan unclear — inspect EXPLAIN above"); + } + } + + await pool.end(); + process.exit(0); +}; + +await main(); diff --git a/server/experiments/explainGetRewards.ts b/server/experiments/explainGetRewards.ts new file mode 100644 index 000000000..b15336e24 --- /dev/null +++ b/server/experiments/explainGetRewards.ts @@ -0,0 +1,151 @@ +import { writeFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { AppEnv, schemas } from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import { drizzle } from "drizzle-orm/node-postgres"; +import pg from "pg"; +import { initDrizzle, prodTestOrgId } from "./experimentEnv"; + +const { rewardRepo, rewardProgramRepo } = await import( + "../src/internal/rewards/repos/index" +); + +// Run with: +// REWARDS_ORG_ID=org_2x3YWDWucn3OSul12pIV6XrcXyo bun run experiments/explainGetRewards.ts + +const main = async () => { + const orgId = process.env.REWARDS_ORG_ID || prodTestOrgId; + const env = AppEnv.Live; + + const { db } = initDrizzle(); + + console.log(`--- Running GET /products/rewards for org=${orgId} env=${env} ---\n`); + + const startBoth = performance.now(); + const [rewards, rewardPrograms] = await Promise.all([ + rewardRepo.list({ db, orgId, env }), + rewardProgramRepo.list({ db, orgId, env }), + ]); + const elapsedBoth = performance.now() - startBoth; + + console.log(`Both queries (Promise.all): ${elapsedBoth.toFixed(2)}ms`); + console.log(` rewards: ${rewards.length} rows`); + console.log(` rewardPrograms: ${rewardPrograms.length} rows\n`); + + const startRewards = performance.now(); + const rewardsAlone = await rewardRepo.list({ db, orgId, env }); + const elapsedRewards = performance.now() - startRewards; + console.log(`rewardRepo.list alone: ${elapsedRewards.toFixed(2)}ms (${rewardsAlone.length} rows)`); + + const startPrograms = performance.now(); + const programsAlone = await rewardProgramRepo.list({ db, orgId, env }); + const elapsedPrograms = performance.now() - startPrograms; + console.log(`rewardProgramRepo.list alone: ${elapsedPrograms.toFixed(2)}ms (${programsAlone.length} rows)\n`); + + // Capture the actual SQL Drizzle generates for rewardRepo.list by wiring up + // a separate drizzle client with a logger. This lets us EXPLAIN ANALYZE the + // real query — not a hand-rolled approximation. + const capturedQueries: { query: string; params: unknown[] }[] = []; + const pool = new pg.Pool({ + connectionString: process.env.DATABASE_URL, + max: 2, + }); + const loggedDb = drizzle(pool, { + schema: schemas, + logger: { + logQuery: (query, params) => capturedQueries.push({ query, params }), + }, + }); + await rewardRepo.list({ db: loggedDb as never, orgId, env }); + const rewardSql = capturedQueries.find((q) => q.query.includes("rewards")); + if (!rewardSql) { + console.log("Could not capture Drizzle SQL — bailing on EXPLAIN.\n"); + } else { + console.log("--- Drizzle-generated SQL for rewardRepo.list ---\n"); + console.log(rewardSql.query); + console.log("\nParams:", rewardSql.params, "\n"); + + console.log("--- EXPLAIN (ANALYZE, BUFFERS) on the actual Drizzle query ---\n"); + const explainQuery = sql.raw( + `EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${rewardSql.query}`, + ); + // Bind the same params Drizzle bound. + const plan = await pool.query( + `EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${rewardSql.query}`, + rewardSql.params as unknown[], + ); + const planLines: string[] = []; + for (const row of plan.rows) { + const line = (row as Record)["QUERY PLAN"]; + if (typeof line === "string") planLines.push(line); + } + console.log(planLines.join("\n")); + + const outDir = resolve(import.meta.dir, "out"); + try { + const { mkdirSync } = await import("node:fs"); + mkdirSync(outDir, { recursive: true }); + } catch { + // ignore + } + const planPath = resolve(outDir, "rewards-explain.txt"); + writeFileSync(planPath, planLines.join("\n")); + console.log(`\nWritten to: ${planPath}\n`); + + const usedIndex = planLines.find((l) => /Index Scan|Index Only Scan/.test(l)); + const seqScan = planLines.find((l) => /Seq Scan on entitlements/.test(l)); + console.log("--- Verdict ---"); + if (seqScan) { + console.log("❌ Still seq-scanning entitlements. Index not picked up."); + console.log(` ${seqScan.trim()}`); + } else if (usedIndex) { + console.log("✅ Index scan in plan:"); + for (const l of planLines.filter((l) => /Index/.test(l))) { + console.log(` ${l.trim()}`); + } + } + console.log(); + + await pool.end(); + } + + console.log("--- Table sizes ---\n"); + const sizes = await db.execute(sql` + SELECT + 'rewards' AS table_name, + (SELECT count(*) FROM rewards) AS total_rows, + (SELECT count(*) FROM rewards WHERE org_id = ${orgId} AND env = ${env}) AS org_rows, + pg_size_pretty(pg_relation_size('rewards')) AS table_size + UNION ALL + SELECT + 'reward_programs', + (SELECT count(*) FROM reward_programs), + (SELECT count(*) FROM reward_programs WHERE org_id = ${orgId} AND env = ${env}), + pg_size_pretty(pg_relation_size('reward_programs')) + UNION ALL + SELECT + 'entitlements', + (SELECT count(*) FROM entitlements), + (SELECT count(*) FROM entitlements WHERE internal_reward_id IS NOT NULL), + pg_size_pretty(pg_relation_size('entitlements')) + `); + for (const row of sizes) { + console.log(row); + } + console.log(); + + console.log("--- Indexes on entitlements ---\n"); + const indexes = await db.execute(sql` + SELECT indexname, indexdef + FROM pg_indexes + WHERE tablename = 'entitlements' + ORDER BY indexname + `); + for (const row of indexes) { + console.log(row); + } + + process.exit(0); +}; + +await main(); diff --git a/server/experiments/out/product-counts-explain-1.txt b/server/experiments/out/product-counts-explain-1.txt new file mode 100644 index 000000000..0c78f0a3b --- /dev/null +++ b/server/experiments/out/product-counts-explain-1.txt @@ -0,0 +1,6 @@ +Index Scan using idx_products_org_env_id_version on products (cost=0.41..2.14 rows=1 width=32) (actual time=0.015..0.015 rows=1.00 loops=1) + Index Cond: ((org_id = 'biu9vSF7vghBLSKW1UTDwxHBAivjnPaK'::text) AND (env = 'live'::text) AND (id = 'legacy_starter'::text)) + Index Searches: 1 + Buffers: shared hit=4 +Planning Time: 0.044 ms +Execution Time: 0.034 ms \ No newline at end of file diff --git a/server/experiments/out/product-counts-explain-2.txt b/server/experiments/out/product-counts-explain-2.txt new file mode 100644 index 000000000..2fdd5a2ca --- /dev/null +++ b/server/experiments/out/product-counts-explain-2.txt @@ -0,0 +1,9 @@ +Aggregate (cost=24.05..24.06 rows=1 width=40) (actual time=0.055..0.055 rows=1.00 loops=1) + Buffers: shared hit=15 + -> Index Only Scan using idx_customer_products_active_counts on customer_products (cost=0.56..15.65 rows=336 width=53) (actual time=0.033..0.041 rows=20.00 loops=1) + Index Cond: (internal_product_id = 'prod_3B80yWLCKa9yCu6HcApN4mhFoh5'::text) + Heap Fetches: 0 + Index Searches: 1 + Buffers: shared hit=15 +Planning Time: 0.214 ms +Execution Time: 0.107 ms \ No newline at end of file diff --git a/server/experiments/out/rewards-explain.txt b/server/experiments/out/rewards-explain.txt new file mode 100644 index 000000000..5e2226e79 --- /dev/null +++ b/server/experiments/out/rewards-explain.txt @@ -0,0 +1,15 @@ +Nested Loop Left Join (cost=8463.99..169378.77 rows=20 width=474) (actual time=0.039..0.622 rows=27.00 loops=1) + Buffers: shared hit=926 + -> Index Scan Backward using coupons_pkey on rewards (cost=0.28..104.12 rows=20 width=442) (actual time=0.020..0.568 rows=27.00 loops=1) + Filter: ((org_id = 'org_2x3YWDWucn3OSul12pIV6XrcXyo'::text) AND (env = 'live'::text)) + Rows Removed by Filter: 1442 + Index Searches: 1 + Buffers: shared hit=898 + -> Aggregate (cost=8463.71..8463.72 rows=1 width=32) (actual time=0.001..0.001 rows=1.00 loops=27) + Buffers: shared hit=28 + -> Index Scan using idx_entitlements_internal_reward_id_c_partial on entitlements rewards_entitlements (cost=0.13..8409.56 rows=10830 width=359) (actual time=0.001..0.001 rows=0.04 loops=27) + Index Cond: (internal_reward_id = rewards.internal_id) + Index Searches: 27 + Buffers: shared hit=28 +Planning Time: 0.111 ms +Execution Time: 0.667 ms \ No newline at end of file diff --git a/server/src/internal/admin/handleListAdminOrgs.ts b/server/src/internal/admin/handleListAdminOrgs.ts index 3cf9f74c1..ed40e27a0 100644 --- a/server/src/internal/admin/handleListAdminOrgs.ts +++ b/server/src/internal/admin/handleListAdminOrgs.ts @@ -1,5 +1,16 @@ import { member, organizations, user, Scopes } from "@autumn/shared"; -import { and, desc, eq, gt, gte, ilike, inArray, lt, or } from "drizzle-orm"; +import { + and, + desc, + eq, + gt, + gte, + ilike, + inArray, + isNull, + lt, + or, +} from "drizzle-orm"; import { createRoute } from "../../honoMiddlewares/routeHandler"; import { getRequestBlockConfigFromSource } from "../misc/requestBlocks/requestBlockStore.js"; @@ -43,6 +54,7 @@ export const handleListAdminOrgs = createRoute({ .from(organizations) .where( and( + isNull(organizations.created_by), searchTerm ? or( eq(organizations.id, searchTerm), diff --git a/server/src/internal/admin/handleListAdminUsers.ts b/server/src/internal/admin/handleListAdminUsers.ts index 82865aab5..8cb2f97f6 100644 --- a/server/src/internal/admin/handleListAdminUsers.ts +++ b/server/src/internal/admin/handleListAdminUsers.ts @@ -1,5 +1,5 @@ import { user, Scopes } from "@autumn/shared"; -import { and, desc, eq, gt, gte, ilike, lt, or } from "drizzle-orm"; +import { and, desc, eq, gt, gte, ilike, isNull, lt, or } from "drizzle-orm"; import { createRoute } from "../../honoMiddlewares/routeHandler"; export const handleListAdminUsers = createRoute({ @@ -42,6 +42,7 @@ export const handleListAdminUsers = createRoute({ .from(user) .where( and( + isNull(user.createdBy), searchTerm ? or( eq(user.id, searchTerm), diff --git a/server/src/internal/customers/cusProducts/CusProdReadService.ts b/server/src/internal/customers/cusProducts/CusProdReadService.ts index ec127a7e3..564732675 100644 --- a/server/src/internal/customers/cusProducts/CusProdReadService.ts +++ b/server/src/internal/customers/cusProducts/CusProdReadService.ts @@ -45,24 +45,27 @@ export class CusProdReadService { }) => { const result = await db .select({ - active: countDistinct( - sql`CASE WHEN ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, - ).as("active"), + active: countDistinct(customerProducts.internal_customer_id).as( + "active", + ), canceled: countDistinct( - sql`CASE WHEN ${isNotNull(customerProducts.canceled_at)} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${isNotNull(customerProducts.canceled_at)} THEN ${customerProducts.internal_customer_id} END`, ).as("canceled"), custom: countDistinct( - sql`CASE WHEN ${eq(customerProducts.is_custom, true)} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${eq(customerProducts.is_custom, true)} THEN ${customerProducts.internal_customer_id} END`, ).as("custom"), trialing: countDistinct( - sql`CASE WHEN ${isNotNull(customerProducts.trial_ends_at)} AND ${sql`${customerProducts.trial_ends_at} > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint`} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${isNotNull(customerProducts.trial_ends_at)} AND ${sql`${customerProducts.trial_ends_at} > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint`} THEN ${customerProducts.internal_customer_id} END`, ).as("trialing"), - all: countDistinct( - sql`CASE WHEN ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, - ).as("all"), + all: countDistinct(customerProducts.internal_customer_id).as("all"), }) .from(customerProducts) - .where(eq(customerProducts.internal_product_id, internalProductId)); + .where( + and( + eq(customerProducts.internal_product_id, internalProductId), + inArray(customerProducts.status, activeStatuses), + ), + ); return result[0]; }; @@ -97,25 +100,29 @@ export class CusProdReadService { const result = await db .select({ - active: countDistinct( - sql`CASE WHEN ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, - ).as("active"), + active: countDistinct(customerProducts.internal_customer_id).as( + "active", + ), canceled: countDistinct( - sql`CASE WHEN ${isNotNull(customerProducts.canceled_at)} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${isNotNull(customerProducts.canceled_at)} THEN ${customerProducts.internal_customer_id} END`, ).as("canceled"), custom: countDistinct( - sql`CASE WHEN ${eq(customerProducts.is_custom, true)} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${eq(customerProducts.is_custom, true)} THEN ${customerProducts.internal_customer_id} END`, ).as("custom"), trialing: countDistinct( - sql`CASE WHEN ${isNotNull(customerProducts.trial_ends_at)} AND ${sql`${customerProducts.trial_ends_at} > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint`} AND ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, + sql`CASE WHEN ${isNotNull(customerProducts.trial_ends_at)} AND ${sql`${customerProducts.trial_ends_at} > (EXTRACT(EPOCH FROM NOW()) * 1000)::bigint`} THEN ${customerProducts.internal_customer_id} END`, ).as("trialing"), - all: countDistinct( - sql`CASE WHEN ${inArray(customerProducts.status, activeStatuses)} THEN ${customerProducts.internal_customer_id} END`, - ).as("all"), + all: countDistinct(customerProducts.internal_customer_id).as("all"), }) .from(customerProducts) .where( - inArray(customerProducts.internal_product_id, internalProductIdsArray), + and( + inArray( + customerProducts.internal_product_id, + internalProductIdsArray, + ), + inArray(customerProducts.status, activeStatuses), + ), ); return result[0]; diff --git a/server/tests/integration/others/rate-limits/rate-limit-overrides.test.ts b/server/tests/integration/others/rate-limits/rate-limit-overrides.test.ts new file mode 100644 index 000000000..c5bd38c5b --- /dev/null +++ b/server/tests/integration/others/rate-limits/rate-limit-overrides.test.ts @@ -0,0 +1,126 @@ +/** + * Asserts the edge-config rate-limit override is actually honored by the + * factory at request time. We point the override at /v1/entities.get for + * the test org (which has a default limit of 100/s) and reduce it to 5/s. + * 6+ parallel requests should then trip the limit. + */ + +import { afterAll, beforeAll, expect, test } from "bun:test"; +import { ApiVersion } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import AutumnError, { AutumnInt } from "@/external/autumn/autumnCli.js"; +import { RateLimitType } from "@/internal/misc/rateLimiter/rateLimitConfigs.js"; +import { _setRateLimitOverridesConfigForTesting } from "@/internal/misc/rateLimiter/rateLimitOverridesStore.js"; + +const testCase = "rate-limit-overrides"; + +const countRateLimited = (results: PromiseSettledResult[]): number => + results.filter( + (result) => + result.status === "rejected" && + result.reason instanceof AutumnError && + result.reason.code === "rate_limit_exceeded", + ).length; + +let orgIdForRevert: string | undefined; + +beforeAll(() => { + _setRateLimitOverridesConfigForTesting({ config: { orgs: {} } }); +}); + +afterAll(() => { + _setRateLimitOverridesConfigForTesting({ config: { orgs: {} } }); +}); + +test(`${chalk.yellowBright(`${testCase}: override lowers effective limit for the org`)}`, async () => { + const customerId = "rate-limit-override-lower"; + const entityId = "entity-rate-limit-override"; + + const { autumnV1, ctx } = await initScenario({ + customerId, + setup: [s.customer({ testClock: false })], + actions: [], + }); + + await autumnV1.entities.create(customerId, { + id: entityId, + name: "Rate Limit Override Entity", + feature_id: TestFeature.Users, + }); + + orgIdForRevert = ctx.org.id; + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + [ctx.org.id]: { + limits: { [RateLimitType.CustomerEntitiesGet]: 5 }, + }, + }, + }, + }); + + const client = new AutumnInt({ + version: ApiVersion.V1_2, + secretKey: ctx.orgSecretKey, + }); + + // 12 requests, override is 5/s — well above the override, well below the + // default of 100. If the override is being honored, several get 429'd. + const results = await Promise.allSettled( + Array.from({ length: 12 }, () => + client.post("/entities.get", { + customer_id: customerId, + entity_id: entityId, + }), + ), + ); + + expect(countRateLimited(results)).toBeGreaterThan(0); +}); + +test(`${chalk.yellowBright(`${testCase}: orgSlug fallback resolves the override`)}`, async () => { + const customerId = "rate-limit-override-slug"; + const entityId = "entity-rate-limit-slug"; + + const { autumnV1, ctx } = await initScenario({ + customerId, + setup: [s.customer({ testClock: false })], + actions: [], + }); + + await autumnV1.entities.create(customerId, { + id: entityId, + name: "Rate Limit Override Slug Entity", + feature_id: TestFeature.Users, + }); + + const slug = ctx.org.slug; + if (!slug) throw new Error("test org has no slug"); + + orgIdForRevert = ctx.org.id; + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + [slug]: { limits: { [RateLimitType.CustomerEntitiesGet]: 5 } }, + }, + }, + }); + + const client = new AutumnInt({ + version: ApiVersion.V1_2, + secretKey: ctx.orgSecretKey, + }); + + const results = await Promise.allSettled( + Array.from({ length: 12 }, () => + client.post("/entities.get", { + customer_id: customerId, + entity_id: entityId, + }), + ), + ); + + expect(countRateLimited(results)).toBeGreaterThan(0); +}); diff --git a/server/tests/unit/rate-limits/rate-limit-overrides.test.ts b/server/tests/unit/rate-limits/rate-limit-overrides.test.ts new file mode 100644 index 000000000..c16414800 --- /dev/null +++ b/server/tests/unit/rate-limits/rate-limit-overrides.test.ts @@ -0,0 +1,126 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { RateLimitType } from "@/internal/misc/rateLimiter/rateLimitConfigs.js"; +import { + _setRateLimitOverridesConfigForTesting, + getOrgRateLimitOverride, +} from "@/internal/misc/rateLimiter/rateLimitOverridesStore.js"; + +const reset = () => { + _setRateLimitOverridesConfigForTesting({ config: { orgs: {} } }); +}; + +describe("getOrgRateLimitOverride", () => { + afterEach(reset); + + test("returns undefined when no override is configured", () => { + reset(); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + type: RateLimitType.CustomerEntitiesGet, + }), + ).toBeUndefined(); + }); + + test("returns the override when matched by orgId", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + org_a: { limits: { [RateLimitType.CustomerEntitiesGet]: 500 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + type: RateLimitType.CustomerEntitiesGet, + }), + ).toBe(500); + }); + + test("falls back to orgSlug when orgId has no entry", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + mintlify: { limits: { [RateLimitType.CustomerEntitiesGet]: 200 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ + orgId: "org_unknown", + orgSlug: "mintlify", + type: RateLimitType.CustomerEntitiesGet, + }), + ).toBe(200); + }); + + test("orgId match wins over orgSlug match", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + org_a: { limits: { [RateLimitType.CustomerEntitiesGet]: 999 } }, + mintlify: { limits: { [RateLimitType.CustomerEntitiesGet]: 1 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + orgSlug: "mintlify", + type: RateLimitType.CustomerEntitiesGet, + }), + ).toBe(999); + }); + + test("overrides are scoped per RateLimitType — other types fall through", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + org_a: { limits: { [RateLimitType.CustomerEntitiesGet]: 500 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + type: RateLimitType.Check, + }), + ).toBeUndefined(); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + type: RateLimitType.Track, + }), + ).toBeUndefined(); + }); + + test("returns undefined when neither orgId nor orgSlug is supplied", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + org_a: { limits: { [RateLimitType.CustomerEntitiesGet]: 500 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ type: RateLimitType.CustomerEntitiesGet }), + ).toBeUndefined(); + }); + + test("supports overriding a value to 0 (effectively block)", () => { + _setRateLimitOverridesConfigForTesting({ + config: { + orgs: { + org_a: { limits: { [RateLimitType.CustomerEntitiesGet]: 0 } }, + }, + }, + }); + expect( + getOrgRateLimitOverride({ + orgId: "org_a", + type: RateLimitType.CustomerEntitiesGet, + }), + ).toBe(0); + }); +}); diff --git a/shared/models/productModels/entModels/entTable.ts b/shared/models/productModels/entModels/entTable.ts index 801109356..f8837b07c 100644 --- a/shared/models/productModels/entModels/entTable.ts +++ b/shared/models/productModels/entModels/entTable.ts @@ -68,6 +68,12 @@ export const entitlements = pgTable( table.internal_reward_id, table.internal_feature_id, ), + // Serves joins on rewards.internal_id (collation C). The plain + // internal_reward_id index above is default-collation and can't be used + // when the join collation is C. + index("idx_entitlements_internal_reward_id_c_partial") + .on(sql`${table.internal_reward_id} COLLATE "C"`) + .where(sql`${table.internal_reward_id} IS NOT NULL`), ], );