latest

2026-05-28 21:06:23 +01:00
parent 44950c1c09
commit aa8f887e11
3 changed files with 466 additions and 32 deletions
--- a/server/experiments/explainIncludeProcessedFilter.ts
+++ b/server/experiments/explainIncludeProcessedFilter.ts
@@ -0,0 +1,349 @@
+import { AppEnv } from "@autumn/shared";
+import { compileFilter } from "@autumn/shared/api/migrations/compiler/compileFilter.js";
+import type { CustomerFilter } from "@autumn/shared/api/migrations/filters/customerFilter.js";
+import { type SQL, sql } from "drizzle-orm";
+import { PgDialect } from "drizzle-orm/pg-core";
+import {
+	buildCustomerCount,
+	buildCustomerSelect,
+	buildProcessedPreviewCount,
+	buildProcessedPreviewSelect,
+} from "@/internal/migrations/v2/filters/customers/buildCustomerSelect.js";
+import { rawWithParamsToDrizzle } from "@/internal/migrations/v2/filters/rawWithParamsToDrizzle.js";
+// Import initDrizzle directly — avoid `experimentEnv` because its
+// `loadLocalEnv()` reads `server/.env` and clobbers env vars injected by
+// `infisical run --env=staging` (e.g. DATABASE_URL).
+import { initDrizzle } from "../src/db/initDrizzle";
+import { FeatureService } from "../src/internal/features/FeatureService.js";
+
+// Why this experiment exists: the "include processed customers" preview
+// (handlePreviewMigrationFilter + buildCustomerSelect) ORs the org/env-scoped
+// compiled filter with `c.internal_id IN (<processed subquery>)`. The OR strips
+// org/env scoping from the second branch, so the planner can't use
+// idx_customers_org_env_internal_id and may seq-scan ALL customers. This script
+// EXPLAINs the current OR query against an equivalent UNION rewrite to confirm
+// the bottleneck and decide whether a new index is needed.
+//
+// Run against a remote env (e.g. staging) via infisical:
+//   infisical run --env=staging --recursive -- \
+//     bun run server/experiments/explainIncludeProcessedFilter.ts
+
+const prodTestOrgId = (() => {
+	const v = process.env.PROD_TEST_ORG_ID;
+	if (!v) throw new Error("PROD_TEST_ORG_ID env var is required");
+	return v;
+})();
+
+const dbUrl = process.env.DATABASE_URL ?? "";
+console.log(
+	"DATABASE URL host:",
+	dbUrl.replace(/:\/\/[^@]+@/, "://***:***@") || "(empty)",
+);
+
+// ─── Configuration ──────────────────────────────────────────────────
+const ORG_ID = prodTestOrgId;
+const ENV = AppEnv.Live;
+const SAMPLE_LIMIT = 10; // matches the default preview page size
+const TRUNCATE_EXPLAIN = true;
+const EXPLAIN_MAX_LINES = 40;
+
+// Optional override. When unset, the script auto-discovers the migration in
+// this org/env with the most live (dry_run = false) customer item runs.
+const MIGRATION_INTERNAL_ID = process.env.MIGRATION_INTERNAL_ID || undefined;
+
+// User-facing migration id (the `id` column, resolved to internal_id like the
+// production handler does). Takes precedence over auto-discovery.
+const MIGRATION_ID = process.env.MIGRATION_ID || "plan_pro-update";
+
+// Filter the live preview applies. Keep it representative of a real migration
+// selection. An empty `{}` matches all customers in the org/env.
+const FILTER: CustomerFilter = {
+	plan: { plan_id: "free" },
+};
+
+// ═════════════════════════════════════════════════════════════════════
+
+const truncateExplainText = (text: string, maxLines: number): string => {
+	const lines = text.split("\n");
+	if (lines.length <= maxLines) return text;
+	const omitted = lines.length - maxLines;
+	return [...lines.slice(0, maxLines), `... (${omitted} more lines truncated)`].join(
+		"\n",
+	);
+};
+
+const printExplainPlan = async ({
+	db,
+	query,
+	label,
+}: {
+	db: ReturnType<typeof initDrizzle>["db"];
+	query: SQL;
+	label: string;
+}) => {
+	console.log(`\n--- EXPLAIN ANALYZE: ${label} ---`);
+	const explainResult = await db.execute(
+		sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`,
+	);
+	const lines: string[] = [];
+	for (const row of explainResult)
+		lines.push(String((row as Record<string, unknown>)["QUERY PLAN"]));
+	const joined = lines.join("\n");
+	console.log(
+		TRUNCATE_EXPLAIN ? truncateExplainText(joined, EXPLAIN_MAX_LINES) : joined,
+	);
+};
+
+const dialect = new PgDialect();
+
+const inlineParams = (text: string, params: readonly unknown[]): string =>
+	text.replace(/\$(\d+)/g, (_, n) => {
+		const v = params[Number(n) - 1];
+		if (v === null || v === undefined) return "NULL";
+		if (typeof v === "number" || typeof v === "boolean") return String(v);
+		return `'${String(v).replace(/'/g, "''")}'`;
+	});
+
+const printSqlQuery = ({ query, label }: { query: SQL; label: string }) => {
+	const { sql: text, params } = dialect.sqlToQuery(query);
+	console.log(`\n--- SQL: ${label} ---`);
+	console.log(inlineParams(text, params));
+};
+
+const runMeasured = async ({
+	db,
+	query,
+	label,
+}: {
+	db: ReturnType<typeof initDrizzle>["db"];
+	query: SQL;
+	label: string;
+}) => {
+	console.log(`\n=== ${label} ===`);
+	printSqlQuery({ query, label });
+	const startedAt = performance.now();
+	const result = await db.execute(query);
+	const elapsedMs = performance.now() - startedAt;
+	console.log(`Rows returned: ${result.length}`);
+	console.log(`Wall-clock: ${elapsedMs.toFixed(2)}ms`);
+	if (label.startsWith("COUNT") && result.length > 0)
+		console.log(`Count: ${(result[0] as Record<string, unknown>).count}`);
+	await printExplainPlan({ db, query, label });
+};
+
+const compiledWhere = ({
+	filter,
+	features,
+}: {
+	filter: CustomerFilter;
+	features: Awaited<ReturnType<typeof FeatureService.list>>;
+}): SQL =>
+	rawWithParamsToDrizzle(
+		compileFilter({
+			filter,
+			ctx: { features },
+			ambient: { orgId: ORG_ID, env: ENV },
+		}),
+	);
+
+// The processed-customers subquery, identical to buildIncludeProcessedOr.
+const processedSubquery = (migrationInternalId: string): SQL => sql`
+	SELECT mir.item_id FROM migration_item_runs mir
+	WHERE mir.migration_internal_id = ${migrationInternalId}
+		AND mir.item_kind = 'customer'
+		AND mir.dry_run = false
+`;
+
+// Proposed UNION rewrite: each branch keeps its own scoping so the planner can
+// use an index per branch instead of seq-scanning all customers.
+const buildUnionSelect = ({
+	where,
+	migrationInternalId,
+	limit,
+}: {
+	where: SQL;
+	migrationInternalId: string;
+	limit: number;
+}): SQL => sql`
+	SELECT u.internal_id, u.id, u.name, u.email
+	FROM (
+		SELECT c.internal_id, c.id, c.name, c.email
+		FROM customers c
+		WHERE (${where})
+		UNION
+		SELECT c.internal_id, c.id, c.name, c.email
+		FROM customers c
+		WHERE c.internal_id IN (${processedSubquery(migrationInternalId)})
+	) u
+	ORDER BY u.internal_id DESC
+	LIMIT ${limit}
+`;
+
+const buildUnionCount = ({
+	where,
+	migrationInternalId,
+}: {
+	where: SQL;
+	migrationInternalId: string;
+}): SQL => sql`
+	SELECT COUNT(*)::bigint AS count
+	FROM (
+		SELECT c.internal_id
+		FROM customers c
+		WHERE (${where})
+		UNION
+		SELECT c.internal_id
+		FROM customers c
+		WHERE c.internal_id IN (${processedSubquery(migrationInternalId)})
+	) u
+`;
+
+// Resolve a user-facing migration `id` to its `internal_id`, scoped to org/env
+// — mirrors migrationRepo.find used by handlePreviewMigrationFilter.
+const resolveMigrationInternalId = async (
+	db: ReturnType<typeof initDrizzle>["db"],
+	id: string,
+): Promise<string | undefined> => {
+	const rows = (await db.execute(sql`
+		SELECT internal_id FROM migrations
+		WHERE org_id = ${ORG_ID} AND env = ${ENV} AND id = ${id}
+		LIMIT 1
+	`)) as Array<{ internal_id: string }>;
+	return rows[0]?.internal_id;
+};
+
+const discoverMigrationInternalId = async (
+	db: ReturnType<typeof initDrizzle>["db"],
+): Promise<string | undefined> => {
+	const rows = (await db.execute(sql`
+		SELECT mir.migration_internal_id AS migration_internal_id, COUNT(*) AS n
+		FROM migration_item_runs mir
+		JOIN migration_runs mr ON mr.migration_internal_id = mir.migration_internal_id
+		WHERE mr.org_id = ${ORG_ID}
+			AND mr.env = ${ENV}
+			AND mir.item_kind = 'customer'
+			AND mir.dry_run = false
+		GROUP BY mir.migration_internal_id
+		ORDER BY n DESC
+		LIMIT 5
+	`)) as Array<{ migration_internal_id: string; n: bigint | number }>;
+
+	if (rows.length === 0) return undefined;
+	console.log("\nMigrations with live customer item runs (top 5):");
+	for (const r of rows)
+		console.log(`  ${r.migration_internal_id} → ${Number(r.n)} processed`);
+	return rows[0].migration_internal_id;
+};
+
+const main = async () => {
+	const replicaUrl = process.env.DATABASE_REPLICA_URL;
+	const usingReplica = Boolean(replicaUrl);
+	if (!usingReplica)
+		console.warn(
+			"DATABASE_REPLICA_URL not set — falling back to DATABASE_URL (primary). Set the replica URL to test against the read replica.",
+		);
+	const { db } = initDrizzle({ replica: usingReplica });
+
+	console.log(
+		`=== INCLUDE-PROCESSED FILTER EXPERIMENT (${usingReplica ? "REPLICA" : "PRIMARY"}) ===`,
+	);
+	console.log(JSON.stringify({ ORG_ID, ENV, FILTER }, null, 2));
+
+	let migrationInternalId = MIGRATION_INTERNAL_ID;
+	if (!migrationInternalId && MIGRATION_ID) {
+		migrationInternalId = await resolveMigrationInternalId(db, MIGRATION_ID);
+		if (migrationInternalId)
+			console.log(`\nResolved MIGRATION_ID '${MIGRATION_ID}' → ${migrationInternalId}`);
+		else
+			console.warn(
+				`\nMIGRATION_ID '${MIGRATION_ID}' not found for this org/env — falling back to auto-discovery.`,
+			);
+	}
+	migrationInternalId ??= await discoverMigrationInternalId(db);
+	if (!migrationInternalId) {
+		console.error(
+			"\nNo migration with live customer item runs found for this org/env. " +
+				"Set MIGRATION_INTERNAL_ID or MIGRATION_ID explicitly to test a specific migration.",
+		);
+		process.exit(1);
+	}
+	console.log(`\nUsing migration_internal_id: ${migrationInternalId}`);
+
+	const orgFeatures = await FeatureService.list({ db, orgId: ORG_ID, env: ENV });
+	console.log(`\nLoaded ${orgFeatures.length} features for resolution context.`);
+	const ctx = { features: orgFeatures };
+	const where = compiledWhere({ filter: FILTER, features: orgFeatures });
+
+	const includeProcessed = { migrationInternalId };
+
+	// 1. Isolated processed subquery — confirms migration_item_runs index coverage.
+	await runMeasured({
+		db,
+		query: sql`SELECT mir.item_id FROM migration_item_runs mir
+			WHERE mir.migration_internal_id = ${migrationInternalId}
+				AND mir.item_kind = 'customer'
+				AND mir.dry_run = false`,
+		label: "SUBQUERY (processed item_ids only)",
+	});
+
+	// 2. Pure filter — exactly what the FILTER STEP (no migrationId) runs.
+	//    Baseline to prove the customer filter alone is fast; only the live
+	//    view's includeProcessed OR is slow.
+	await runMeasured({
+		db,
+		query: buildCustomerCount({ orgId: ORG_ID, env: ENV, filter: FILTER, ctx }),
+		label: "COUNT [filter only — filter step]",
+	});
+	await runMeasured({
+		db,
+		query: buildCustomerSelect({
+			orgId: ORG_ID,
+			env: ENV,
+			filter: FILTER,
+			ctx,
+			limit: SAMPLE_LIMIT,
+		}),
+		label: `SELECT [filter only — filter step] (limit ${SAMPLE_LIMIT})`,
+	});
+
+	// 3. Live-view path: the dedicated preview builders (filter ∪ processed).
+	await runMeasured({
+		db,
+		query: buildProcessedPreviewCount({
+			orgId: ORG_ID,
+			env: ENV,
+			filter: FILTER,
+			ctx,
+			includeProcessed,
+		}),
+		label: "COUNT [preview builder]",
+	});
+	await runMeasured({
+		db,
+		query: buildProcessedPreviewSelect({
+			orgId: ORG_ID,
+			env: ENV,
+			filter: FILTER,
+			ctx,
+			includeProcessed,
+			limit: SAMPLE_LIMIT,
+		}),
+		label: `SELECT [preview builder] (limit ${SAMPLE_LIMIT})`,
+	});
+
+	// 4. Hand-written UNION reference (sanity check the builder matches this).
+	await runMeasured({
+		db,
+		query: buildUnionCount({ where, migrationInternalId }),
+		label: "COUNT [UNION — proposed]",
+	});
+	await runMeasured({
+		db,
+		query: buildUnionSelect({ where, migrationInternalId, limit: SAMPLE_LIMIT }),
+		label: `SELECT [UNION — proposed] (limit ${SAMPLE_LIMIT})`,
+	});
+
+	process.exit(0);
+};
+
+await main();
--- a/server/src/internal/migrations/v2/filters/customers/buildCustomerSelect.ts
+++ b/server/src/internal/migrations/v2/filters/customers/buildCustomerSelect.ts
@@ -15,7 +15,6 @@ export type CustomerQueryArgs = {
 	ctx: ResolutionContext;
 	checkpoint?: CustomerCheckpointExclusion;
 	search?: string;
-	includeProcessed?: IncludeProcessed;
 };

 const compileWhere = ({ orgId, env, filter, ctx }: CustomerQueryArgs): SQL =>
@@ -68,22 +67,38 @@ const buildSearchWhere = (search: string | undefined): SQL => {
 	return sql`AND (c.name ILIKE ${pattern} OR c.email ILIKE ${pattern} OR c.id ILIKE ${pattern})`;
 };

-const buildIncludeProcessedOr = (
-	includeProcessed: IncludeProcessed | undefined,
-): SQL => {
-	if (!includeProcessed) return sql``;
-	return sql`OR c.internal_id IN (
+const buildProcessedIn = (includeProcessed: IncludeProcessed): SQL => sql`
+	c.internal_id IN (
 		SELECT mir.item_id FROM migration_item_runs mir
 		WHERE mir.migration_internal_id = ${includeProcessed.migrationInternalId}
 			AND mir.item_kind = 'customer'
 			AND mir.dry_run = false
 	)`;
+
+// Predicates shared by both UNION branches (and the single-branch query).
+// Rebuilt per call so a branch never reuses another's SQL chunk instance.
+const buildCommonWhere = ({
+	checkpoint,
+	search,
+	afterInternalId,
+}: {
+	checkpoint?: CustomerCheckpointExclusion;
+	search?: string;
+	afterInternalId?: string;
+}): SQL => {
+	const cursor = afterInternalId
+		? sql`AND c.internal_id < ${afterInternalId}`
+		: sql``;
+	return sql`${buildCheckpointWhere(checkpoint)} ${buildSearchWhere(search)} ${cursor}`;
 };

 /**
 * Full SELECT. Returns `{ internal_id, id }` rows newest-first via keyset
 * pagination on `c.internal_id DESC`, so successive iterations over an
 * unchanged customer set yield rows in the same order.
+ *
+ * Pure filter set only — the run path. To also surface already-processed
+ * customers (preview live view), use `buildProcessedPreviewSelect`.
 */
 export const buildCustomerSelect = ({
 	orgId,
@@ -92,7 +107,6 @@ export const buildCustomerSelect = ({
 	ctx,
 	checkpoint,
 	search,
-	includeProcessed,
 	limit,
 	afterInternalId,
 }: CustomerQueryArgs & {
@@ -100,17 +114,11 @@ export const buildCustomerSelect = ({
 	afterInternalId?: string;
 }): SQL => {
 	const where = compileWhere({ orgId, env, filter, ctx });
-	const checkpointWhere = buildCheckpointWhere(checkpoint);
-	const searchWhere = buildSearchWhere(search);
-	const processedOr = buildIncludeProcessedOr(includeProcessed);
-	const cursor = afterInternalId
-		? sql`AND c.internal_id < ${afterInternalId}`
-		: sql``;
 	const limitClause = limit !== undefined ? sql`LIMIT ${limit}` : sql``;
 	return sql`
 		SELECT c.internal_id, c.id, c.name, c.email
 		FROM customers c
-		WHERE ((${where}) ${processedOr}) ${checkpointWhere} ${searchWhere} ${cursor}
+		WHERE (${where}) ${buildCommonWhere({ checkpoint, search, afterInternalId })}
 		ORDER BY c.internal_id DESC
 		${limitClause}
 	`;
@@ -124,15 +132,81 @@ export const buildCustomerCount = ({
 	ctx,
 	checkpoint,
 	search,
-	includeProcessed,
 }: CustomerQueryArgs): SQL => {
 	const where = compileWhere({ orgId, env, filter, ctx });
-	const checkpointWhere = buildCheckpointWhere(checkpoint);
-	const searchWhere = buildSearchWhere(search);
-	const processedOr = buildIncludeProcessedOr(includeProcessed);
 	return sql`
 		SELECT COUNT(*)::bigint AS count
 		FROM customers c
-		WHERE ((${where}) ${processedOr}) ${checkpointWhere} ${searchWhere}
+		WHERE (${where}) ${buildCommonWhere({ checkpoint, search })}
+	`;
+};
+
+// ─── Preview-only: filter set ∪ already-processed set ────────────────
+// The live view surfaces customers an in-flight migration already ran for,
+// which the live filter no longer matches. We UNION the two scoped sets
+// rather than OR them: an `OR ... IN (...)` strips org/env scoping from the
+// customers scan and forces a full-table seq scan, whereas each UNION branch
+// keeps its own index. Equivalent to `(filter OR processed) AND <common>`
+// because `<common>` (checkpoint/search/cursor) is applied per branch.
+
+type ProcessedPreviewArgs = CustomerQueryArgs & {
+	includeProcessed: IncludeProcessed;
+};
+
+export const buildProcessedPreviewSelect = ({
+	orgId,
+	env,
+	filter,
+	ctx,
+	checkpoint,
+	search,
+	includeProcessed,
+	limit,
+	afterInternalId,
+}: ProcessedPreviewArgs & {
+	limit?: number;
+	afterInternalId?: string;
+}): SQL => {
+	const where = compileWhere({ orgId, env, filter, ctx });
+	const processed = buildProcessedIn(includeProcessed);
+	const limitClause = limit !== undefined ? sql`LIMIT ${limit}` : sql``;
+	return sql`
+		SELECT u.internal_id, u.id, u.name, u.email
+		FROM (
+			SELECT c.internal_id, c.id, c.name, c.email
+			FROM customers c
+			WHERE (${where}) ${buildCommonWhere({ checkpoint, search, afterInternalId })}
+			UNION
+			SELECT c.internal_id, c.id, c.name, c.email
+			FROM customers c
+			WHERE (${processed}) ${buildCommonWhere({ checkpoint, search, afterInternalId })}
+		) u
+		ORDER BY u.internal_id DESC
+		${limitClause}
+	`;
+};
+
+export const buildProcessedPreviewCount = ({
+	orgId,
+	env,
+	filter,
+	ctx,
+	checkpoint,
+	search,
+	includeProcessed,
+}: ProcessedPreviewArgs): SQL => {
+	const where = compileWhere({ orgId, env, filter, ctx });
+	const processed = buildProcessedIn(includeProcessed);
+	return sql`
+		SELECT COUNT(*)::bigint AS count
+		FROM (
+			SELECT c.internal_id
+			FROM customers c
+			WHERE (${where}) ${buildCommonWhere({ checkpoint, search })}
+			UNION
+			SELECT c.internal_id
+			FROM customers c
+			WHERE (${processed}) ${buildCommonWhere({ checkpoint, search })}
+		) u
 	`;
 };
--- a/server/src/internal/migrations/v2/filters/customers/filterCustomers.ts
+++ b/server/src/internal/migrations/v2/filters/customers/filterCustomers.ts
@@ -4,6 +4,8 @@ import { iterateOverFilterResults } from "../iterateOverFilterResults.js";
 import {
 	buildCustomerCount,
 	buildCustomerSelect,
+	buildProcessedPreviewCount,
+	buildProcessedPreviewSelect,
 	type CustomerCheckpointExclusion,
 	type IncludeProcessed,
 } from "./buildCustomerSelect.js";
@@ -40,13 +42,19 @@ export const filterCustomers = ({
 		filter,
 		checkpoint,
 		search,
-		includeProcessed,
 		ctx: { features: ctx.features },
 	};
 	return iterateOverFilterResults<CustomerRow>({
 		db: ctx.db,
 		buildSelect: ({ limit, afterInternalId }) =>
-			buildCustomerSelect({ ...args, limit, afterInternalId }),
+			includeProcessed
+				? buildProcessedPreviewSelect({
+						...args,
+						includeProcessed,
+						limit,
+						afterInternalId,
+					})
+				: buildCustomerSelect({ ...args, limit, afterInternalId }),
 		batchSize,
 	});
 };
@@ -65,16 +73,19 @@ export const countCustomers = async ({
 	search?: string;
 	includeProcessed?: IncludeProcessed;
 }): Promise<number> => {
-	const [{ count }] = (await ctx.db.execute(
-		buildCustomerCount({
-			orgId: ctx.org.id,
-			env: ctx.env,
-			filter,
-			checkpoint,
-			search,
-			includeProcessed,
-			ctx: { features: ctx.features },
-		}),
-	)) as Array<{ count: bigint | number }>;
+	const args = {
+		orgId: ctx.org.id,
+		env: ctx.env,
+		filter,
+		checkpoint,
+		search,
+		ctx: { features: ctx.features },
+	};
+	const query = includeProcessed
+		? buildProcessedPreviewCount({ ...args, includeProcessed })
+		: buildCustomerCount(args);
+	const [{ count }] = (await ctx.db.execute(query)) as Array<{
+		count: bigint | number;
+	}>;
 	return Number(count);
 };