feat(analytics): query timeout limits based on plan

We update the schema to enforce  a max of the growth plan limit.

Author: nams1570

apps/backend/src/app/api/latest/internal/analytics/query/route.ts

@@ -1,13 +1,16 @@
 import { getClickhouseExternalClient } from "@/lib/clickhouse";
+import { getBillingTeamId } from "@/lib/plan-entitlements";
 import { createSmartRouteHandler } from "@/route-handlers/smart-route-handler";
+import { getStackServerApp } from "@/stack";
 import { KnownErrors } from "@stackframe/stack-shared";
+import { ITEM_IDS, PLAN_LIMITS } from "@stackframe/stack-shared/dist/plans";
 import { adaptSchema, adminAuthTypeSchema, jsonSchema, yupBoolean, yupMixed, yupNumber, yupObject, yupRecord, yupString } from "@stackframe/stack-shared/dist/schema-fields";
 import { getNodeEnvironment } from "@stackframe/stack-shared/dist/utils/env";
 import { captureError, StackAssertionError } from "@stackframe/stack-shared/dist/utils/errors";
 import { Result } from "@stackframe/stack-shared/dist/utils/results";
 import { randomUUID } from "crypto";
 
-const MAX_QUERY_TIMEOUT_MS = 120_000;
+const MAX_QUERY_TIMEOUT_MS = Math.max(...Object.values(PLAN_LIMITS).map(p => p.analyticsTimeoutSeconds)) * 1000;
 const DEFAULT_QUERY_TIMEOUT_MS = 10_000;
 
 export const POST = createSmartRouteHandler({
@@ -36,6 +39,16 @@ export const POST = createSmartRouteHandler({
     if (body.include_all_branches) {
       throw new StackAssertionError("include_all_branches is not supported yet");
     }
+
+    let effectiveTimeoutMs = body.timeout_ms;
+    const billingTeamId = getBillingTeamId(auth.tenancy.project);
+    if (billingTeamId != null) {
+      const app = getStackServerApp();
+      const timeoutItem = await app.getItem({ itemId: ITEM_IDS.analyticsTimeoutSeconds, teamId: billingTeamId });
+      const maxAllowedMs = timeoutItem.quantity * 1000;
+      effectiveTimeoutMs = Math.min(body.timeout_ms, maxAllowedMs);
+    }
+
     const client = getClickhouseExternalClient();
     const queryId = `${auth.tenancy.project.id}:${auth.tenancy.branchId}:${randomUUID()}`;
     const resultSet = await Result.fromPromise(client.query({
@@ -45,7 +58,7 @@ export const POST = createSmartRouteHandler({
       clickhouse_settings: {
         SQL_project_id: auth.tenancy.project.id,
         SQL_branch_id: auth.tenancy.branchId,
-        max_execution_time: body.timeout_ms / 1000,
+        max_execution_time: effectiveTimeoutMs / 1000,
         readonly: "1",
         allow_ddl: 0,
         max_result_rows: MAX_RESULT_ROWS.toString(),

apps/e2e/tests/backend/endpoints/api/v1/analytics-query.test.ts

@@ -1,7 +1,8 @@
+import { PLAN_LIMITS, PlanId } from "@stackframe/stack-shared/dist/plans";
 import { wait } from "@stackframe/stack-shared/dist/utils/promises";
 import { deindent } from "@stackframe/stack-shared/dist/utils/strings";
 import { it } from "../../../../helpers";
-import { Project, User, niceBackendFetch } from "../../../backend-helpers";
+import { backendContext, InternalProjectKeys, Project, User, niceBackendFetch } from "../../../backend-helpers";
 
 async function runQuery(body: { query: string, params?: Record<string, string>, timeout_ms?: number }) {
   await Project.createAndSwitch({ config: { magic_link_enabled: true } });
@@ -15,6 +16,33 @@ async function runQuery(body: { query: string, params?: Record<string, string>,
   return response;
 }
 
+async function runQueryWithPlan(planId: PlanId, body: { query: string, params?: Record<string, string>, timeout_ms?: number }) {
+  const { createProjectResponse, adminAccessToken } = await Project.createAndSwitch({ config: { magic_link_enabled: true } });
+  const ownerTeamId = createProjectResponse.body.owner_team_id;
+
+  if (planId !== "free") {
+    const savedKeys = backendContext.value.projectKeys;
+    backendContext.set({ projectKeys: InternalProjectKeys });
+    const grantResponse = await niceBackendFetch(`/api/v1/payments/products/team/${ownerTeamId}`, {
+      method: "POST",
+      accessType: "server",
+      body: { product_id: planId },
+    });
+    if (grantResponse.status !== 200) {
+      throw new Error(`Failed to grant plan '${planId}' to team '${ownerTeamId}': ${JSON.stringify(grantResponse.body)}`);
+    }
+    backendContext.set({ projectKeys: savedKeys });
+  }
+
+  const response = await niceBackendFetch("/api/v1/internal/analytics/query", {
+    method: "POST",
+    accessType: "admin",
+    body,
+  });
+
+  return response;
+}
+
 type ExpectLike = ((value: unknown) => { toEqual: (value: unknown) => void }) & {
   any: (constructor: unknown) => unknown,
 };
@@ -154,10 +182,11 @@ it("can execute a query with custom timeout", async ({ expect }) => {
   `);
 });
 
-it("rejects timeouts longer than 2 minutes", async ({ expect }) => {
+it("rejects timeouts longer than max plan limit", async ({ expect }) => {
+  const maxSchemaMs = Math.max(...Object.values(PLAN_LIMITS).map(p => p.analyticsTimeoutSeconds)) * 1000;
   const response = await runQuery({
     query: "SELECT 1 as value",
-    timeout_ms: 120_001,
+    timeout_ms: maxSchemaMs + 1,
   });
 
   expect(stripQueryId(response, expect)).toMatchInlineSnapshot(`
@@ -168,12 +197,12 @@ it("rejects timeouts longer than 2 minutes", async ({ expect }) => {
         "details": {
           "message": deindent\`
             Request validation failed on POST /api/v1/internal/analytics/query:
-              - body.timeout_ms must be less than or equal to 120000
+              - body.timeout_ms must be less than or equal to ${maxSchemaMs}
           \`,
         },
         "error": deindent\`
           Request validation failed on POST /api/v1/internal/analytics/query:
-            - body.timeout_ms must be less than or equal to 120000
+            - body.timeout_ms must be less than or equal to ${maxSchemaMs}
         \`,
       },
       "headers": Headers {
@@ -1524,6 +1553,51 @@ it("does not allow input() function", async ({ expect }) => {
   `);
 });
 
+it("clamps timeout to free plan limit", async ({ expect }) => {
+  const response = await runQueryWithPlan("free", {
+    query: "SELECT getSetting('max_execution_time') as max_execution_time",
+    timeout_ms: 120000,
+  });
+
+  expect(response.status).toBe(200);
+  const maxExecutionTime = Number((response.body?.result as any)?.[0]?.max_execution_time);
+  expect(maxExecutionTime).toBe(PLAN_LIMITS.free.analyticsTimeoutSeconds);
+});
+
+it("clamps timeout to team plan limit", async ({ expect }) => {
+  const response = await runQueryWithPlan("team", {
+    query: "SELECT getSetting('max_execution_time') as max_execution_time",
+    timeout_ms: 120000,
+  });
+
+  expect(response.status).toBe(200);
+  const maxExecutionTime = Number((response.body?.result as any)?.[0]?.max_execution_time);
+  expect(maxExecutionTime).toBe(PLAN_LIMITS.team.analyticsTimeoutSeconds);
+});
+
+it("clamps timeout to growth plan limit", async ({ expect }) => {
+  const maxSchemaMs = Math.max(...Object.values(PLAN_LIMITS).map(p => p.analyticsTimeoutSeconds)) * 1000;
+  const response = await runQueryWithPlan("growth", {
+    query: "SELECT getSetting('max_execution_time') as max_execution_time",
+    timeout_ms: maxSchemaMs,
+  });
+
+  expect(response.status).toBe(200);
+  const maxExecutionTime = Number((response.body?.result as any)?.[0]?.max_execution_time);
+  expect(maxExecutionTime).toBe(PLAN_LIMITS.growth.analyticsTimeoutSeconds);
+});
+
+it("does not clamp timeout below the plan limit", async ({ expect }) => {
+  const response = await runQueryWithPlan("team", {
+    query: "SELECT getSetting('max_execution_time') as max_execution_time",
+    timeout_ms: 5000,
+  });
+
+  expect(response.status).toBe(200);
+  const maxExecutionTime = Number((response.body?.result as any)?.[0]?.max_execution_time);
+  expect(maxExecutionTime).toBe(5);
+});
+
 it("does not allow numbers table function with large values", async ({ expect }) => {
   const response = await runQuery({
     query: "SELECT * FROM numbers(1000000000)",