diff --git a/packages/db/src/heartbeat-context-snapshot-index-migration.test.ts b/packages/db/src/heartbeat-context-snapshot-index-migration.test.ts new file mode 100644 index 0000000000..e4147bc6dc --- /dev/null +++ b/packages/db/src/heartbeat-context-snapshot-index-migration.test.ts @@ -0,0 +1,65 @@ +import { readFile } from "node:fs/promises"; +import { fileURLToPath } from "node:url"; +import { describe, expect, it, afterEach } from "vitest"; +import postgres from "postgres"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./test-embedded-postgres.js"; + +const cleanups: Array<() => Promise> = []; +const support = await getEmbeddedPostgresTestSupport(); +const d = support.supported ? describe : describe.skip; + +afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); +}); + +d("heartbeat context_snapshot expression index migration", () => { + it("applies full migration chain and uses the new indexes", async () => { + const dbh = await startEmbeddedPostgresTestDatabase("pap16575-idx-"); + cleanups.push(() => dbh.cleanup()); + const sql = postgres(dbh.connectionString, { max: 1 }); + cleanups.push(async () => { await sql.end(); }); + + const idx = await sql`SELECT indexname FROM pg_indexes WHERE tablename IN ('heartbeat_runs','agent_wakeup_requests')`; + const names = idx.map((r) => r.indexname as string); + expect(names).toContain("heartbeat_runs_company_ctx_issue_created_idx"); + expect(names).toContain("heartbeat_runs_company_ctx_task_created_idx"); + expect(names).toContain("agent_wakeup_requests_company_payload_issue_idx"); + + await sql.unsafe("SET enable_seqscan = off"); + const plan = await sql.unsafe( + "EXPLAIN SELECT id FROM heartbeat_runs WHERE company_id = '00000000-0000-0000-0000-000000000001' AND context_snapshot ->> 'issueId' = 'x' ORDER BY created_at DESC, id DESC LIMIT 1", + ); + const planText = plan.map((r) => Object.values(r)[0]).join("\n"); + expect(planText).toContain("heartbeat_runs_company_ctx_issue_created_idx"); + + const taskPlan = await sql.unsafe( + "EXPLAIN SELECT id FROM heartbeat_runs WHERE company_id = '00000000-0000-0000-0000-000000000001' AND context_snapshot ->> 'taskId' = 'x' ORDER BY created_at DESC, id DESC LIMIT 1", + ); + const taskText = taskPlan.map((r) => Object.values(r)[0]).join("\n"); + expect(taskText).toContain("heartbeat_runs_company_ctx_task_created_idx"); + + const wakePlan = await sql.unsafe( + "EXPLAIN SELECT id FROM agent_wakeup_requests WHERE company_id = '00000000-0000-0000-0000-000000000001' AND status = 'deferred_issue_execution' AND payload ->> 'issueId' = 'x' LIMIT 1", + ); + const wakeText = wakePlan.map((r) => Object.values(r)[0]).join("\n"); + expect(wakeText).toContain("agent_wakeup_requests_company_payload_issue_idx"); + + // Idempotency: re-running the migration statements against an already + // migrated database must be a no-op, not an error. + const migrationSql = await readFile( + fileURLToPath(new URL("./migrations/0209_heartbeat_context_snapshot_indexes.sql", import.meta.url)), + "utf8", + ); + const statements = migrationSql + .split("--> statement-breakpoint") + .map((s) => s.trim()) + .filter((s) => s.length > 0); + expect(statements.length).toBeGreaterThan(0); + for (const statement of statements) { + await sql.unsafe(statement); + } + }, 240_000); +}); diff --git a/packages/db/src/migrations/0209_heartbeat_context_snapshot_indexes.sql b/packages/db/src/migrations/0209_heartbeat_context_snapshot_indexes.sql new file mode 100644 index 0000000000..9331d63450 --- /dev/null +++ b/packages/db/src/migrations/0209_heartbeat_context_snapshot_indexes.sql @@ -0,0 +1,4 @@ +CREATE INDEX IF NOT EXISTS "heartbeat_runs_company_ctx_issue_created_idx" ON "heartbeat_runs" USING btree ("company_id", ("context_snapshot" ->> 'issueId'), "created_at" DESC);--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "heartbeat_runs_company_ctx_task_created_idx" ON "heartbeat_runs" USING btree ("company_id", ("context_snapshot" ->> 'taskId'), "created_at" DESC);--> statement-breakpoint +-- paperclip:migration-safety-ignore large-create-index-not-concurrently: Drizzle migrations run transactionally, so CONCURRENTLY is unavailable. This expression index removes per-issue full scans of agent_wakeup_requests in recovery sweeps (hasActiveExecutionPath), which currently dominate server load; the one-time build lock is the lesser cost. +CREATE INDEX IF NOT EXISTS "agent_wakeup_requests_company_payload_issue_idx" ON "agent_wakeup_requests" USING btree ("company_id", ("payload" ->> 'issueId')); diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 07ed0c584e..06f3456dba 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -1450,6 +1450,13 @@ "when": 1785988680096, "tag": "0208_keen_sharon_carter", "breakpoints": true + }, + { + "idx": 209, + "version": "7", + "when": 1786020026023, + "tag": "0209_heartbeat_context_snapshot_indexes", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/agent_wakeup_requests.ts b/packages/db/src/schema/agent_wakeup_requests.ts index 7b0ec64fa9..194e1732a3 100644 --- a/packages/db/src/schema/agent_wakeup_requests.ts +++ b/packages/db/src/schema/agent_wakeup_requests.ts @@ -40,5 +40,9 @@ export const agentWakeupRequests = pgTable( reviewPathRecoveryIdempotencyUq: uniqueIndex("agent_wakeup_requests_review_path_recovery_idempotency_uq") .on(table.companyId, table.idempotencyKey) .where(sql`${table.idempotencyKey} LIKE 'issue_review_path_lost:%' AND ${table.status} <> 'skipped'`), + companyPayloadIssueIdx: index("agent_wakeup_requests_company_payload_issue_idx").on( + table.companyId, + sql`(${table.payload} ->> 'issueId')`, + ), }), ); diff --git a/packages/db/src/schema/heartbeat_runs.ts b/packages/db/src/schema/heartbeat_runs.ts index 0f3590bb0f..49d3dd8f61 100644 --- a/packages/db/src/schema/heartbeat_runs.ts +++ b/packages/db/src/schema/heartbeat_runs.ts @@ -1,3 +1,4 @@ +import { sql } from "drizzle-orm"; import { type AnyPgColumn, pgTable, uuid, text, timestamp, jsonb, index, integer, bigint, boolean } from "drizzle-orm/pg-core"; import { companies } from "./companies.js"; import { agents } from "./agents.js"; @@ -88,5 +89,15 @@ export const heartbeatRuns = pgTable( table.companyId, table.createdAt.desc(), ), + companyCtxIssueCreatedIdx: index("heartbeat_runs_company_ctx_issue_created_idx").on( + table.companyId, + sql`(${table.contextSnapshot} ->> 'issueId')`, + table.createdAt.desc(), + ), + companyCtxTaskCreatedIdx: index("heartbeat_runs_company_ctx_task_created_idx").on( + table.companyId, + sql`(${table.contextSnapshot} ->> 'taskId')`, + table.createdAt.desc(), + ), }), );