rivetkit 2.3.4 → 2.3.5-rc.1

package/src/common/database/config.ts

@@ -23,6 +23,11 @@ export interface SqliteExecuteResult extends SqliteQueryResult {
 	lastInsertRowId?: number | null;
 }
 
+export interface SqliteBatchStatement {
+	sql: string;
+	params?: SqliteBindings;
+}
+
 export interface SqliteNativeMetrics {
 	requestBuildNs: number;
 	serializeNs: number;
@@ -43,6 +48,10 @@ export interface SqliteDatabase {
 		callback?: (row: unknown[], columns: string[]) => void,
 	): Promise<void>;
 	execute(sql: string, params?: SqliteBindings): Promise<SqliteExecuteResult>;
+	executeBatch(
+		statements: SqliteBatchStatement[],
+	): Promise<SqliteExecuteResult[]>;
+	beginTransaction(timeoutMs?: number): Promise<SqliteTransactionDatabase>;
 	run(sql: string, params?: SqliteBindings): Promise<void>;
 	query(sql: string, params?: SqliteBindings): Promise<SqliteQueryResult>;
 	nativeMetrics?():
@@ -52,6 +61,16 @@ export interface SqliteDatabase {
 	close(): Promise<void>;
 }
 
+export interface SqliteTransactionDatabase {
+	exec(
+		sql: string,
+		callback?: (row: unknown[], columns: string[]) => void,
+	): Promise<void>;
+	execute(sql: string, params?: SqliteBindings): Promise<SqliteExecuteResult>;
+	commit(): Promise<void>;
+	rollback(): Promise<void>;
+}
+
 /**
  * Provider for opening native databases from the active runtime.
  */
@@ -143,6 +162,11 @@ export type RawAccess = {
 	 * Executes a raw SQL query.
 	 */
 	execute: ExecuteFunction;
+	/** Runs a callback in an isolated SQLite transaction. */
+	transaction: <T>(
+		callback: (tx: RawAccess) => Promise<T> | T,
+		options?: { timeout?: number },
+	) => Promise<T>;
 	/**
 	 * Returns native SQLite metrics when the active runtime supports them.
 	 */

package/src/common/database/mod.test.ts

@@ -1,17 +1,21 @@
-import { describe, expect, test } from "vitest";
+import { pino } from "pino";
+import { beforeEach, describe, expect, test } from "vitest";
+import { configureBaseLogger } from "@/common/log";
 import type {
 	DatabaseProviderContext,
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteExecuteResult,
+	SqliteTransactionDatabase,
 } from "./config";
 import { db } from "./mod";
 
+let logLines: string[];
+
 class FakeSqliteDatabase implements SqliteDatabase {
-	executeCalls: {
-		sql: string;
-		params?: SqliteBindings;
-	}[] = [];
+	failSql = new Map<string, Error>();
+	executeCalls: { sql: string; params?: SqliteBindings }[] = [];
+	transactionTimeouts: Array<number | undefined> = [];
 
 	async exec(): Promise<void> {}
 
@@ -19,18 +23,36 @@ class FakeSqliteDatabase implements SqliteDatabase {
 		sql: string,
 		params?: SqliteBindings,
 	): Promise<SqliteExecuteResult> {
-		this.executeCalls.push({
-			sql,
-			params,
-		});
+		this.record(sql, params);
+		return emptyResult();
+	}
+
+	async beginTransaction(
+		timeoutMs?: number,
+	): Promise<SqliteTransactionDatabase> {
+		this.transactionTimeouts.push(timeoutMs);
+		this.record("BEGIN");
 		return {
-			columns: [],
-			rows: [],
-			changes: 0,
-			lastInsertRowId: null,
+			exec: async () => {},
+			execute: async (sql, params) => {
+				this.record(sql, params);
+				return emptyResult();
+			},
+			commit: async () => this.record("COMMIT"),
+			rollback: async () => this.record("ROLLBACK"),
 		};
 	}
 
+	async executeBatch(
+		statements: Array<{ sql: string; params?: SqliteBindings }>,
+	): Promise<SqliteExecuteResult[]> {
+		const results: SqliteExecuteResult[] = [];
+		for (const statement of statements) {
+			results.push(await this.execute(statement.sql, statement.params));
+		}
+		return results;
+	}
+
 	async run(sql: string, params?: SqliteBindings): Promise<void> {
 		await this.execute(sql, params);
 	}
@@ -41,6 +63,21 @@ class FakeSqliteDatabase implements SqliteDatabase {
 	}
 
 	async close(): Promise<void> {}
+
+	private record(sql: string, params?: SqliteBindings): void {
+		this.executeCalls.push({ sql, params });
+		const error = this.failSql.get(sql);
+		if (error) throw error;
+	}
+}
+
+function emptyResult(): SqliteExecuteResult {
+	return {
+		columns: [],
+		rows: [],
+		changes: 0,
+		lastInsertRowId: null,
+	};
 }
 
 function testProviderContext(
@@ -54,84 +91,136 @@ function testProviderContext(
 			batchDelete: async () => {},
 			deleteRange: async () => {},
 		},
-		nativeDatabaseProvider: {
-			open: async () => database,
-		},
+		nativeDatabaseProvider: { open: async () => database },
 	};
 }
 
 describe("db", () => {
-	test("runs onMigrate inside a sqlite savepoint", async () => {
+	beforeEach(() => {
+		logLines = [];
+		configureBaseLogger(
+			pino(
+				{ level: "warn", base: {}, timestamp: false },
+				{ write: (line: string) => logLines.push(line) },
+			),
+		);
+	});
+
+	test("runs onMigrate through the shared transaction handle", async () => {
 		const nativeDb = new FakeSqliteDatabase();
 		const provider = db({
 			onMigrate: async (client) => {
 				await client.execute(
-					"CREATE TABLE items(id INTEGER PRIMARY KEY, value TEXT)",
+					"CREATE TABLE items(id INTEGER PRIMARY KEY)",
 				);
-				await client.execute("SELECT COUNT(*) AS count FROM items");
 			},
 		});
 		const client = await provider.createClient(
 			testProviderContext(nativeDb),
 		);
-
 		await provider.onMigrate(client);
 
-		expect(nativeDb.executeCalls).toEqual([
-			{
-				sql: "SAVEPOINT __rivet_on_migrate",
-				params: undefined,
-			},
-			{
-				sql: "CREATE TABLE items(id INTEGER PRIMARY KEY, value TEXT)",
-				params: undefined,
-			},
-			{
-				sql: "SELECT COUNT(*) AS count FROM items",
-				params: undefined,
-			},
-			{
-				sql: "RELEASE SAVEPOINT __rivet_on_migrate",
-				params: undefined,
-			},
+		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
+			"BEGIN",
+			"SAVEPOINT __rivet_on_migrate",
+			"CREATE TABLE items(id INTEGER PRIMARY KEY)",
+			"RELEASE SAVEPOINT __rivet_on_migrate",
+			"COMMIT",
 		]);
+		expect(nativeDb.transactionTimeouts).toEqual([300_000]);
 	});
 
-	test("rolls back the migration savepoint when onMigrate fails", async () => {
+	test("rolls back migrations when onMigrate fails", async () => {
 		const nativeDb = new FakeSqliteDatabase();
 		const provider = db({
-			onMigrate: async (client) => {
-				await client.execute(
-					"CREATE TABLE items(id INTEGER PRIMARY KEY, value TEXT)",
-				);
+			onMigrate: async () => {
 				throw new Error("migration failed");
 			},
 		});
 		const client = await provider.createClient(
 			testProviderContext(nativeDb),
 		);
-
 		await expect(provider.onMigrate(client)).rejects.toThrow(
 			"migration failed",
 		);
+		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
+			"BEGIN",
+			"SAVEPOINT __rivet_on_migrate",
+			"ROLLBACK TO SAVEPOINT __rivet_on_migrate",
+			"RELEASE SAVEPOINT __rivet_on_migrate",
+			"ROLLBACK",
+		]);
+	});
 
-		expect(nativeDb.executeCalls).toEqual([
-			{
-				sql: "SAVEPOINT __rivet_on_migrate",
-				params: undefined,
-			},
-			{
-				sql: "CREATE TABLE items(id INTEGER PRIMARY KEY, value TEXT)",
-				params: undefined,
-			},
-			{
-				sql: "ROLLBACK TO SAVEPOINT __rivet_on_migrate",
-				params: undefined,
-			},
-			{
-				sql: "RELEASE SAVEPOINT __rivet_on_migrate",
-				params: undefined,
+	test("commits transaction work and forwards timeout", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		const client = await db().createClient(testProviderContext(nativeDb));
+		const value = await client.transaction(
+			async (tx) => {
+				await tx.execute(
+					"INSERT INTO items(value) VALUES (?)",
+					"inside",
+				);
+				return 42;
 			},
+			{ timeout: 120_000 },
+		);
+		expect(value).toBe(42);
+		expect(nativeDb.transactionTimeouts).toEqual([120_000]);
+		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
+			"BEGIN",
+			"INSERT INTO items(value) VALUES (?)",
+			"COMMIT",
+		]);
+	});
+
+	test("rolls back a transaction when the callback throws", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		const client = await db().createClient(testProviderContext(nativeDb));
+		await expect(
+			client.transaction(async (tx) => {
+				await tx.execute(
+					"INSERT INTO items(value) VALUES (?)",
+					"inside",
+				);
+				throw new Error("callback failed");
+			}),
+		).rejects.toThrow("callback failed");
+		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
+			"BEGIN",
+			"INSERT INTO items(value) VALUES (?)",
+			"ROLLBACK",
 		]);
 	});
+
+	test("validates transaction timeouts", async () => {
+		const client = await db().createClient(
+			testProviderContext(new FakeSqliteDatabase()),
+		);
+		for (const timeout of [0, -1, Number.NaN, Number.POSITIVE_INFINITY]) {
+			await expect(
+				client.transaction(async () => {}, { timeout }),
+			).rejects.toThrow("positive finite");
+		}
+	});
+
+	test("warns once for manual cross-call transactions and names the opt-out", async () => {
+		const client = await db().createClient(
+			testProviderContext(new FakeSqliteDatabase()),
+		);
+		await client.execute("BEGIN");
+		await client.execute("COMMIT");
+		expect(logLines).toHaveLength(1);
+		expect(JSON.parse(logLines[0] ?? "{}").msg).toContain(
+			"Set warnOnManualTransactions: false",
+		);
+	});
+
+	test("can suppress the manual transaction warning", async () => {
+		const client = await db({
+			warnOnManualTransactions: false,
+		}).createClient(testProviderContext(new FakeSqliteDatabase()));
+		await client.execute("BEGIN");
+		expect(logLines).toHaveLength(0);
+	});
 });

package/src/common/database/mod.ts

@@ -1,10 +1,23 @@
-import type { DatabaseProvider, RawAccess, SqliteDatabase } from "./config";
-import { isSqliteBindingObject, toSqliteBindings } from "./shared";
+import { getLogger } from "@/common/log";
+import type {
+	DatabaseProvider,
+	RawAccess,
+	SqliteDatabase,
+	SqliteTransactionDatabase,
+} from "./config";
+import {
+	isManualTransactionControl,
+	isSqliteBindingObject,
+	MIGRATION_TRANSACTION_TIMEOUT_MS,
+	toSqliteBindings,
+	validateTransactionTimeout,
+} from "./shared";
 
 export type { RawAccess } from "./config";
 
 interface DatabaseFactoryConfig {
 	onMigrate?: (db: RawAccess) => Promise<void> | void;
+	warnOnManualTransactions?: boolean;
 }
 
 function hasMultipleStatements(query: string): boolean {
@@ -14,6 +27,7 @@ function hasMultipleStatements(query: string): boolean {
 
 export function db({
 	onMigrate,
+	warnOnManualTransactions = true,
 }: DatabaseFactoryConfig = {}): DatabaseProvider<RawAccess> {
 	return {
 		createClient: async (ctx) => {
@@ -26,6 +40,7 @@ export function db({
 
 			const db = await nativeDatabaseProvider.open(ctx.actorId);
 			let closed = false;
+			let manualTransactionWarned = false;
 			const ensureOpen = () => {
 				if (closed) {
 					throw new Error(
@@ -34,7 +49,10 @@ export function db({
 				}
 			};
 
-			const client: RawAccess = {
+			const createClient = (
+				target: SqliteDatabase | SqliteTransactionDatabase,
+				transactionScoped = false,
+			): RawAccess => ({
 				execute: async <
 					TRow extends Record<string, unknown> = Record<
 						string,
@@ -45,6 +63,19 @@ export function db({
 					...args: unknown[]
 				): Promise<TRow[]> => {
 					ensureOpen();
+					if (
+						!transactionScoped &&
+						warnOnManualTransactions &&
+						!manualTransactionWarned &&
+						!hasMultipleStatements(query) &&
+						isManualTransactionControl(query)
+					) {
+						manualTransactionWarned = true;
+						getLogger("database").warn(
+							{ actorId: ctx.actorId },
+							"Manual cross-call SQLite transactions can interleave with other actor work. Use db.transaction() for coordinated transactions. Set warnOnManualTransactions: false in your db(...) configuration to disable this warning.",
+						);
+					}
 
 					const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
 					const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
@@ -57,7 +88,7 @@ export function db({
 								isSqliteBindingObject(args[0])
 									? toSqliteBindings(args[0])
 									: toSqliteBindings(args);
-							const { rows, columns } = await db.execute(
+							const { rows, columns } = await target.execute(
 								query,
 								bindings,
 							);
@@ -67,13 +98,16 @@ export function db({
 						}
 
 						if (!hasMultipleStatements(query)) {
-							const { rows, columns } = await db.execute(query);
+							const { rows, columns } = await target.execute(
+								query,
+								undefined,
+							);
 							return rows.map((row) =>
 								rowToObject<TRow>(row, columns),
 							);
 						}
 
-						return await execMultiStatement<TRow>(db, query);
+						return await execMultiStatement<TRow>(target, query);
 					} finally {
 						const durationMs = performance.now() - start;
 						ctx.metrics?.trackSql(query, durationMs);
@@ -92,6 +126,28 @@ export function db({
 						}
 					}
 				},
+				transaction: async <T>(
+					callback: (tx: RawAccess) => Promise<T> | T,
+					options?: { timeout?: number },
+				): Promise<T> => {
+					validateTransactionTimeout(options?.timeout);
+					const transaction = await db.beginTransaction(
+						options?.timeout,
+					);
+					const tx = createClient(transaction, true);
+					try {
+						const result = await callback(tx);
+						await transaction.commit();
+						return result;
+					} catch (error) {
+						try {
+							await transaction.rollback();
+						} catch {
+							// Preserve the callback or commit error after expiry cleanup.
+						}
+						throw error;
+					}
+				},
 				close: async () => {
 					if (!closed) {
 						closed = true;
@@ -99,12 +155,15 @@ export function db({
 					}
 				},
 				nativeMetrics: () => db.nativeMetrics?.() ?? null,
-			};
+			});
+			const client = createClient(db);
 			return client;
 		},
 		onMigrate: async (client) => {
 			if (onMigrate) {
-				await withMigrationSavepoint(client, () => onMigrate(client));
+				await withMigrationSavepoint(client, (leased) =>
+					onMigrate(leased),
+				);
 			}
 		},
 	};
@@ -122,7 +181,7 @@ function rowToObject<TRow extends Record<string, unknown>>(
 }
 
 async function execMultiStatement<TRow extends Record<string, unknown>>(
-	db: SqliteDatabase,
+	db: Pick<SqliteDatabase, "exec">,
 	query: string,
 ): Promise<TRow[]> {
 	const results: Record<string, unknown>[] = [];
@@ -138,19 +197,30 @@ async function execMultiStatement<TRow extends Record<string, unknown>>(
 
 async function withMigrationSavepoint<T>(
 	client: RawAccess,
-	callback: () => Promise<T> | T,
+	callback: (leased: RawAccess) => Promise<T> | T,
 ): Promise<T> {
-	await client.execute("SAVEPOINT __rivet_on_migrate");
-	try {
-		const result = await callback();
-		await client.execute("RELEASE SAVEPOINT __rivet_on_migrate");
-		return result;
-	} catch (error) {
-		try {
-			await client.execute("ROLLBACK TO SAVEPOINT __rivet_on_migrate");
-		} finally {
-			await client.execute("RELEASE SAVEPOINT __rivet_on_migrate");
-		}
-		throw error;
-	}
+	return await client.transaction(
+		async (transaction) => {
+			await transaction.execute("SAVEPOINT __rivet_on_migrate");
+			try {
+				const result = await callback(transaction);
+				await transaction.execute(
+					"RELEASE SAVEPOINT __rivet_on_migrate",
+				);
+				return result;
+			} catch (error) {
+				try {
+					await transaction.execute(
+						"ROLLBACK TO SAVEPOINT __rivet_on_migrate",
+					);
+				} finally {
+					await transaction.execute(
+						"RELEASE SAVEPOINT __rivet_on_migrate",
+					);
+				}
+				throw error;
+			}
+		},
+		{ timeout: MIGRATION_TRANSACTION_TIMEOUT_MS },
+	);
 }

package/src/common/database/native-database.test.ts

@@ -19,6 +19,15 @@ function deferred<T>() {
 }
 
 class FakeNativeDatabase implements JsNativeDatabaseLike {
+	async beginTransaction() {
+		return {
+			exec: async (_sql: string) => this.exec(),
+			execute: async (sql: string, params?: NativeParams) =>
+				this.execute(sql, params),
+			commit: async () => {},
+			rollback: async () => {},
+		};
+	}
 	active = 0;
 	maxActive = 0;
 	closed = false;