diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 8c480ccc8e5..d6062cd9818 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -615,6 +615,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([ "0436_legal_reavers.sql", "0447_legacy_controller_process.sql", "0450_petite_frightful_four.sql", + "0453_curious_ben_parker.sql", "0384_funny_adam_warlock.sql", diff --git a/packages/db/src/heartbeat-workspace-companion-index.test.ts b/packages/db/src/heartbeat-workspace-companion-index.test.ts new file mode 100644 index 00000000000..438517501ea --- /dev/null +++ b/packages/db/src/heartbeat-workspace-companion-index.test.ts @@ -0,0 +1,63 @@ +import { createHash, randomUUID } from "node:crypto"; +import { readFile } from "node:fs/promises"; +import postgres from "postgres"; +import { describe, expect, it } from "vitest"; +import { applyPendingMigrations } from "./client.js"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./test-embedded-postgres.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const indexName = "heartbeat_runs_company_native_workspace_idx"; +const migrationFile = "0453_curious_ben_parker.sql"; + +(support.supported ? describe : describe.skip)("workspace companion expression index", () => { + it("builds outside a transaction and preserves ordered prepared/finalized references", async () => { + const database = await startEmbeddedPostgresTestDatabase("workspace-companion-index-"); + const sql = postgres(database.connectionString, { max: 1, onnotice: () => {} }); + try { + const migration = await readFile(new URL(`./migrations/${migrationFile}`, import.meta.url), "utf8"); + const hash = createHash("sha256").update(migration).digest("hex"); + await sql`DELETE FROM drizzle.__drizzle_migrations WHERE hash=${hash}`; + await sql.unsafe(`DROP INDEX ${indexName}`); + const companies = await sql`INSERT INTO companies(name,issue_prefix) VALUES ('Index fixture','CIX1'),('Foreign fixture','CIX2') RETURNING id`; + const company = companies[0]!.id, foreignCompany = companies[1]!.id; + const [agent] = await sql`INSERT INTO agents(company_id,name) VALUES (${company},'Local fixture') RETURNING id`; + const [foreignAgent] = await sql`INSERT INTO agents(company_id,name) VALUES (${foreignCompany},'Foreign fixture') RETURNING id`; + const issues = await sql`INSERT INTO issues(company_id,title) VALUES (${company},'Source issue'),(${company},'Other source issue') RETURNING id`; + const workspace = randomUUID(), expected: string[] = []; + for (const [state, issue] of [["prepared", issues[0]!.id], ["finalized", issues[1]!.id], ["prepared", null], ["finalized", null]] as const) { + const id = randomUUID(); expected.push(id); + await sql`INSERT INTO heartbeat_runs(id,company_id,agent_id,native_issue_id,runner_profile_json) + VALUES (${id},${company},${agent!.id},${issue},${sql.json({ nativeWorkspaceSync: { workspaceId: workspace, state } })})`; + } + await sql`INSERT INTO heartbeat_runs(company_id,agent_id,runner_profile_json) + VALUES (${foreignCompany},${foreignAgent!.id},${sql.json({ nativeWorkspaceSync: { workspaceId: workspace, state: "prepared" } })})`; + for (const sync of [{ workspaceId: { malformed: true }, state: "prepared" }, { workspaceId: workspace.toUpperCase(), state: "prepared" }, { workspaceId: workspace, state: "other" }, { workspaceId: workspace }, "malformed", null]) { + await sql`INSERT INTO heartbeat_runs(company_id,agent_id,runner_profile_json) + VALUES (${company},${agent!.id},${sql.json({ nativeWorkspaceSync: sync })})`; + } + const before = await sql`SELECT id,native_issue_id,runner_profile_json#>>'{nativeWorkspaceSync,workspaceId}' workspace_id, + runner_profile_json#>>'{nativeWorkspaceSync,state}' state FROM heartbeat_runs + WHERE company_id=${company} AND runner_profile_json#>>'{nativeWorkspaceSync,workspaceId}'=${workspace} + AND runner_profile_json#>>'{nativeWorkspaceSync,state}' IN ('prepared','finalized') ORDER BY id`; + expect(before.map(row => row.id)).toEqual(expected.sort()); + await applyPendingMigrations(database.connectionString); + const after = await sql`SELECT id,native_issue_id,runner_profile_json->'nativeWorkspaceSync'->>'workspaceId' workspace_id, + runner_profile_json->'nativeWorkspaceSync'->>'state' state FROM heartbeat_runs + WHERE company_id=${company} AND runner_profile_json->'nativeWorkspaceSync'->>'workspaceId'=${workspace} + AND runner_profile_json->'nativeWorkspaceSync'->>'state' IN ('prepared','finalized') ORDER BY id`; + expect(after).toEqual(before); + const [index] = await sql`SELECT indisvalid,indisready,pg_get_expr(indpred,indrelid) predicate,pg_get_indexdef(indexrelid) definition + FROM pg_index WHERE indexrelid=${indexName}::regclass`; + expect(index).toMatchObject({ indisvalid: true, indisready: true, predicate: null }); + expect(index!.definition).toContain("company_id"); + expect(index!.definition).toContain("nativeWorkspaceSync"); + expect(index!.definition).toContain("workspaceId"); + await applyPendingMigrations(database.connectionString); + const [history] = await sql`SELECT count(*)::int count FROM drizzle.__drizzle_migrations WHERE hash=${hash}`; + expect(history!.count).toBe(1); + } finally { + await sql.end(); + await database.cleanup(); + } + }, 90_000); +}); diff --git a/packages/db/src/migrations/0453_curious_ben_parker.sql b/packages/db/src/migrations/0453_curious_ben_parker.sql new file mode 100644 index 00000000000..176986a82d9 --- /dev/null +++ b/packages/db/src/migrations/0453_curious_ben_parker.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "heartbeat_runs_company_native_workspace_idx" ON "heartbeat_runs" USING btree ("company_id",("runner_profile_json"->'nativeWorkspaceSync'->>'workspaceId')); diff --git a/packages/db/src/migrations/meta/0448_snapshot.json b/packages/db/src/migrations/meta/0453_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0448_snapshot.json rename to packages/db/src/migrations/meta/0453_snapshot.json index 05c25835806..8082bdb3ed7 100644 --- a/packages/db/src/migrations/meta/0448_snapshot.json +++ b/packages/db/src/migrations/meta/0453_snapshot.json @@ -1,6 +1,6 @@ { - "id": "360fef2e-1ac3-4ae8-92a2-245e544b61da", - "prevId": "76f06e29-53a8-4ecc-809a-d52b6bbb6520", + "id": "d54bfeb4-dedb-4f8b-9195-a4d4aa67d71f", + "prevId": "64e8be1e-e0cb-4cd8-85e0-0fdcce5eafde", "version": "7", "dialect": "postgresql", "tables": { @@ -4102,6 +4102,28 @@ "method": "btree", "with": {} }, + "agent_wakeup_requests_issue_blocker_stalled_uq": { + "name": "agent_wakeup_requests_issue_blocker_stalled_uq", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "idempotency_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"agent_wakeup_requests\".\"idempotency_key\" LIKE 'issue_blocker_stalled:%' AND \"agent_wakeup_requests\".\"status\" IN ('queued', 'deferred_issue_execution', 'claimed', 'coalesced', 'completed')", + "concurrently": true, + "method": "btree", + "with": {} + }, "agent_wakeup_requests_company_payload_issue_idx": { "name": "agent_wakeup_requests_company_payload_issue_idx", "columns": [ @@ -7860,6 +7882,150 @@ }, "isRLSEnabled": false }, + "public.budget_resume_claims": { + "name": "budget_resume_claims", + "schema": "", + "columns": { + "run_id": { + "name": "run_id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "incident_id": { + "name": "incident_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "actor_user_id": { + "name": "actor_user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "actor_source": { + "name": "actor_source", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "claimed_at": { + "name": "claimed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "settled_at": { + "name": "settled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "reason": { + "name": "reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "request_id": { + "name": "request_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "budget_resume_claims_pending_idx": { + "name": "budget_resume_claims_pending_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"budget_resume_claims\".\"state\" = 'pending'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "budget_resume_claims_run_id_heartbeat_runs_id_fk": { + "name": "budget_resume_claims_run_id_heartbeat_runs_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "heartbeat_runs", + "columnsFrom": [ + "run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "budget_resume_claims_company_id_companies_id_fk": { + "name": "budget_resume_claims_company_id_companies_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "budget_resume_claims_incident_id_budget_incidents_id_fk": { + "name": "budget_resume_claims_incident_id_budget_incidents_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "budget_incidents", + "columnsFrom": [ + "incident_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.built_in_managed_resources": { "name": "built_in_managed_resources", "schema": "", @@ -17233,7 +17399,7 @@ "columnsTo": [ "id" ], - "onDelete": "set null", + "onDelete": "cascade", "onUpdate": "no action" }, "company_secret_proposals_secret_proposal_id_company_secret_proposals_id_fk": { @@ -21977,6 +22143,112 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.controller_instances": { + "name": "controller_instances", + "schema": "", + "columns": { + "boot_id": { + "name": "boot_id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "epoch": { + "name": "epoch", + "type": "bigserial", + "primaryKey": false, + "notNull": true + }, + "process_space": { + "name": "process_space", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "pid": { + "name": "pid", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "started_at": { + "name": "started_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "renewed_at": { + "name": "renewed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "lease_expires_at": { + "name": "lease_expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "fenced_at": { + "name": "fenced_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "fenced_by_epoch": { + "name": "fenced_by_epoch", + "type": "bigint", + "primaryKey": false, + "notNull": false + }, + "stopped_at": { + "name": "stopped_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "controller_instances_epoch_idx": { + "name": "controller_instances_epoch_idx", + "columns": [ + { + "expression": "epoch", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "controller_instances_lease_idx": { + "name": "controller_instances_lease_idx", + "columns": [ + { + "expression": "lease_expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.cost_events": { "name": "cost_events", "schema": "", @@ -35290,6 +35562,27 @@ "method": "btree", "with": {} }, + "heartbeat_runs_company_native_workspace_idx": { + "name": "heartbeat_runs_company_native_workspace_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"runner_profile_json\"->'nativeWorkspaceSync'->>'workspaceId')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_company_issue_created_idx": { "name": "heartbeat_runs_company_issue_created_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 8d913de1669..2b28994fd41 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -3158,6 +3158,13 @@ "when": 1794205579544, "tag": "0452_secret_proposal_secret_cascade", "breakpoints": true + }, + { + "idx": 453, + "version": "7", + "when": 1794260223713, + "tag": "0453_curious_ben_parker", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/heartbeat_runs.ts b/packages/db/src/schema/heartbeat_runs.ts index 34cb7475ee4..5b88537f32a 100644 --- a/packages/db/src/schema/heartbeat_runs.ts +++ b/packages/db/src/schema/heartbeat_runs.ts @@ -190,6 +190,9 @@ export const heartbeatRuns = pgTable( table.agentId, table.startedAt, ), + companyNativeWorkspaceIdx: index("heartbeat_runs_company_native_workspace_idx") + .on(table.companyId, sql`(${table.runnerProfileJson}->'nativeWorkspaceSync'->>'workspaceId')`) + .concurrently(), companyIssueCreatedIdx: index("heartbeat_runs_company_issue_created_idx").on( table.companyId, table.issueId, diff --git a/server/src/services/online-workspace-source.ts b/server/src/services/online-workspace-source.ts index 25ef69d479c..2c8d4a64cdb 100644 --- a/server/src/services/online-workspace-source.ts +++ b/server/src/services/online-workspace-source.ts @@ -202,11 +202,11 @@ export async function inspectOnlineWorkspaceSource(db: ReadDb, workspace: Worksp if (recipe && (!recipe.restoredAt || recipe.contents && !recipe.contentsRestoredAt || recipe.restorePoint && !recipe.restorePoint.expiredAt && !recipe.restorePoint.appliedAt)) return null; const runs = await db.select({ id: heartbeatRuns.id, issueId: heartbeatRuns.nativeIssueId, - workspaceId: sql`${heartbeatRuns.runnerProfileJson}#>>'{nativeWorkspaceSync,workspaceId}'`, - state: sql`${heartbeatRuns.runnerProfileJson}#>>'{nativeWorkspaceSync,state}'` }).from(heartbeatRuns).where(and( + workspaceId: sql`${heartbeatRuns.runnerProfileJson}->'nativeWorkspaceSync'->>'workspaceId'`, + state: sql`${heartbeatRuns.runnerProfileJson}->'nativeWorkspaceSync'->>'state'` }).from(heartbeatRuns).where(and( eq(heartbeatRuns.companyId, workspace.companyId), - inArray(sql`${heartbeatRuns.runnerProfileJson}#>>'{nativeWorkspaceSync,workspaceId}'`, included.map(node => node.row.id)), - sql`${heartbeatRuns.runnerProfileJson}#>>'{nativeWorkspaceSync,state}' in ('prepared','finalized')`, + inArray(sql`${heartbeatRuns.runnerProfileJson}->'nativeWorkspaceSync'->>'workspaceId'`, included.map(node => node.row.id)), + sql`${heartbeatRuns.runnerProfileJson}->'nativeWorkspaceSync'->>'state' in ('prepared','finalized')`, )).orderBy(heartbeatRuns.id); const runsByWorkspace = new Map(); for (const run of runs) { diff --git a/server/src/services/remote-workspace-online-source.test.ts b/server/src/services/remote-workspace-online-source.test.ts index 7d508fcfe5b..85527595038 100644 --- a/server/src/services/remote-workspace-online-source.test.ts +++ b/server/src/services/remote-workspace-online-source.test.ts @@ -797,7 +797,7 @@ describe.skipIf(!support.supported)("online ordinary physical source fence", () const selected = await stored(f.child.id), debug = db.$client.options.debug; const scans: string[] = []; db.$client.options.debug = (_connection, query) => { - if (query.includes('from "heartbeat_runs"') && query.includes("nativeWorkspaceSync,workspaceId")) scans.push(query); + if (query.includes('from "heartbeat_runs"') && query.includes("->'nativeWorkspaceSync'->>'workspaceId'")) scans.push(query); }; try { const scope = await inspectOnlineWorkspaceSource(db, selected);