bun-sqlite/tests/queue.test.ts

import { describe, expect, beforeAll, beforeEach, test } from "bun:test";
import { Database } from "bun:sqlite";
import fs from "fs";
import os from "os";
import path from "path";
import createDBSchema from "../src/lib/schema/create-db-schema";
import appendDefaultFieldsToDbSchema from "../src/utils/append-default-fields-to-db-schema";
import DbInsert from "../src/lib/sqlite/db-insert";
import DbSelect from "../src/lib/sqlite/db-select";
import DbUpdate from "../src/lib/sqlite/db-update";
import DbDelete from "../src/lib/sqlite/db-delete";
import {
    QueueAdd,
    QueueClaim,
    QueueComplete,
    QueueFail,
    QueueReleaseStale,
} from "../src/lib/sqlite/db-queue";
import type {
    BunSQLiteConfig,
    BUN_SQLITE_DatabaseSchemaType,
} from "../src/types";

const db = new Database(":memory:");

const dbSchema: BUN_SQLITE_DatabaseSchemaType = {
    dbName: "test",
    tables: [
        {
            tableName: "emails",
            isQueue: { maxAttempts: 2, backoff: 1000, lockTimeout: 60_000 },
            fields: [{ fieldName: "recipient", dataType: "TEXT" }],
        },
        {
            tableName: "pings",
            isQueue: { deleteOnComplete: true },
            fields: [{ fieldName: "url", dataType: "TEXT" }],
        },
        {
            tableName: "users",
            fields: [{ fieldName: "name", dataType: "TEXT" }],
        },
    ],
};

global.CONFIG = { db_name: "test" };
global.DB_SCHEMA = dbSchema;
global.SQLITE_CLIENT = db;

beforeAll(() => {
    createDBSchema({
        db,
        db_schema: appendDefaultFieldsToDbSchema({ dbSchema }),
    });
});

beforeEach(() => {
    db.run(`DELETE FROM emails`);
    db.run(`DELETE FROM pings`);
});

describe("queue schema", () => {
    test("adds the queue columns and index", () => {
        const columns = (
            db.query(`PRAGMA table_info(emails)`).all() as { name: string }[]
        ).map((c) => c.name);

        expect(columns).toEqual([
            "id",
            "created_at",
            "updated_at",
            "recipient",
            "status",
            "attempts",
            "available_at",
            "locked_at",
            "locked_by",
            "last_error",
        ]);

        const indexes = (
            db.query(`PRAGMA index_list(emails)`).all() as { name: string }[]
        ).map((i) => i.name);

        expect(indexes).toContain("idx_emails_queue");
    });

    test("leaves non-queue tables alone", () => {
        const columns = (
            db.query(`PRAGMA table_info(users)`).all() as { name: string }[]
        ).map((c) => c.name);

        expect(columns).not.toContain("status");
    });

    test("rejects a table that is both a queue and a vector table", () => {
        expect(() =>
            createDBSchema({
                db: new Database(":memory:"),
                db_schema: {
                    tables: [
                        {
                            tableName: "bad",
                            isQueue: true,
                            isVector: true,
                            fields: [],
                        },
                    ],
                },
            }),
        ).toThrow("can't be both a queue and a vector table");
    });

    test("queue ops reject non-queue tables", async () => {
        const res = await QueueClaim({ table: "users" });

        expect(res.success).toBe(false);
        expect(res.error).toContain("not a queue table");
    });
});

describe("queue ops", () => {
    test("add, claim, complete", async () => {
        const added = await QueueAdd({
            table: "emails",
            data: [{ recipient: "[email protected]" }],
        });

        expect(added.success).toBe(true);
        expect(added.singleRes?.status).toBe("pending");

        const claimed = await QueueClaim({ table: "emails", worker: "w1" });

        expect(claimed.payload?.length).toBe(1);
        expect(claimed.singleRes?.status).toBe("processing");
        expect(claimed.singleRes?.attempts).toBe(1);
        expect(claimed.singleRes?.locked_by).toBe("w1");

        const empty = await QueueClaim({ table: "emails" });

        expect(empty.success).toBe(true);
        expect(empty.payload?.length).toBe(0);

        const done = await QueueComplete({
            table: "emails",
            id: claimed.singleRes!.id,
        });

        expect(done.success).toBe(true);
        expect(done.singleRes?.status).toBe("completed");
        expect(done.singleRes?.locked_by).toBeNull();
    });

    test("claims in FIFO order and respects limit", async () => {
        await QueueAdd({
            table: "emails",
            data: [{ recipient: "1" }, { recipient: "2" }, { recipient: "3" }],
        });

        const first = await QueueClaim({ table: "emails", limit: 2 });
        const second = await QueueClaim({ table: "emails", limit: 2 });

        expect(first.payload?.map((r) => r.recipient)).toEqual(["1", "2"]);
        expect(second.payload?.map((r) => r.recipient)).toEqual(["3"]);
    });

    test("concurrent claims never return the same job", async () => {
        await QueueAdd({
            table: "emails",
            data: Array.from({ length: 20 }, (_, i) => ({ recipient: String(i) })),
        });

        const results = await Promise.all(
            Array.from({ length: 10 }, (_, i) =>
                QueueClaim({ table: "emails", worker: `w${i}`, limit: 3 }),
            ),
        );

        const ids = results.flatMap((r) => r.payload!.map((row) => row.id));

        expect(ids.length).toBe(20);
        expect(new Set(ids).size).toBe(20);
    });

    test("delayed jobs are not claimed early", async () => {
        await QueueAdd({
            table: "emails",
            data: [{ recipient: "later" }],
            delay: 60_000,
        });

        const res = await QueueClaim({ table: "emails" });

        expect(res.payload?.length).toBe(0);
    });

    test("fail retries with backoff, then marks failed", async () => {
        await QueueAdd({ table: "emails", data: [{ recipient: "flaky" }] });

        const first = await QueueClaim({ table: "emails" });
        const id = first.singleRes!.id;

        const before = Date.now();
        const retry = await QueueFail({ table: "emails", id, error: "boom" });

        expect(retry.singleRes?.status).toBe("pending");
        expect(retry.singleRes?.last_error).toBe("boom");
        expect(retry.singleRes?.available_at).toBeGreaterThanOrEqual(
            before + 1000,
        );

        // Skip the backoff
        db.run(`UPDATE emails SET available_at = 0 WHERE id = ?`, [id]);

        const second = await QueueClaim({ table: "emails" });
        expect(second.singleRes?.attempts).toBe(2);

        const dead = await QueueFail({ table: "emails", id, error: "boom 2" });

        expect(dead.singleRes?.status).toBe("failed");
        expect(dead.singleRes?.last_error).toBe("boom 2");
    });

    test("complete and fail only touch processing jobs", async () => {
        const added = await QueueAdd({ table: "emails", data: [{ recipient: "x" }] });
        const id = added.singleRes!.id;

        expect((await QueueComplete({ table: "emails", id })).success).toBe(
            false,
        );
        expect((await QueueFail({ table: "emails", id })).success).toBe(false);
    });

    test("stale locks are released", async () => {
        await QueueAdd({ table: "emails", data: [{ recipient: "stuck" }] });
        const claimed = await QueueClaim({ table: "emails" });
        const id = claimed.singleRes!.id;

        db.run(`UPDATE emails SET locked_at = ? WHERE id = ?`, [
            Date.now() - 120_000,
            id,
        ]);

        const released = await QueueReleaseStale({ table: "emails" });

        expect(released.singleRes?.id).toBe(id);
        expect(released.singleRes?.status).toBe("pending");
        expect(released.singleRes?.locked_at).toBeNull();
    });

    test("stale locks out of attempts are marked failed", async () => {
        await QueueAdd({ table: "emails", data: [{ recipient: "stuck" }] });
        const claimed = await QueueClaim({ table: "emails" });
        const id = claimed.singleRes!.id;

        db.run(`UPDATE emails SET attempts = 2, locked_at = 0 WHERE id = ?`, [
            id,
        ]);

        // claim releases stale locks before claiming
        const res = await QueueClaim({ table: "emails" });
        expect(res.payload?.length).toBe(0);

        const row = db.query(`SELECT * FROM emails WHERE id = ?`).get(id) as any;
        expect(row.status).toBe("failed");
        expect(row.last_error).toBe("Lock timed out");
    });

    test("deleteOnComplete removes the job", async () => {
        await QueueAdd({ table: "pings", data: [{ url: "https://x" }] });
        const claimed = await QueueClaim({ table: "pings" });

        const done = await QueueComplete({
            table: "pings",
            id: claimed.singleRes!.id,
        });

        expect(done.success).toBe(true);
        expect(
            (db.query(`SELECT COUNT(*) AS n FROM pings`).get() as any).n,
        ).toBe(0);
    });
});

describe("queue across processes", () => {
    test("workers in separate processes never claim the same job", async () => {
        const db_dir = fs.mkdtempSync(path.join(os.tmpdir(), "bun-sqlite-queue-"));

        try {
            const config: BunSQLiteConfig = {
                db_name: "queue.db",
                db_dir,
                wal_mode: true,
            };
            const schema: BUN_SQLITE_DatabaseSchemaType = {
                tables: [
                    {
                        tableName: "jobs",
                        isQueue: true,
                        fields: [{ fieldName: "n", dataType: "INTEGER" }],
                    },
                ],
            };

            const fileDb = new Database(path.join(db_dir, "queue.db"));
            fileDb.run("PRAGMA journal_mode = WAL;");
            createDBSchema({
                db: fileDb,
                db_schema: appendDefaultFieldsToDbSchema({ dbSchema: schema }),
            });

            const jobCount = 500;
            const insert = fileDb.prepare(
                "INSERT INTO jobs (n, status, attempts, available_at) VALUES (?, 'pending', 0, 0)",
            );
            fileDb.transaction(() => {
                for (let i = 0; i < jobCount; i++) insert.run(i);
            })();
            fileDb.close();

            const workers = Array.from({ length: 4 }, (_, i) =>
                Bun.spawn(
                    ["bun", path.join(import.meta.dir, "fixtures/queue-worker.ts")],
                    {
                        env: {
                            ...process.env,
                            QUEUE_TEST_CONFIG: JSON.stringify(config),
                            QUEUE_TEST_SCHEMA: JSON.stringify(schema),
                            QUEUE_TEST_WORKER: `w${i}`,
                        },
                        stdout: "pipe",
                        stderr: "pipe",
                    },
                ),
            );

            // Wait until every worker is up, then start them together
            while (
                fs.readdirSync(db_dir).filter((f) => f.startsWith("ready-"))
                    .length < workers.length
            ) {
                await Bun.sleep(10);
            }
            fs.writeFileSync(path.join(db_dir, "go"), "");

            const outputs = await Promise.all(
                workers.map(async (proc) => {
                    const [stdout, stderr, code] = await Promise.all([
                        new Response(proc.stdout).text(),
                        new Response(proc.stderr).text(),
                        proc.exited,
                    ]);

                    expect(stderr).toBe("");
                    expect(code).toBe(0);

                    return JSON.parse(stdout.trim().split("\n").pop()!) as number[];
                }),
            );

            const ids = outputs.flat();

            expect(ids.length).toBe(jobCount);
            expect(new Set(ids).size).toBe(jobCount);
            // More than one worker got jobs, so the claims really did
            // interleave (the split is uneven: a worker holding the write
            // lock tends to keep it)
            expect(outputs.filter((o) => o.length > 0).length).toBeGreaterThan(1);

            const check = new Database(path.join(db_dir, "queue.db"));
            const counts = check
                .query("SELECT status, COUNT(*) AS n FROM jobs GROUP BY status")
                .all();
            check.close();

            expect(counts).toEqual([{ status: "completed", n: jobCount }]);
        } finally {
            fs.rmSync(db_dir, { recursive: true, force: true });
        }
    }, 30_000);
});

describe("reserved-word names", () => {
    test("insert, select, update, upsert and delete work on `to` / `order`", async () => {
        db.run(
            "CREATE TABLE `order` (id INTEGER PRIMARY KEY AUTOINCREMENT, `to` TEXT UNIQUE, `group` TEXT, created_at INTEGER, updated_at INTEGER)",
        );

        global.DB_SCHEMA = {
            ...dbSchema,
            tables: [
                ...dbSchema.tables,
                {
                    tableName: "order",
                    fields: [
                        { fieldName: "to", dataType: "TEXT", unique: true },
                        { fieldName: "group", dataType: "TEXT" },
                    ],
                },
            ],
        };

        try {
            const inserted = await DbInsert({
                table: "order",
                data: [{ to: "[email protected]", group: "x" }],
            });
            expect(inserted.success).toBe(true);
            expect(inserted.singleRes?.to).toBe("[email protected]");

            const upserted = await DbInsert({
                table: "order",
                data: [{ to: "[email protected]", group: "y" }],
                update_on_duplicate: true,
            });
            expect(upserted.success).toBe(true);
            expect(upserted.singleRes?.group).toBe("y");

            const selected = await DbSelect({
                table: "order",
                query: {
                    query: { to: { value: "[email protected]" } },
                    selectFields: ["to", "group"],
                    order: { field: "group", strategy: "DESC" },
                    group: "to",
                },
            });
            expect(selected.success).toBe(true);
            expect(selected.payload).toEqual([
                { to: "[email protected]", group: "y" },
            ]);

            const updated = await DbUpdate({
                table: "order",
                data: { group: "z" },
                query: { query: { to: { value: "[email protected]" } } },
            });
            expect(updated.singleRes?.group).toBe("z");

            const deleted = await DbDelete({
                table: "order",
                query: { query: { to: { value: "[email protected]" } } },
            });
            expect(deleted.success).toBe(true);
            expect(
                (db.query("SELECT COUNT(*) AS n FROM `order`").get() as any).n,
            ).toBe(0);
        } finally {
            global.DB_SCHEMA = dbSchema;
        }
    });

    test("unknown columns fail instead of matching nothing", async () => {
        const res = await DbSelect({
            table: "emails",
            query: { query: { nope: { value: "x" } } },
        });

        expect(res.success).toBe(false);
        expect(res.error).toContain("no such column");
    });
});