From 7c470e0225ce5ade49f40abf3a2f90ef30f1408b Mon Sep 17 00:00:00 2001 From: Benjamin Toby Date: Sun, 18 Oct 2026 15:43:22 +0100 Subject: [PATCH] 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 --- README.md | 122 +++++++++ src/index.ts | 17 ++ src/lib/mariadb/db-delete.ts | 7 +- src/lib/mariadb/db-insert.ts | 8 + src/lib/mariadb/db-update.ts | 12 +- src/lib/queue/queue-ack.ts | 83 ++++++ src/lib/queue/queue-cancel.ts | 54 ++++ src/lib/queue/queue-claim.ts | 122 +++++++++ src/lib/queue/queue-extend.ts | 64 +++++ src/lib/queue/queue-fail.ts | 71 +++++ src/lib/queue/queue-push.ts | 33 +++ src/lib/queue/queue-utils.ts | 98 +++++++ src/types/index.ts | 109 ++++++++ .../append-default-fields-to-db-schema.ts | 34 ++- tests/db.test.ts | 3 +- tests/queue.test.ts | 258 ++++++++++++++++++ 16 files changed, 1091 insertions(+), 4 deletions(-) create mode 100644 src/lib/queue/queue-ack.ts create mode 100644 src/lib/queue/queue-cancel.ts create mode 100644 src/lib/queue/queue-claim.ts create mode 100644 src/lib/queue/queue-extend.ts create mode 100644 src/lib/queue/queue-fail.ts create mode 100644 src/lib/queue/queue-push.ts create mode 100644 src/lib/queue/queue-utils.ts create mode 100644 tests/queue.test.ts diff --git a/README.md b/README.md index 6244dc3..f5483f1 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ A schema-driven MariaDB manager for [Bun](https://bun.sh), featuring automatic s - [Update](#update) - [Delete](#delete) - [Raw SQL](#raw-sql) +- [Queue Tables](#queue-tables) - [Query API Reference](#query-api-reference) - [Vector Support](#vector-support) - [HTML Sanitization](#html-sanitization) @@ -243,6 +244,7 @@ interface BUN_MARIADB_TableSchemaType { parentTableName?: string; // inherit / merge fields from another table tableNameOld?: string; // rename: old name triggers ALTER TABLE RENAME collation?: "utf8mb4_bin" | "utf8mb4_unicode_520_ci"; + queue_table?: boolean | BUN_MARIADB_QueueTableOptions; // job queue, see Queue Tables } ``` @@ -688,6 +690,126 @@ console.log(res.payload); --- +## Queue Tables + +Set `queue_table: true` on a table to use it as a job queue. Any number of workers can pull from it at the same time: claims use `SELECT ... FOR UPDATE SKIP LOCKED` (MariaDB 10.6+), so no job is handed to two workers and workers don't wait on each other. + +```ts +{ + tableName: "email_jobs", + queue_table: true, + fields: [ + { fieldName: "to", dataType: "VARCHAR", dataLength: 255 }, + { fieldName: "subject", dataType: "VARCHAR", dataLength: 255 }, + ], +} +``` + +Pass an options object instead of `true` to set the queue's defaults. Arguments to the `queue` methods override them. + +```ts +queue_table: { + max_attempts: 5, // attempts before a job is marked `failed` (default 5) + lease_ms: 60_000, // how long a claim lasts (default 30s) + keep_completed: true, // `ack` marks jobs `done` instead of deleting them (default false) +} +``` + +`bunx bun-mariadb schema` adds these columns, plus the `idx_queue_claim` and `idx_queue_lease` indexes. They are managed by the library: declaring a field or index with one of these names on a queue table is an error. + +| Field | Type | Description | +| ------------------------ | --------------------------------------- | --------------------------------------------------- | +| `queue_status` | `ENUM(...)` | Job state, see below | +| `queue_available_at` | `BIGINT` (ms) | The job can't be claimed before this time | +| `queue_attempts` | `INT` | How many times the job has been claimed | +| `queue_lease_expires_at` | `BIGINT` (ms) | When the current claim expires | +| `queue_lock_token` | `VARCHAR(64)` | Token of the current claim | +| `queue_last_error` | `TEXT` | Error from the last failed attempt | + +Statuses: + +| Status | Meaning | +| ------------ | ----------------------------------------------------------------- | +| `pending` | Waiting to be claimed. The only status `claim` picks up | +| `processing` | Claimed by a worker | +| `done` | Acknowledged and kept (`keep_completed`). Otherwise jobs are deleted | +| `failed` | Used up its attempts (dead letter) | +| `cancelled` | Taken out of the queue with `queue.cancel` | + +The statuses are fixed because the queue's SQL depends on them. For your own job outcomes (`sent`, `bounced`, ...), add a field of your own and set it before calling `ack`. + +### Worker loop + +```ts +await BunMariaDB.queue.push({ + table: "email_jobs", + data: [{ to: "a@example.com", subject: "Hi" }], + delay_ms: 5_000, // optional +}); + +const { payload: jobs } = await BunMariaDB.queue.claim({ + table: "email_jobs", + limit: 10, + lease_ms: 60_000, +}); + +for (const job of jobs || []) { + try { + await sendEmail(job); + await BunMariaDB.queue.ack({ + table: "email_jobs", + id: job.id, + lock_token: job.queue_lock_token, + }); + } catch (e: any) { + await BunMariaDB.queue.fail({ + table: "email_jobs", + id: job.id, + lock_token: job.queue_lock_token, + error: e.message, + retry_delay_ms: 30_000, + }); + } +} +``` + +| Method | Description | +| ------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- | +| `queue.push({ table, data, delay_ms? })` | Insert jobs as `pending` | +| `queue.claim({ table, limit?, lease_ms?, max_attempts? })` | Claim up to `limit` (default 1) jobs, oldest first, for `lease_ms` (default 30s) | +| `queue.ack({ table, id, lock_token })` | Job done: deletes it, or marks it `done` with `keep_completed` | +| `queue.fail({ table, id, lock_token, error?, retry_delay_ms?, max_attempts? })` | Back to `pending` after `retry_delay_ms`, or `failed` once `max_attempts` (default 5) is reached | +| `queue.extend({ table, id, lock_token, lease_ms? })` | Extend the lease for a long-running job | +| `queue.cancel({ table, id })` | Take a `pending` job out of the queue. Claimed jobs can't be cancelled | + +Delivery is **at least once**. If a worker crashes or runs past its lease, the next `claim` puts the job back in the queue (or marks it `failed` on its last attempt), and the next claim gets a new `queue_lock_token`. After that, `ack` / `fail` / `extend` with the old token return `success: false`. Make job handlers idempotent. Lease times use the database clock, so worker clock drift doesn't matter. + +How regular CRUD treats queue tables: + +- `insert` / `update` refuse to set `queue_status`, `queue_attempts`, `queue_lease_expires_at`, `queue_lock_token` or `queue_last_error`. Only the queue API changes those. +- `update` / `delete` skip jobs that are `processing`, so a worker's claimed job can't change or disappear while it runs. +- `select` is unchanged. Filter on `queue_status` to get the rows you want: + +```ts +// Dead letters +await BunMariaDB.select({ + table: "email_jobs", + query: { query: { queue_status: { value: "failed" } } }, +}); + +// Everything still in the queue +await BunMariaDB.select({ + table: "email_jobs", + count: true, + query: { query: { queue_status: { value: ["pending", "processing"], equality: "IN" } } }, +}); +``` + +- To retry a dead-lettered job, delete it and push it again, or reset it with raw `sql`. +- With `keep_completed`, `done` rows stay until you delete them, e.g. `delete` with a `queue_status` / `updated_at` filter. + +--- + ## Query API Reference `query` on `select` / `update` / `delete` accepts `ServerQueryParam`: diff --git a/src/index.ts b/src/index.ts index db08514..7635563 100644 --- a/src/index.ts +++ b/src/index.ts @@ -4,6 +4,12 @@ import DbInsert from "./lib/mariadb/db-insert"; import DbSelect from "./lib/mariadb/db-select"; import DbSQL from "./lib/mariadb/db-sql"; import DbUpdate from "./lib/mariadb/db-update"; +import QueueAck from "./lib/queue/queue-ack"; +import QueueCancel from "./lib/queue/queue-cancel"; +import QueueClaim from "./lib/queue/queue-claim"; +import QueueExtend from "./lib/queue/queue-extend"; +import QueueFail from "./lib/queue/queue-fail"; +import QueuePush from "./lib/queue/queue-push"; import grabDbSchema from "./utils/grab-db-schema"; import grabJoinFieldsFromQueryObject from "./utils/grab-join-fields-from-query-object"; @@ -24,6 +30,14 @@ const BunMariaDB = { update: withInit(DbUpdate), delete: withInit(DbDelete), sql: withInit(DbSQL), + queue: { + push: withInit(QueuePush), + claim: withInit(QueueClaim), + ack: withInit(QueueAck), + fail: withInit(QueueFail), + extend: withInit(QueueExtend), + cancel: withInit(QueueCancel), + }, utils: { grab_db_schema: withInit(grabDbSchema), grab_join_fields_from_query_object: grabJoinFieldsFromQueryObject, @@ -50,5 +64,8 @@ export { DockerComposeServices, IndexTypes, DefaultFields, + QueueStatuses, + QueueFields, + QueueProtectedFields, RequiredENVs, } from "./types"; diff --git a/src/lib/mariadb/db-delete.ts b/src/lib/mariadb/db-delete.ts index 8ac5455..009fbef 100644 --- a/src/lib/mariadb/db-delete.ts +++ b/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; diff --git a/src/lib/mariadb/db-insert.ts b/src/lib/mariadb/db-insert.ts index a8f190c..e629aab 100644 --- a/src/lib/mariadb/db-insert.ts +++ b/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, diff --git a/src/lib/mariadb/db-update.ts b/src/lib/mariadb/db-update.ts index bcebed2..2c4a73a 100644 --- a/src/lib/mariadb/db-update.ts +++ b/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, diff --git a/src/lib/queue/queue-ack.ts b/src/lib/queue/queue-ack.ts new file mode 100644 index 0000000..07b9828 --- /dev/null +++ b/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: 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, + id, + lock_token, + config, +}: Params
): Promise { + 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, + }; + } +} diff --git a/src/lib/queue/queue-cancel.ts b/src/lib/queue/queue-cancel.ts new file mode 100644 index 0000000..51c6c9a --- /dev/null +++ b/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: 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, + id, + config, +}: Params
): Promise { + 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, + }; + } +} diff --git a/src/lib/queue/queue-claim.ts b/src/lib/queue/queue-claim.ts new file mode 100644 index 0000000..7d9251d --- /dev/null +++ b/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: 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
): Promise> { + 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; + } + + 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; + } + + 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({ + 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, + }; + } +} diff --git a/src/lib/queue/queue-extend.ts b/src/lib/queue/queue-extend.ts new file mode 100644 index 0000000..41c2458 --- /dev/null +++ b/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: 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, + id, + lock_token, + lease_ms, + config, +}: Params
): Promise { + 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, + }; + } +} diff --git a/src/lib/queue/queue-fail.ts b/src/lib/queue/queue-fail.ts new file mode 100644 index 0000000..05af54c --- /dev/null +++ b/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: 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, + id, + lock_token, + error, + retry_delay_ms = 0, + max_attempts, + config, +}: Params
): Promise { + 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, + }; + } +} diff --git a/src/lib/queue/queue-push.ts b/src/lib/queue/queue-push.ts new file mode 100644 index 0000000..6a7d2f0 --- /dev/null +++ b/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): Promise { + const jobs = data.map((d) => ({ + ...(delay_ms ? { queue_available_at: Date.now() + delay_ms } : {}), + ...d, + })); + + return DbInsert({ table, data: jobs, config }); +} diff --git a/src/lib/queue/queue-utils.ts b/src/lib/queue/queue-utils.ts new file mode 100644 index 0000000..532df01 --- /dev/null +++ b/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; + +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`, + ); + } +} diff --git a/src/types/index.ts b/src/types/index.ts index 9adae99..300f1fe 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -96,6 +96,25 @@ export interface BUN_MARIADB_TableSchemaType { */ childTableDbId?: string | number; collation?: (typeof MariaDBCollations)[number]; + /** + * Make this table a job queue. Adds the `QueueFields` columns and + * `QueueIndexes`, and routes state changes through `BunMariaDB.queue`. + * Pass an options object to set the queue's defaults. + */ + queue_table?: boolean | BUN_MARIADB_QueueTableOptions; +} + +/** + * Defaults for a queue table. Per-call arguments to `BunMariaDB.queue` + * override them. + */ +export interface BUN_MARIADB_QueueTableOptions { + /** Attempts before a job is marked `failed`. Defaults to 5 */ + max_attempts?: number; + /** How long a claim lasts before another worker may take the job. Defaults to 30s */ + lease_ms?: number; + /** `ack` marks jobs `done` and keeps them instead of deleting them */ + keep_completed?: boolean; } /** @@ -1639,6 +1658,96 @@ export const DefaultFields: BUN_MARIADB_FieldSchemaType[] = [ }, ]; +/** + * Lifecycle of a queue row. Only `pending` jobs are claimed. + * - `done`: acknowledged, kept because of `keep_completed` (otherwise deleted) + * - `failed`: used up its attempts (dead letter) + * - `cancelled`: removed from the queue with `queue.cancel` + */ +export const QueueStatuses = [ + "pending", + "processing", + "done", + "failed", + "cancelled", +] as const; + +export type QueueStatus = (typeof QueueStatuses)[number]; + +/** + * Fields added to every table with `queue_table: true`. + */ +export const QueueFields: BUN_MARIADB_FieldSchemaType[] = [ + { + fieldName: "queue_status", + dataType: "ENUM", + options: [...QueueStatuses], + notNullValue: true, + defaultValue: "pending", + fieldDescription: "Queue state of the job.", + }, + { + fieldName: "queue_available_at", + dataType: "BIGINT", + notNullValue: true, + defaultValue: 0, + fieldDescription: + "The job cannot be claimed before this time. (Unix Timestamp, ms)", + }, + { + fieldName: "queue_attempts", + dataType: "INT", + notNullValue: true, + defaultValue: 0, + fieldDescription: "How many times the job has been claimed.", + }, + { + fieldName: "queue_lease_expires_at", + dataType: "BIGINT", + fieldDescription: + "When the current claim expires and the job can be reclaimed. (Unix Timestamp, ms)", + }, + { + fieldName: "queue_lock_token", + dataType: "VARCHAR", + dataLength: 64, + fieldDescription: + "Token of the current claim. Required to ack, fail or extend the job.", + }, + { + fieldName: "queue_last_error", + dataType: "TEXT", + fieldDescription: "Error message from the last failed attempt.", + }, +]; + +/** + * Queue fields only the queue API may write. `queue_available_at` can be set + * on insert and update. + */ +export const QueueProtectedFields = [ + "queue_status", + "queue_attempts", + "queue_lease_expires_at", + "queue_lock_token", + "queue_last_error", +] as const; + +/** + * Indexes added to queue tables. Claims scan `idx_queue_claim` in order, so + * `LIMIT` stops the scan early and only the claimed rows get locked. + */ +export const QueueIndexes: BUN_MARIADB_IndexSchemaType[] = [ + { + indexName: "idx_queue_claim", + indexTableFields: ["queue_status", "queue_available_at"], + }, + { + indexName: "idx_queue_lease", + indexTableFields: ["queue_status", "queue_lease_expires_at"], + }, +]; + export type BunMariaDBQueryFieldValues< F extends string = string, T extends string = string, diff --git a/src/utils/append-default-fields-to-db-schema.ts b/src/utils/append-default-fields-to-db-schema.ts index 87a02c9..0bd9f1c 100644 --- a/src/utils/append-default-fields-to-db-schema.ts +++ b/src/utils/append-default-fields-to-db-schema.ts @@ -1,10 +1,18 @@ import _ from "lodash"; -import { DefaultFields, type BUN_MARIADB_DatabaseSchemaType } from "../types"; +import { + DefaultFields, + QueueFields, + QueueIndexes, + type BUN_MARIADB_DatabaseSchemaType, +} from "../types"; type Params = { dbSchema: BUN_MARIADB_DatabaseSchemaType; }; +const queueFieldNames = new Set(QueueFields.map((f) => f.fieldName)); +const queueIndexNames = new Set(QueueIndexes.map((i) => i.indexName)); + export default function ({ dbSchema }: Params): BUN_MARIADB_DatabaseSchemaType { const finaldbSchema = _.cloneDeep(dbSchema); finaldbSchema.tables = finaldbSchema.tables.map((t) => { @@ -13,6 +21,30 @@ export default function ({ dbSchema }: Params): BUN_MARIADB_DatabaseSchemaType { (f) => !f.fieldName?.match(/^(id|created_at|updated_at)$/), ); newTable.fields.unshift(...DefaultFields); + + if (newTable.queue_table) { + const reserved = [ + ...newTable.fields + .map((f) => f.fieldName) + .filter((name) => queueFieldNames.has(name)), + ...(newTable.indexes || []) + .map((i) => i.indexName) + .filter((name) => queueIndexNames.has(name)), + ]; + + if (reserved.length > 0) { + throw new Error( + `\`${newTable.tableName}\` is a queue table: ${reserved.join(", ")} ${reserved.length > 1 ? "are" : "is"} managed by bun-mariadb and can't be declared in the schema`, + ); + } + + newTable.fields.push(..._.cloneDeep(QueueFields)); + newTable.indexes = [ + ...(newTable.indexes || []), + ..._.cloneDeep(QueueIndexes), + ]; + } + return newTable; }); diff --git a/tests/db.test.ts b/tests/db.test.ts index ee29b30..6894a41 100644 --- a/tests/db.test.ts +++ b/tests/db.test.ts @@ -36,7 +36,8 @@ describe.skipIf(!TEST_DB)("MariaDB integration", () => { appendDefaults = (await import("../src/utils/append-default-fields-to-db-schema")).default; schemaLib = await import("../src/lib/schema/run-schema-query"); - global.CONFIG = { db_name: TEST_DB! }; + // The test container uses a self-signed certificate + global.CONFIG = { db_name: TEST_DB!, ssl_reject_unauthorized: false }; global.DB_SCHEMA = schema; setMariaDBClient({ config: global.CONFIG }); diff --git a/tests/queue.test.ts b/tests/queue.test.ts new file mode 100644 index 0000000..44eb7bc --- /dev/null +++ b/tests/queue.test.ts @@ -0,0 +1,258 @@ +/** + * Queue table integration tests against a real MariaDB. + * + * Skipped unless BUN_MARIADB_TEST_DB is set (see tests/db.test.ts). + */ +import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test"; +import type { BUN_MARIADB_DatabaseSchemaType } from "../src/types"; + +const TEST_DB = process.env.BUN_MARIADB_TEST_DB; +const TABLE = "bmt_jobs"; +const KEPT_TABLE = "bmt_jobs_kept"; + +const schema: BUN_MARIADB_DatabaseSchemaType = { + tables: [ + { + tableName: TABLE, + queue_table: true, + fields: [{ fieldName: "name", dataType: "VARCHAR" }], + }, + { + tableName: KEPT_TABLE, + queue_table: { keep_completed: true, max_attempts: 2 }, + fields: [{ fieldName: "name", dataType: "VARCHAR" }], + }, + ], +} as BUN_MARIADB_DatabaseSchemaType; + +test("queue tables refuse user-declared queue fields and indexes", async () => { + const appendDefaults = (await import("../src/utils/append-default-fields-to-db-schema")).default; + + expect(() => + appendDefaults({ + dbSchema: { + tables: [ + { + tableName: "q", + queue_table: true, + fields: [{ fieldName: "queue_status", dataType: "VARCHAR" }], + indexes: [{ indexName: "idx_queue_claim", indexTableFields: ["id"] }], + }, + ], + }, + }), + ).toThrow("queue_status, idx_queue_claim are managed by bun-mariadb"); +}); + +describe.skipIf(!TEST_DB)("Queue tables", () => { + let DB: typeof import("../src/index").default; + + beforeAll(async () => { + const setMariaDBClient = (await import("../src/functions/set-mariadb-client")).default; + const createDBSchema = (await import("../src/lib/schema/create-db-schema")).default; + const appendDefaults = (await import("../src/utils/append-default-fields-to-db-schema")).default; + + /** + * The test container uses a self-signed certificate. `db_config` gives + * this file its own client, since other test files close theirs. The + * client is created before `global.CONFIG` changes, so the previous + * file's closed global client isn't reused. + */ + const config = { + db_name: TEST_DB!, + ssl_reject_unauthorized: false, + db_config: { max: 10 }, + }; + setMariaDBClient({ config }); + global.CONFIG = config; + global.DB_SCHEMA = schema; + + DB = (await import("../src/index")).default; + + await createDBSchema({ + db_schema: appendDefaults({ dbSchema: schema }), + config: global.CONFIG, + }); + }); + + beforeEach(async () => { + await DB.sql({ sql: `DELETE FROM ${TABLE}` }); + await DB.sql({ sql: `DELETE FROM ${KEPT_TABLE}` }); + }); + + afterAll(async () => { + await DB?.sql({ sql: `DROP TABLE IF EXISTS ${TABLE}, ${KEPT_TABLE}` }); + await DB?.sql({ + sql: `DELETE FROM __db_schema_manager__ WHERE table_name IN ('${TABLE}', '${KEPT_TABLE}')`, + }); + await global.MARIADB_CLIENT?.close(); + }); + + test("schema adds queue columns and indexes", async () => { + const cols = await DB.sql({ + sql: "SELECT COLUMN_NAME FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = ?", + values: [TABLE], + }); + const names = cols.payload?.map((c) => c.COLUMN_NAME); + expect(names).toContain("queue_status"); + expect(names).toContain("queue_lock_token"); + + const idx = await DB.sql({ + sql: "SELECT COLUMN_NAME FROM information_schema.STATISTICS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = ? AND INDEX_NAME = 'idx_queue_claim' ORDER BY SEQ_IN_INDEX", + values: [TABLE], + }); + expect(idx.payload?.map((c) => c.COLUMN_NAME)).toEqual(["queue_status", "queue_available_at"]); + }); + + test("insert and update refuse protected queue fields", async () => { + const ins = await DB.insert({ table: TABLE, data: [{ name: "x", queue_status: "processing" }] }); + expect(ins.success).toBe(false); + expect(ins.error).toContain("queue_status"); + + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + const upd = await DB.update({ table: TABLE, query: { query: { name: { value: "a" } } }, data: { queue_attempts: 0 } }); + expect(upd.success).toBe(false); + }); + + test("claim, then ack with the lock token deletes the job", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + + const claim = await DB.queue.claim({ table: TABLE }); + expect(claim.success).toBe(true); + const job = claim.payload![0]; + expect(job.queue_status).toBe("processing"); + expect(job.queue_attempts).toBe(1); + expect(job.queue_lock_token).toBeTruthy(); + + expect((await DB.queue.claim({ table: TABLE })).payload).toEqual([]); + + const bad = await DB.queue.ack({ table: TABLE, id: job.id, lock_token: "nope" }); + expect(bad.success).toBe(false); + + const ack = await DB.queue.ack({ table: TABLE, id: job.id, lock_token: job.queue_lock_token }); + expect(ack.success).toBe(true); + expect((await DB.select({ table: TABLE })).payload).toEqual([]); + }); + + test("concurrent claims never return the same job twice", async () => { + await DB.queue.push({ table: TABLE, data: Array.from({ length: 20 }, (_, i) => ({ name: `j${i}` })) }); + + const claims = await Promise.all( + Array.from({ length: 8 }, () => DB.queue.claim({ table: TABLE, limit: 3 })), + ); + + expect(claims.every((c) => c.success)).toBe(true); + const ids = claims.flatMap((c) => c.payload!.map((j) => j.id)); + expect(new Set(ids).size).toBe(ids.length); + expect(ids.length).toBe(20); + }); + + test("update and delete skip claimed jobs", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + const job = (await DB.queue.claim({ table: TABLE })).payload![0]; + + const upd = await DB.update({ table: TABLE, targetId: job.id, data: { name: "b" } }); + expect(upd.success).toBe(false); + + await DB.delete({ table: TABLE, targetId: job.id }); + expect((await DB.select({ table: TABLE })).payload?.length).toBe(1); + }); + + test("fail retries, then dead-letters after max_attempts", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + + let job = (await DB.queue.claim({ table: TABLE })).payload![0]; + const retry = await DB.queue.fail({ table: TABLE, id: job.id, lock_token: job.queue_lock_token, error: "boom", max_attempts: 2 }); + expect(retry.success).toBe(true); + expect(retry.single_res?.queue_status).toBe("pending"); + expect(retry.single_res?.queue_last_error).toBe("boom"); + + job = (await DB.queue.claim({ table: TABLE })).payload![0]; + expect(job.queue_attempts).toBe(2); + const dead = await DB.queue.fail({ table: TABLE, id: job.id, lock_token: job.queue_lock_token, error: "boom", max_attempts: 2 }); + expect(dead.single_res?.queue_status).toBe("failed"); + + expect((await DB.queue.claim({ table: TABLE })).payload).toEqual([]); + }); + + test("an expired lease is reclaimed and the old token stops working", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + + const first = (await DB.queue.claim({ table: TABLE, lease_ms: 0 })).payload![0]; + const second = (await DB.queue.claim({ table: TABLE })).payload![0]; + expect(second.id).toBe(first.id); + expect(second.queue_attempts).toBe(2); + + expect((await DB.queue.ack({ table: TABLE, id: first.id, lock_token: first.queue_lock_token })).success).toBe(false); + expect((await DB.queue.extend({ table: TABLE, id: second.id, lock_token: second.queue_lock_token })).success).toBe(true); + }); + + test("an expired lease on the last attempt is dead-lettered", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }] }); + await DB.queue.claim({ table: TABLE, lease_ms: 0 }); + + expect((await DB.queue.claim({ table: TABLE, max_attempts: 1 })).payload).toEqual([]); + const row = (await DB.select({ table: TABLE })).single_res; + expect(row?.queue_status).toBe("failed"); + }); + + test("claims in push order and holds back delayed jobs", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "later" }], delay_ms: 60_000 }); + await DB.queue.push({ table: TABLE, data: [{ name: "first" }, { name: "second" }] }); + + const claim = await DB.queue.claim({ table: TABLE, limit: 5 }); + expect(claim.payload?.map((j) => j.name)).toEqual(["first", "second"]); + }); + + test("keep_completed: ack marks the job done and keeps it", async () => { + await DB.queue.push({ table: KEPT_TABLE, data: [{ name: "a" }] }); + const job = (await DB.queue.claim({ table: KEPT_TABLE })).payload![0]; + + const ack = await DB.queue.ack({ table: KEPT_TABLE, id: job.id, lock_token: job.queue_lock_token }); + expect(ack.success).toBe(true); + expect(ack.single_res?.queue_status).toBe("done"); + expect(ack.single_res?.queue_lock_token).toBeNull(); + + expect((await DB.queue.claim({ table: KEPT_TABLE })).payload).toEqual([]); + expect((await DB.queue.ack({ table: KEPT_TABLE, id: job.id, lock_token: job.queue_lock_token })).success).toBe(false); + }); + + test("table max_attempts is the default for fail", async () => { + await DB.queue.push({ table: KEPT_TABLE, data: [{ name: "a" }] }); + + for (const expected of ["pending", "failed"]) { + const job = (await DB.queue.claim({ table: KEPT_TABLE })).payload![0]; + const res = await DB.queue.fail({ table: KEPT_TABLE, id: job.id, lock_token: job.queue_lock_token }); + expect(res.single_res?.queue_status).toBe(expected); + } + }); + + test("cancel takes pending jobs out of the queue, but not claimed ones", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }, { name: "b" }] }); + const rows = (await DB.select({ table: TABLE, query: { order: { field: "id", strategy: "ASC" } } })).payload!; + + const cancel = await DB.queue.cancel({ table: TABLE, id: rows[0].id }); + expect(cancel.single_res?.queue_status).toBe("cancelled"); + + const claimed = (await DB.queue.claim({ table: TABLE, limit: 5 })).payload!; + expect(claimed.map((j) => j.name)).toEqual(["b"]); + + expect((await DB.queue.cancel({ table: TABLE, id: claimed[0].id })).success).toBe(false); + }); + + test("select filters queue rows by status", async () => { + await DB.queue.push({ table: TABLE, data: [{ name: "a" }, { name: "b" }, { name: "c" }] }); + const job = (await DB.queue.claim({ table: TABLE })).payload![0]; + await DB.queue.fail({ table: TABLE, id: job.id, lock_token: job.queue_lock_token, max_attempts: 1 }); + + const failed = await DB.select({ table: TABLE, query: { query: { queue_status: { value: "failed" } } } }); + expect(failed.payload?.map((r) => r.name)).toEqual(["a"]); + + const open = await DB.select({ + table: TABLE, + count: true, + query: { query: { queue_status: { value: ["pending", "processing"], equality: "IN" } } }, + }); + expect(open.count).toBe(2); + }); +});