diff --git a/src/selfhost/pg-queue.ts b/src/selfhost/pg-queue.ts index 5df8f442d2..4454d1843a 100644 --- a/src/selfhost/pg-queue.ts +++ b/src/selfhost/pg-queue.ts @@ -527,7 +527,15 @@ export function createPgQueue( ): Promise { for (const m of messages) await enqueue(m.body, m.delaySeconds ?? 0); }, - } as unknown as Queue; + async snapshot() { + const res = await pool.query( + `SELECT payload, status, run_after FROM ${TABLE} WHERE status IN ('pending','processing','dead')`, + ); + return buildSelfHostQueueSnapshot( + res.rows as Array<{ payload: string; status: string; run_after: string | number }>, + ); + }, + } as unknown as Queue & { snapshot(): Promise }; return { binding, @@ -573,14 +581,7 @@ export function createPgQueue( async stats() { return readQueueStats(); }, - async snapshot() { - const res = await pool.query( - `SELECT payload, status, run_after FROM ${TABLE} WHERE status IN ('pending','processing','dead')`, - ); - return buildSelfHostQueueSnapshot( - res.rows as Array<{ payload: string; status: string; run_after: string | number }>, - ); - }, + snapshot: binding.snapshot, }; async function reclaimExpiredProcessingJobs(): Promise { diff --git a/src/selfhost/sqlite-queue.ts b/src/selfhost/sqlite-queue.ts index cefcc36bbb..5528a10a60 100644 --- a/src/selfhost/sqlite-queue.ts +++ b/src/selfhost/sqlite-queue.ts @@ -473,7 +473,15 @@ export function createSqliteQueue( ): Promise { for (const m of messages) enqueue(m.body, m.delaySeconds ?? 0); }, - } as unknown as Queue; + snapshot() { + return buildSelfHostQueueSnapshot( + driver.query( + `SELECT payload, status, run_after FROM ${TABLE} WHERE status IN ('pending','processing','dead')`, + [], + ).rows as Array<{ payload: string; status: string; run_after: number }>, + ); + }, + } as unknown as Queue & { snapshot(): SelfHostQueueSnapshot }; return { binding, @@ -521,14 +529,7 @@ export function createSqliteQueue( stats() { return readQueueStats(driver); }, - snapshot() { - return buildSelfHostQueueSnapshot( - driver.query( - `SELECT payload, status, run_after FROM ${TABLE} WHERE status IN ('pending','processing','dead')`, - [], - ).rows as Array<{ payload: string; status: string; run_after: number }>, - ); - }, + snapshot: binding.snapshot, }; } diff --git a/test/unit/selfhost-pg-queue.test.ts b/test/unit/selfhost-pg-queue.test.ts index e4f68b76f5..d4d60051a0 100644 --- a/test/unit/selfhost-pg-queue.test.ts +++ b/test/unit/selfhost-pg-queue.test.ts @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { Pool, QueryResult } from "pg"; import { createPgQueue } from "../../src/selfhost/pg-queue"; +import { queueSnapshotFromBinding } from "../../src/selfhost/queue-common"; import { renderMetrics, resetMetrics } from "../../src/selfhost/metrics"; import { RetryableJobError } from "../../src/queue/retryable"; import type { JobMessage } from "../../src/types"; @@ -1338,8 +1339,18 @@ describe("createPgQueue (durable #977)", () => { ], rowCount: 4, }); + m.fn.mockResolvedValueOnce({ + rows: [ + { payload: JSON.stringify(msg("agent-regate-pr")), status: "pending", run_after: String(now - 1) }, + { payload: JSON.stringify(msg("agent-regate-pr")), status: "processing", run_after: String(now - 1) }, + { payload: JSON.stringify(msg("github-webhook")), status: "pending", run_after: String(now + 60_000) }, + { payload: JSON.stringify(msg("rag-index-repo")), status: "dead", run_after: String(now - 1) }, + ], + rowCount: 4, + }); const snapshot = await q.snapshot(); + const bindingSnapshot = await queueSnapshotFromBinding(q.binding); expect(snapshot.totals).toMatchObject({ pending: 2, processing: 1, dead: 1 }); expect(snapshot.byType).toEqual( @@ -1350,5 +1361,6 @@ describe("createPgQueue (durable #977)", () => { { type: "rag-index-repo", status: "dead", count: 1, due: 0 }, ]), ); + expect(bindingSnapshot).toEqual(snapshot); }); }); diff --git a/test/unit/selfhost-sqlite-queue.test.ts b/test/unit/selfhost-sqlite-queue.test.ts index 298918e6d5..b6cf34a7dd 100644 --- a/test/unit/selfhost-sqlite-queue.test.ts +++ b/test/unit/selfhost-sqlite-queue.test.ts @@ -2,6 +2,7 @@ import { DatabaseSync } from "node:sqlite"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { nodeSqliteDriver } from "../../src/selfhost/d1-adapter"; import { createSqliteQueue } from "../../src/selfhost/sqlite-queue"; +import { queueSnapshotFromBinding } from "../../src/selfhost/queue-common"; import { renderMetrics, resetMetrics } from "../../src/selfhost/metrics"; import { RetryableJobError } from "../../src/queue/retryable"; import type { JobMessage } from "../../src/types"; @@ -837,6 +838,7 @@ describe("createSqliteQueue (durable #980)", () => { ); const snapshot = q.snapshot(); + const bindingSnapshot = await queueSnapshotFromBinding(q.binding); expect(snapshot.totals).toMatchObject({ pending: 2, processing: 1, dead: 1 }); expect(snapshot.byType).toEqual( @@ -847,6 +849,7 @@ describe("createSqliteQueue (durable #980)", () => { { type: "rag-index-repo", status: "dead", count: 1, due: 0 }, ]), ); + expect(bindingSnapshot).toEqual(snapshot); }); it("coalesces recurring maintenance jobs by semantic scope and keeps distinct scopes separate", async () => {