Add queue tables

A table with `queue_table` set works as a job queue that many workers can
pull from at once. Claims use SELECT ... FOR UPDATE SKIP LOCKED over the
idx_queue_claim index order, so each job goes to one worker and only the
claimed rows are locked.

- Schema: adds the queue_* columns and the idx_queue_claim /
  idx_queue_lease indexes. Declaring these names yourself is an error.
- `queue_table` accepts `true` or { max_attempts, lease_ms, keep_completed }
- BunMariaDB.queue: push, claim, ack, fail, extend, cancel
- Statuses: pending, processing, done, failed, cancelled
- Leases use the database clock. Expired leases are reclaimed (or
  dead-lettered on the last attempt) and old lock tokens stop working.
- insert/update can't set queue state fields. update/delete skip claimed
  jobs.

Tests: allow the test container's self-signed certificate, and give the
queue test its own client so it doesn't reuse one closed by db.test.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

src/lib/mariadb/db-delete.ts

@@ -7,6 +7,7 @@ import type {
 } from "../../types";
 import sqlGenerator from "../../utils/sql-generator";
 import { quoteIdentifier } from "../../utils/sql-identifier";
+import { isQueueTable, QUEUE_NOT_CLAIMED_SQL } from "../queue/queue-utils";
 
 type Params<
     Schema extends { [k: string]: any } = { [k: string]: any },
@@ -38,7 +39,7 @@ export default async function DbDelete<
             targetId,
         });
 
-        const whereClause = sqlObj.where;
+        let whereClause = sqlObj.where;
 
         if (!whereClause) {
             return {
@@ -50,6 +51,10 @@ export default async function DbDelete<
             };
         }
 
+        if (isQueueTable(table)) {
+            whereClause += ` AND ${QUEUE_NOT_CLAIMED_SQL}`;
+        }
+
         let sql = `DELETE FROM ${quoteIdentifier(table)} ${whereClause}`;
         sql += ` RETURNING *`;
         sqlObj.string = sql;

src/lib/mariadb/db-insert.ts

@@ -7,6 +7,10 @@ import type {
 import sqlInsertGenerator from "../../utils/sql-insert-generator";
 import { sanitizeHtmlFieldsBatch } from "../../utils/sanitize-html-fields";
 import grabDuplicateSafeInsertSql from "../grab-duplicate-safe-insert-sql";
+import {
+    assertNoProtectedQueueFields,
+    isQueueTable,
+} from "../queue/queue-utils";
 
 type Params<
     Schema extends { [k: string]: any } = { [k: string]: any },
@@ -30,6 +34,10 @@ export default async function DbInsert<
     let sqlObj: SQLInsertGenReturn | null = null;
 
     try {
+        if (isQueueTable(table)) {
+            assertNoProtectedQueueFields(table, data);
+        }
+
         const sanitizedData = sanitizeHtmlFieldsBatch({
             table,
             data,

src/lib/mariadb/db-update.ts

@@ -10,6 +10,11 @@ import sanitizeHtmlFields from "../../utils/sanitize-html-fields";
 import { quoteIdentifier } from "../../utils/sql-identifier";
 import sqlValuePlaceholder from "../../utils/sql-value-placeholder";
 import grabTableFields from "../../utils/grab-table-fields";
+import {
+    assertNoProtectedQueueFields,
+    isQueueTable,
+    QUEUE_NOT_CLAIMED_SQL,
+} from "../queue/queue-utils";
 
 type Params<
     Schema extends { [k: string]: any } = { [k: string]: any },
@@ -55,7 +60,7 @@ export default async function DbUpdate<
             targetId,
         });
 
-        const whereClause = sqlQueryObj.where;
+        let whereClause = sqlQueryObj.where;
 
         if (!whereClause) {
             return {
@@ -64,6 +69,11 @@ export default async function DbUpdate<
             };
         }
 
+        if (isQueueTable(table)) {
+            assertNoProtectedQueueFields(table, [data]);
+            whereClause += ` AND ${QUEUE_NOT_CLAIMED_SQL}`;
+        }
+
         const sanitizedData = sanitizeHtmlFields({
             table,
             data,

src/lib/queue/queue-ack.ts

@@ -0,0 +1,83 @@
+import dbHandler from "../db-handler";
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Mark a claimed job as done: deletes it, or sets it to `done` when the
+ * table has `keep_completed`.
+ *
+ * Fails when the lock token no longer matches, i.e. the lease expired and
+ * another worker reclaimed the job.
+ */
+export default async function QueueAck<Table extends string = string>({
+    table,
+    id,
+    lock_token,
+    config,
+}: Params<Table>): Promise<DBResponseObject> {
+    try {
+        const options = requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+        const claimedWhere = `WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`;
+
+        if (!options.keep_completed) {
+            const res = await dbHandler({
+                query: `DELETE FROM ${tableSql} ${claimedWhere} RETURNING *`,
+                values: [id, lock_token],
+                config,
+            });
+
+            if (!res.success) {
+                return { ...res, success: false };
+            }
+
+            if (!res.payload?.length) {
+                return {
+                    ...res,
+                    success: false,
+                    msg: "Job is not claimed with this lock token",
+                };
+            }
+
+            return res;
+        }
+
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = 'done', queue_lease_expires_at = NULL, queue_lock_token = NULL, updated_at = ${QUEUE_NOW_SQL} ${claimedWhere}`,
+            values: [id, lock_token],
+            config,
+        });
+
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+
+        return await dbHandler({
+            query: `SELECT * FROM ${tableSql} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    } catch (error: any) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}

src/lib/queue/queue-cancel.ts

@@ -0,0 +1,54 @@
+import dbHandler from "../db-handler";
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Take a `pending` job out of the queue without deleting it. Jobs a worker
+ * has already claimed can't be cancelled.
+ */
+export default async function QueueCancel<Table extends string = string>({
+    table,
+    id,
+    config,
+}: Params<Table>): Promise<DBResponseObject> {
+    try {
+        requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = 'cancelled', updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'pending'`,
+            values: [id],
+            config,
+        });
+
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not pending",
+            };
+        }
+
+        return await dbHandler({
+            query: `SELECT * FROM ${tableSql} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    } catch (error: any) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}

src/lib/queue/queue-claim.ts

@@ -0,0 +1,122 @@
+import dbHandler, { holdEventLoop, resolveClient } from "../db-handler";
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+
+type Params<Table extends string = string> = {
+    table: Table;
+    /** Max jobs to claim. Defaults to 1 */
+    limit?: number;
+    /** How long the claim lasts before another worker may take the job. Defaults to the table's `lease_ms` */
+    lease_ms?: number;
+    /** Jobs whose lease expired after this many attempts are marked `failed`. Defaults to the table's `max_attempts` */
+    max_attempts?: number;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Claim the next available jobs, oldest `queue_available_at` first. Each
+ * returned row carries a `queue_lock_token` that must be passed to `ack`,
+ * `fail` or `extend`.
+ *
+ * `FOR UPDATE SKIP LOCKED` lets concurrent workers claim different jobs
+ * without blocking on each other or claiming the same job twice.
+ */
+export default async function QueueClaim<
+    Schema extends { [k: string]: any } = { [k: string]: any },
+    Table extends string = string,
+>({
+    table,
+    limit = 1,
+    lease_ms,
+    max_attempts,
+    config,
+}: Params<Table>): Promise<DBResponseObject<Schema>> {
+    try {
+        const options = requireQueueTable(table);
+
+        const tableSql = quoteIdentifier(table);
+        const finalLimit = Math.max(1, Math.floor(Number(limit) || 1));
+        const leaseMs = Math.max(
+            0,
+            Math.floor(Number(lease_ms ?? options.lease_ms) || 0),
+        );
+        const maxAttempts = Math.max(
+            1,
+            Math.floor(Number(max_attempts ?? options.max_attempts) || 1),
+        );
+
+        /**
+         * Release jobs whose worker died or stalled past its lease: back to
+         * `pending`, or `failed` if that was the last attempt.
+         */
+        const expired_res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = IF(queue_attempts >= ?, 'failed', 'pending'), queue_lease_expires_at = NULL, queue_lock_token = NULL, queue_last_error = 'Lease expired', updated_at = ${QUEUE_NOW_SQL} WHERE queue_status = 'processing' AND queue_lease_expires_at <= ${QUEUE_NOW_SQL}`,
+            values: [maxAttempts],
+            config,
+        });
+
+        if (!expired_res.success) {
+            return { ...expired_res, success: false } as DBResponseObject<Schema>;
+        }
+
+        const CLIENT = resolveClient(config);
+
+        if (!CLIENT) {
+            throw new Error(`Couldn't grab MariaDB Client.`);
+        }
+
+        const lock_token = crypto.randomUUID();
+
+        return await holdEventLoop(() =>
+            CLIENT.begin(async (tx) => {
+                /**
+                 * The ORDER BY must match `idx_queue_claim`. A filesort would
+                 * lock every pending row before LIMIT applies, and SKIP LOCKED
+                 * would then hide the whole queue from other workers.
+                 */
+                const matched_res = await dbHandler<{ id: number }>({
+                    query: `SELECT id FROM ${tableSql} FORCE INDEX (idx_queue_claim) WHERE queue_status = 'pending' AND queue_available_at <= ${QUEUE_NOW_SQL} ORDER BY queue_status, queue_available_at, id LIMIT ? FOR UPDATE SKIP LOCKED`,
+                    values: [finalLimit],
+                    client: tx,
+                });
+
+                if (!matched_res.success) {
+                    return {
+                        ...matched_res,
+                        success: false,
+                    } as DBResponseObject<Schema>;
+                }
+
+                const ids = (matched_res.payload || []).map((row) => row.id);
+
+                if (ids.length === 0) {
+                    return { success: true, payload: [] };
+                }
+
+                const idPlaceholders = ids.map(() => "?").join(",");
+
+                const update_res = await dbHandler({
+                    query: `UPDATE ${tableSql} SET queue_status = 'processing', queue_attempts = queue_attempts + 1, queue_lease_expires_at = ${QUEUE_NOW_SQL} + ?, queue_lock_token = ?, updated_at = ${QUEUE_NOW_SQL} WHERE id IN (${idPlaceholders})`,
+                    values: [leaseMs, lock_token, ...ids],
+                    client: tx,
+                });
+
+                if (!update_res.success) {
+                    throw new Error(update_res.error || "Queue claim failed");
+                }
+
+                return await dbHandler<Schema>({
+                    query: `SELECT * FROM ${tableSql} WHERE id IN (${idPlaceholders}) ORDER BY queue_available_at, id`,
+                    values: ids,
+                    client: tx,
+                });
+            }),
+        );
+    } catch (error: any) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}

src/lib/queue/queue-extend.ts

@@ -0,0 +1,64 @@
+import dbHandler from "../db-handler";
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    /** New lease length, counted from now. Defaults to the table's `lease_ms` */
+    lease_ms?: number;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Extend the lease on a claimed job, for jobs that run longer than the
+ * lease given to `claim`.
+ */
+export default async function QueueExtend<Table extends string = string>({
+    table,
+    id,
+    lock_token,
+    lease_ms,
+    config,
+}: Params<Table>): Promise<DBResponseObject> {
+    try {
+        const options = requireQueueTable(table);
+
+        const leaseMs = Math.max(
+            0,
+            Math.floor(Number(lease_ms ?? options.lease_ms) || 0),
+        );
+
+        const res = await dbHandler({
+            query: `UPDATE ${quoteIdentifier(table)} SET queue_lease_expires_at = ${QUEUE_NOW_SQL} + ?, updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`,
+            values: [leaseMs, id, lock_token],
+            config,
+        });
+
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+
+        return await dbHandler({
+            query: `SELECT * FROM ${quoteIdentifier(table)} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    } catch (error: any) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}

src/lib/queue/queue-fail.ts

@@ -0,0 +1,71 @@
+import dbHandler from "../db-handler";
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    error?: string;
+    /** Wait this long before the job can be retried. Defaults to 0 */
+    retry_delay_ms?: number;
+    /** Mark the job `failed` instead of retrying once it reaches this many attempts. Defaults to the table's `max_attempts` */
+    max_attempts?: number;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Release a claimed job after a failed attempt. It goes back to `pending`
+ * (after `retry_delay_ms`) or to `failed` once `max_attempts` is reached.
+ */
+export default async function QueueFail<Table extends string = string>({
+    table,
+    id,
+    lock_token,
+    error,
+    retry_delay_ms = 0,
+    max_attempts,
+    config,
+}: Params<Table>): Promise<DBResponseObject> {
+    try {
+        const options = requireQueueTable(table);
+
+        const tableSql = quoteIdentifier(table);
+        const retryDelayMs = Math.max(0, Math.floor(Number(retry_delay_ms) || 0));
+        const maxAttempts = Math.max(
+            1,
+            Math.floor(Number(max_attempts ?? options.max_attempts) || 1),
+        );
+
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = IF(queue_attempts >= ?, 'failed', 'pending'), queue_available_at = ${QUEUE_NOW_SQL} + ?, queue_lease_expires_at = NULL, queue_lock_token = NULL, queue_last_error = ?, updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`,
+            values: [maxAttempts, retryDelayMs, error ?? null, id, lock_token],
+            config,
+        });
+
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+
+        return await dbHandler({
+            query: `SELECT * FROM ${quoteIdentifier(table)} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    } catch (error: any) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}

src/lib/queue/queue-push.ts

@@ -0,0 +1,33 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+import DbInsert from "../mariadb/db-insert";
+
+type Params<
+    Schema extends { [k: string]: any } = { [k: string]: any },
+    Table extends string = string,
+> = {
+    table: Table;
+    data: Schema[];
+    /** Wait this long before the jobs can be claimed */
+    delay_ms?: number;
+    config?: BunMariaDBConfig;
+};
+
+/**
+ * Add jobs to a queue table.
+ */
+export default async function QueuePush<
+    Schema extends { [k: string]: any } = { [k: string]: any },
+    Table extends string = string,
+>({
+    table,
+    data,
+    delay_ms,
+    config,
+}: Params<Schema, Table>): Promise<DBResponseObject> {
+    const jobs = data.map((d) => ({
+        ...(delay_ms ? { queue_available_at: Date.now() + delay_ms } : {}),
+        ...d,
+    }));
+
+    return DbInsert({ table, data: jobs, config });
+}

src/lib/queue/queue-utils.ts

@@ -0,0 +1,98 @@
+import _ from "lodash";
+import { readLiveSchema } from "../../functions/live-schema";
+import {
+    QueueProtectedFields,
+    type BUN_MARIADB_DatabaseSchemaType,
+    type BUN_MARIADB_QueueTableOptions,
+} from "../../types";
+
+/**
+ * Current time in ms from the database clock. Leases are set and checked
+ * with this so workers on hosts with drifting clocks agree on expiry.
+ */
+export const QUEUE_NOW_SQL = "CAST(UNIX_TIMESTAMP(NOW(3)) * 1000 AS SIGNED)";
+
+/**
+ * Extra WHERE condition for generic update/delete on queue tables: rows a
+ * worker has claimed are only changed through the queue API.
+ */
+export const QUEUE_NOT_CLAIMED_SQL = "`queue_status` <> 'processing'";
+
+export type ResolvedQueueTableOptions = Required<BUN_MARIADB_QueueTableOptions>;
+
+const QUEUE_TABLE_DEFAULTS: ResolvedQueueTableOptions = {
+    max_attempts: 5,
+    lease_ms: 30_000,
+    keep_completed: false,
+};
+
+/**
+ * A table's `queue_table` options with defaults filled in (looked up
+ * directly or through its parent table), or `undefined` if it is not a
+ * queue table.
+ */
+export function grabQueueTableOptions(
+    tableName: string,
+): ResolvedQueueTableOptions | undefined {
+    const dbSchema: BUN_MARIADB_DatabaseSchemaType | undefined =
+        global.DB_SCHEMA || readLiveSchema();
+
+    const table = dbSchema?.tables?.find((t) => t.tableName === tableName);
+
+    if (!table) return undefined;
+
+    if (table.queue_table === undefined) {
+        return table.parentTableName
+            ? grabQueueTableOptions(table.parentTableName)
+            : undefined;
+    }
+
+    if (!table.queue_table) return undefined;
+
+    return {
+        ...QUEUE_TABLE_DEFAULTS,
+        ...(typeof table.queue_table === "object"
+            ? _.omitBy(table.queue_table, _.isUndefined)
+            : {}),
+    };
+}
+
+/**
+ * Whether a table is declared with `queue_table` (directly or through its
+ * parent table).
+ */
+export function isQueueTable(tableName: string): boolean {
+    return Boolean(grabQueueTableOptions(tableName));
+}
+
+/**
+ * Queue options for a table the queue API was called on. Throws if the table
+ * is not a queue table.
+ */
+export function requireQueueTable(tableName: string): ResolvedQueueTableOptions {
+    const options = grabQueueTableOptions(tableName);
+
+    if (!options) {
+        throw new Error(`\`${tableName}\` is not a queue table`);
+    }
+
+    return options;
+}
+
+/**
+ * Throw if `data` sets queue state the queue API owns.
+ */
+export function assertNoProtectedQueueFields(
+    table: string,
+    data: { [k: string]: any }[],
+): void {
+    const protectedFields = QueueProtectedFields.filter((field) =>
+        data.some((d) => d && d[field] !== undefined),
+    );
+
+    if (protectedFields.length > 0) {
+        throw new Error(
+            `\`${table}\` is a queue table: ${protectedFields.join(", ")} can only be changed through BunMariaDB.queue`,
+        );
+    }
+}