@axiom-lattice/pg-stores 3.3.13 → 3.4.0

package/src/__tests__/PostgreSQLChannelReplyJobStore.test.ts

@@ -0,0 +1,333 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+import type { CreateChannelReplyJobInput } from "@axiom-lattice/protocols";
+import { PostgreSQLChannelReplyJobStore } from "../stores/PostgreSQLChannelReplyJobStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+const SCHEMA = `channel_reply_job_test_${randomUUID().replace(/-/g, "").slice(0, 16)}`;
+
+function input(
+  overrides: Partial<CreateChannelReplyJobInput> = {},
+): CreateChannelReplyJobInput {
+  return {
+    runId: randomUUID(),
+    tenantId: randomUUID(),
+    assistantId: "assistant-1",
+    threadId: "thread-1",
+    adapterChannel: "lark",
+    channelInstallationId: "install-1",
+    replyTarget: { chatId: "chat-1" },
+    inputMessageId: "input-1",
+    availableAt: Date.now() - 1_000,
+    ...overrides,
+  };
+}
+
+maybeDescribe("PostgreSQLChannelReplyJobStore (PostgreSQL)", () => {
+  let store: PostgreSQLChannelReplyJobStore;
+  let pool: Pool;
+
+  beforeAll(async () => {
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`CREATE SCHEMA ${SCHEMA}`);
+    await admin.end();
+
+    const poolConfig = {
+      connectionString: TEST_DATABASE_URL!,
+      options: `-c search_path=${SCHEMA}`,
+    };
+    store = new PostgreSQLChannelReplyJobStore({ poolConfig });
+    await store.initialize();
+    pool = new Pool(poolConfig);
+  });
+
+  beforeEach(async () => {
+    await pool.query("DELETE FROM lattice_channel_reply_jobs");
+  });
+
+  afterAll(async () => {
+    await pool.end();
+    await store.dispose();
+
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`DROP SCHEMA IF EXISTS ${SCHEMA} CASCADE`);
+    await admin.end();
+  });
+
+  it("creates the table, indexes and required column types through its migration", async () => {
+    const table = await pool.query(
+      "SELECT to_regclass('lattice_channel_reply_jobs') AS name",
+    );
+    expect(table.rows[0].name).toBe("lattice_channel_reply_jobs");
+
+    const indexes = await pool.query(
+      `SELECT indexname FROM pg_indexes
+        WHERE schemaname = $1 AND tablename = 'lattice_channel_reply_jobs'`,
+      [SCHEMA],
+    );
+    const names = indexes.rows.map((row) => row.indexname);
+    expect(names).toEqual(
+      expect.arrayContaining([
+        "idx_lattice_channel_reply_jobs_due",
+        "idx_lattice_channel_reply_jobs_tenant_thread",
+      ]),
+    );
+
+    const columns = await pool.query(
+      `SELECT column_name, data_type, is_nullable
+         FROM information_schema.columns
+        WHERE table_schema = $1 AND table_name = 'lattice_channel_reply_jobs'`,
+      [SCHEMA],
+    );
+    const byName = Object.fromEntries(
+      columns.rows.map((row) => [row.column_name, row]),
+    );
+    expect(byName.run_id.data_type).toBe("text");
+    expect(byName.tenant_id.data_type).toBe("character varying");
+    for (const column of [
+      "adapter_channel",
+      "channel_installation_id",
+      "input_message_id",
+      "created_at",
+      "updated_at",
+    ]) {
+      expect(byName[column].is_nullable).toBe("NO");
+    }
+  });
+
+  it("creates a pending job with database-clock timestamps and reads it back scoped by tenant", async () => {
+    const created = input();
+    const before = Date.now();
+    const job = await store.create(created);
+
+    expect(job).toMatchObject({
+      ...created,
+      streamCursor: null,
+      status: "pending",
+      attempts: 0,
+      leaseOwner: null,
+      leaseExpiresAt: null,
+      lastError: null,
+      outputMessageId: null,
+      deliveredAt: null,
+    });
+    expect(job.createdAt).toBeGreaterThanOrEqual(before - 1_000);
+    expect(job.updatedAt).toBe(job.createdAt);
+    await expect(store.get(created.tenantId, created.runId)).resolves.toEqual(job);
+    await expect(store.get(randomUUID(), created.runId)).resolves.toBeNull();
+  });
+
+  it("supports non-UUID tenant and run ids", async () => {
+    const created = await store.create(
+      input({ tenantId: "default", runId: "run-default-1" }),
+    );
+
+    expect(created.tenantId).toBe("default");
+    expect(created.runId).toBe("run-default-1");
+    await expect(store.get("default", "run-default-1")).resolves.toMatchObject({
+      tenantId: "default",
+      runId: "run-default-1",
+    });
+    await expect(
+      store.claim("default", "run-default-1", "owner-a", 1_000),
+    ).resolves.not.toBeNull();
+  });
+
+  it("round-trips nested reply targets through JSONB", async () => {
+    const replyTarget = {
+      chatId: "chat-1",
+      frame: { id: "frame-1", nested: [1, 2, { ok: true }] },
+      text: "héllo",
+    };
+    const created = await store.create(input({ replyTarget }));
+
+    const loaded = await store.get(created.tenantId, created.runId);
+    expect(loaded?.replyTarget).toEqual(replyTarget);
+  });
+
+  it("rejects duplicate run ids and missing required fields", async () => {
+    const created = await store.create(input());
+
+    await expect(
+      store.create(input({ tenantId: created.tenantId, runId: created.runId })),
+    ).rejects.toThrow();
+    await expect(
+      store.create(input({ inputMessageId: undefined as unknown as string })),
+    ).rejects.toThrow(/inputMessageId/);
+    await expect(
+      store.create(input({ adapterChannel: "" })),
+    ).rejects.toThrow(/adapterChannel/);
+  });
+
+  it("lets exactly one concurrent claim win using the database clock", async () => {
+    const created = await store.create(input());
+    const before = Date.now();
+
+    const claims = await Promise.all([
+      store.claim(created.tenantId, created.runId, "owner-a", 1_000, 1),
+      store.claim(created.tenantId, created.runId, "owner-b", 1_000, 1),
+    ]);
+
+    const winners = claims.filter((job) => job !== null);
+    expect(winners).toHaveLength(1);
+    const winner = winners[0]!;
+    expect(winner.leaseExpiresAt).toBeGreaterThanOrEqual(before + 500);
+    expect((await store.get(created.tenantId, created.runId))?.leaseOwner).toBe(
+      winner.leaseOwner,
+    );
+  });
+
+  it("does not claim a job before it is available", async () => {
+    const created = await store.create(input({ availableAt: Date.now() + 60_000 }));
+
+    await expect(
+      store.claim(created.tenantId, created.runId, "owner-a", 1_000),
+    ).resolves.toBeNull();
+  });
+
+  it("reschedules for the lease owner and then allows a later claim", async () => {
+    const created = await store.create(input());
+    await store.claim(created.tenantId, created.runId, "owner-a", 1_000, 1);
+
+    await expect(
+      store.reschedule(created.tenantId, created.runId, "owner-b", "nope", Date.now(), 1),
+    ).resolves.toBeNull();
+
+    const availableAt = Date.now() - 1;
+    const rescheduled = await store.reschedule(
+      created.tenantId,
+      created.runId,
+      "owner-a",
+      "temporary failure",
+      availableAt,
+      1,
+    );
+    expect(rescheduled).toMatchObject({
+      status: "pending",
+      attempts: 1,
+      leaseOwner: null,
+      leaseExpiresAt: null,
+      lastError: "temporary failure",
+      availableAt,
+    });
+    expect(rescheduled!.updatedAt).toBeGreaterThan(1_000);
+
+    await expect(
+      store.claim(created.tenantId, created.runId, "owner-b", 1_000),
+    ).resolves.not.toBeNull();
+  });
+
+  it("marks a job delivered only for the lease owner", async () => {
+    const created = await store.create(input());
+    await store.claim(created.tenantId, created.runId, "owner-a", 1_000, 1);
+
+    await expect(
+      store.markDelivered(created.tenantId, created.runId, "owner-b", "out-1", 1),
+    ).resolves.toBeNull();
+
+    const delivered = await store.markDelivered(
+      created.tenantId,
+      created.runId,
+      "owner-a",
+      "out-1",
+      1,
+    );
+    expect(delivered).toMatchObject({
+      status: "delivered",
+      outputMessageId: "out-1",
+      leaseOwner: null,
+      leaseExpiresAt: null,
+    });
+    expect(delivered!.deliveredAt).toBeGreaterThan(1_000);
+    expect(delivered!.updatedAt).toBe(delivered!.deliveredAt);
+  });
+
+  it("marks a job failed and deletes it by tenant scope", async () => {
+    const created = await store.create(input());
+    await store.claim(created.tenantId, created.runId, "owner-a", 1_000, 1);
+
+    await expect(
+      store.markFailed(created.tenantId, created.runId, "owner-b", "nope", 1),
+    ).resolves.toBeNull();
+
+    const failed = await store.markFailed(
+      created.tenantId,
+      created.runId,
+      "owner-a",
+      "permanent failure",
+      1,
+    );
+    expect(failed).toMatchObject({
+      status: "failed",
+      leaseOwner: null,
+      leaseExpiresAt: null,
+      lastError: "permanent failure",
+      deliveredAt: null,
+    });
+    expect(failed!.updatedAt).toBeGreaterThan(1_000);
+
+    await expect(store.delete(randomUUID(), created.runId)).resolves.toBe(false);
+    await expect(store.delete(created.tenantId, created.runId)).resolves.toBe(true);
+    await expect(store.get(created.tenantId, created.runId)).resolves.toBeNull();
+    await expect(store.delete(created.tenantId, created.runId)).resolves.toBe(false);
+  });
+
+  it("prunes delivered and failed jobs older than the retention threshold", async () => {
+    const delivered = await store.create(input());
+    await store.claim(delivered.tenantId, delivered.runId, "owner-a", 1_000, 1);
+    await store.markDelivered(delivered.tenantId, delivered.runId, "owner-a", "run-1", 1_000);
+
+    const failed = await store.create(input());
+    await store.claim(failed.tenantId, failed.runId, "owner-a", 1_000, 1);
+    await store.markFailed(failed.tenantId, failed.runId, "owner-a", "dead", 1_000);
+
+    const fresh = await store.create(input());
+    await store.claim(fresh.tenantId, fresh.runId, "owner-a", 1_000, 1);
+    await store.markDelivered(fresh.tenantId, fresh.runId, "owner-a", "run-2", 5_000);
+
+    const pending = await store.create(input());
+
+    // PostgreSQL uses the database clock for updates; pin timestamps so the
+    // retention boundary is deterministic.
+    await pool.query(
+      "UPDATE lattice_channel_reply_jobs SET updated_at = 1000 WHERE run_id = ANY($1)",
+      [[delivered.runId, failed.runId]],
+    );
+    await pool.query(
+      "UPDATE lattice_channel_reply_jobs SET updated_at = 5000 WHERE run_id = $1",
+      [fresh.runId],
+    );
+
+    await expect(store.deleteTerminalBefore(2_000, 10)).resolves.toBe(2);
+    await expect(store.get(delivered.tenantId, delivered.runId)).resolves.toBeNull();
+    await expect(store.get(failed.tenantId, failed.runId)).resolves.toBeNull();
+    await expect(store.get(fresh.tenantId, fresh.runId)).resolves.toMatchObject({ status: "delivered" });
+    await expect(store.get(pending.tenantId, pending.runId)).resolves.toMatchObject({ status: "pending" });
+    await expect(store.deleteTerminalBefore(9_000, 10)).resolves.toBe(1);
+  });
+
+  it("lists due pending jobs with consistent limit clamping", async () => {
+    const now = Date.now();
+    const first = await store.create(input({ availableAt: now - 3_000 }));
+    const second = await store.create(input({ availableAt: now - 2_000 }));
+    const third = await store.create(input({ availableAt: now - 1_000 }));
+    await store.create(input({ availableAt: now + 60_000 }));
+
+    expect((await store.listPending(now, 0)).map((job) => job.runId)).toEqual([
+      first.runId,
+    ]);
+    expect((await store.listPending(now, -5)).map((job) => job.runId)).toEqual([
+      first.runId,
+    ]);
+    expect((await store.listPending(now, 2)).map((job) => job.runId)).toEqual([
+      first.runId,
+      second.runId,
+    ]);
+    expect((await store.listPending(now, 1_000)).map((job) => job.runId)).toEqual([
+      first.runId,
+      second.runId,
+      third.runId,
+    ]);
+  });
+});