rivetkit 2.3.4 → 2.3.5

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

@@ -1,9 +1,11 @@
 import { decodeBridgeRivetError } from "@/actor/errors";
 import type {
+	SqliteBatchStatement,
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteExecuteResult,
 	SqliteNativeMetrics,
+	SqliteTransactionDatabase,
 } from "./config";
 
 type NativeBindNoValues = {
@@ -65,12 +67,21 @@ interface NativeExecuteResult {
 	lastInsertRowId?: number | null;
 }
 
+interface NativeBatchStatement {
+	sql: string;
+	params?: NativeBindParam[] | null;
+}
+
 export interface JsNativeDatabaseLike {
 	exec(sql: string): Promise<NativeExecResult>;
 	execute(
 		sql: string,
 		params?: NativeBindParam[] | null,
 	): Promise<NativeExecuteResult>;
+	executeBatch?(
+		statements: NativeBatchStatement[],
+	): Promise<NativeExecuteResult[]>;
+	beginTransaction(timeoutMs?: number): Promise<JsNativeTransactionLike>;
 	query(
 		sql: string,
 		params?: NativeBindParam[] | null,
@@ -84,6 +95,16 @@ export interface JsNativeDatabaseLike {
 	close(): Promise<void>;
 }
 
+export interface JsNativeTransactionLike {
+	exec(sql: string): Promise<NativeExecResult>;
+	execute(
+		sql: string,
+		params?: NativeBindParam[] | null,
+	): Promise<NativeExecuteResult>;
+	commit(): Promise<void>;
+	rollback(): Promise<void>;
+}
+
 function shouldAttachNativeKvError(message: string): boolean {
 	return /i\/o error|unable to open database file/i.test(message);
 }
@@ -207,6 +228,15 @@ function toNativeBindings(
 	});
 }
 
+function toNativeBatchStatement(
+	statement: SqliteBatchStatement,
+): NativeBatchStatement {
+	return {
+		sql: statement.sql,
+		params: toNativeBindings(statement.sql, statement.params),
+	};
+}
+
 function normalizeNativeMetrics(
 	metrics: SqliteNativeMetrics | null | undefined,
 ): SqliteNativeMetrics | null {
@@ -343,6 +373,68 @@ export function wrapJsNativeDatabase(
 		): Promise<SqliteExecuteResult> {
 			return await executeNative(sql, params);
 		},
+		async executeBatch(
+			statements: SqliteBatchStatement[],
+		): Promise<SqliteExecuteResult[]> {
+			const nativeStatements = statements.map(toNativeBatchStatement);
+			const release = gate.enter();
+			let transaction: JsNativeTransactionLike | undefined;
+			try {
+				const results = database.executeBatch
+					? await database.executeBatch(nativeStatements)
+					: await (async () => {
+							transaction = await database.beginTransaction();
+							const transactionResults: NativeExecuteResult[] =
+								[];
+							for (const statement of nativeStatements) {
+								transactionResults.push(
+									await transaction.execute(
+										statement.sql,
+										statement.params,
+									),
+								);
+							}
+							await transaction.commit();
+							transaction = undefined;
+							return transactionResults;
+						})();
+				for (const result of results) {
+					if (result.lastInsertRowId !== undefined) {
+						lastInsertRowId = result.lastInsertRowId;
+					}
+				}
+				return results;
+			} catch (error) {
+				if (transaction) {
+					try {
+						await transaction.rollback();
+					} catch {
+						// The original batch error is more actionable than cleanup failure.
+					}
+				}
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+		},
+		async beginTransaction(
+			timeoutMs?: number,
+		): Promise<SqliteTransactionDatabase> {
+			const release = gate.enter();
+			let transaction: JsNativeTransactionLike;
+			try {
+				transaction = await database.beginTransaction(timeoutMs);
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+			return wrapTransaction(database, transaction, gate, (result) => {
+				if (result.lastInsertRowId !== undefined) {
+					lastInsertRowId = result.lastInsertRowId;
+				}
+			});
+		},
 		async run(sql: string, params?: SqliteBindings): Promise<void> {
 			await executeNative(sql, params);
 		},
@@ -360,6 +452,65 @@ export function wrapJsNativeDatabase(
 	};
 }
 
+function wrapTransaction(
+	database: JsNativeDatabaseLike,
+	transaction: JsNativeTransactionLike,
+	gate: NativeCloseGate,
+	onExecute: (result: NativeExecuteResult) => void,
+): SqliteTransactionDatabase {
+	return {
+		async exec(sql, callback) {
+			const release = gate.enter();
+			let result: NativeExecResult;
+			try {
+				result = await transaction.exec(sql);
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+			if (callback) {
+				for (const row of result.rows) callback(row, result.columns);
+			}
+		},
+		async execute(sql, params) {
+			const release = gate.enter();
+			try {
+				const result = await transaction.execute(
+					sql,
+					toNativeBindings(sql, params),
+				);
+				onExecute(result);
+				return result;
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+		},
+		async commit() {
+			const release = gate.enter();
+			try {
+				await transaction.commit();
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+		},
+		async rollback() {
+			const release = gate.enter();
+			try {
+				await transaction.rollback();
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+		},
+	};
+}
+
 function lastInsertRowIdColumnName(sql: string): string | undefined {
 	const match = sql.match(
 		/^\s*SELECT\s+last_insert_rowid\s*\(\s*\)\s*(?:AS\s+("[^"]+"|`[^`]+`|\[[^\]]+\]|\w+))?\s*;?\s*$/i,

package/src/common/database/shared.ts

@@ -1,5 +1,20 @@
 import type { SqliteBindings } from "./config";
 
+/** Migrations may legitimately do substantially more work than request transactions. */
+export const MIGRATION_TRANSACTION_TIMEOUT_MS = 5 * 60_000;
+
+export function isManualTransactionControl(query: string): boolean {
+	return /^\s*(?:BEGIN|SAVEPOINT|COMMIT|END|ROLLBACK)\b/i.test(query);
+}
+
+export function validateTransactionTimeout(timeout: number | undefined): void {
+	if (timeout !== undefined && (!Number.isFinite(timeout) || timeout <= 0)) {
+		throw new Error(
+			"db.transaction() timeout must be a positive finite number of milliseconds",
+		);
+	}
+}
+
 type SqliteBindingObject = Record<string, unknown>;
 
 function isSqliteBindingValue(value: unknown): boolean {

package/src/common/encoding.ts

@@ -473,27 +473,34 @@ function base64DecodeToArrayBuffer(base64: string): ArrayBuffer {
 }
 
 /** Stringifies with compat for values that BARE & CBOR supports. */
-export function jsonStringifyCompat(input: any): string {
-	return JSON.stringify(input, (_key, value) => {
-		if (typeof value === "bigint") {
-			return [JSON_COMPAT_BIGINT, value.toString()];
-		}
-		if (value instanceof ArrayBuffer) {
-			return [JSON_COMPAT_ARRAY_BUFFER, base64EncodeArrayBuffer(value)];
-		}
-		if (value instanceof Uint8Array) {
-			return [JSON_COMPAT_UINT8_ARRAY, base64EncodeUint8Array(value)];
-		}
-		if (
-			Array.isArray(value) &&
-			value.length === 2 &&
-			typeof value[0] === "string" &&
-			value[0].startsWith("$")
-		) {
-			return [`$${value[0]}`, value[1]];
-		}
-		return value;
-	});
+export function jsonStringifyCompat(
+	input: any,
+	space?: string | number,
+): string {
+	return JSON.stringify(
+		input,
+		(_key, value) => {
+			if (typeof value === "bigint") {
+				return [JSON_COMPAT_BIGINT, value.toString()];
+			}
+			if (value instanceof ArrayBuffer) {
+				return [JSON_COMPAT_ARRAY_BUFFER, base64EncodeArrayBuffer(value)];
+			}
+			if (value instanceof Uint8Array) {
+				return [JSON_COMPAT_UINT8_ARRAY, base64EncodeUint8Array(value)];
+			}
+			if (
+				Array.isArray(value) &&
+				value.length === 2 &&
+				typeof value[0] === "string" &&
+				value[0].startsWith("$")
+			) {
+				return [`$${value[0]}`, value[1]];
+			}
+			return value;
+		},
+		space,
+	);
 }
 
 /** Parses JSON with compat for values that BARE & CBOR supports. */

package/src/db/drizzle.test.ts

@@ -0,0 +1,157 @@
+import { describe, expect, test } from "vitest";
+import type {
+	DatabaseProviderContext,
+	SqliteBatchStatement,
+	SqliteBindings,
+	SqliteDatabase,
+	SqliteExecuteResult,
+	SqliteTransactionDatabase,
+} from "@/common/database/config";
+import { db } from "./drizzle";
+
+class FakeSqliteDatabase implements SqliteDatabase {
+	executeCalls: Array<{ sql: string; params?: SqliteBindings }> = [];
+	transactionTimeouts: Array<number | undefined> = [];
+
+	async exec(): Promise<void> {}
+
+	async execute(
+		sql: string,
+		params?: SqliteBindings,
+	): Promise<SqliteExecuteResult> {
+		this.executeCalls.push({ sql, params });
+		return emptyResult();
+	}
+
+	async executeBatch(
+		statements: SqliteBatchStatement[],
+	): Promise<SqliteExecuteResult[]> {
+		const transaction = await this.beginTransaction();
+		try {
+			const results: SqliteExecuteResult[] = [];
+			for (const statement of statements) {
+				results.push(
+					await transaction.execute(statement.sql, statement.params),
+				);
+			}
+			await transaction.commit();
+			return results;
+		} catch (error) {
+			await transaction.rollback();
+			throw error;
+		}
+	}
+
+	async beginTransaction(
+		timeoutMs?: number,
+	): Promise<SqliteTransactionDatabase> {
+		this.transactionTimeouts.push(timeoutMs);
+		this.executeCalls.push({ sql: "BEGIN" });
+		return {
+			exec: async () => {},
+			execute: async (sql, params) => {
+				this.executeCalls.push({ sql, params });
+				return emptyResult();
+			},
+			commit: async () => {
+				this.executeCalls.push({ sql: "COMMIT" });
+			},
+			rollback: async () => {
+				this.executeCalls.push({ sql: "ROLLBACK" });
+			},
+		};
+	}
+
+	async run(sql: string, params?: SqliteBindings): Promise<void> {
+		await this.execute(sql, params);
+	}
+
+	async query(sql: string, params?: SqliteBindings) {
+		const { columns, rows } = await this.execute(sql, params);
+		return { columns, rows };
+	}
+
+	async close(): Promise<void> {}
+}
+
+function emptyResult(): SqliteExecuteResult {
+	return {
+		columns: [],
+		rows: [],
+		changes: 0,
+		lastInsertRowId: null,
+	};
+}
+
+function testProviderContext(
+	database: SqliteDatabase,
+): DatabaseProviderContext {
+	return {
+		actorId: "actor-a",
+		kv: {
+			batchPut: async () => {},
+			batchGet: async (keys) => keys.map(() => null),
+			batchDelete: async () => {},
+			deleteRange: async () => {},
+		},
+		nativeDatabaseProvider: { open: async () => database },
+	};
+}
+
+describe("Drizzle database transactions", () => {
+	test("runs migrations in the shared transaction with a generous timeout", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		const provider = db({
+			onMigrate: async (client) => {
+				await client.execute(
+					"CREATE TABLE items(id INTEGER PRIMARY KEY)",
+				);
+			},
+		});
+		const client = await provider.createClient(
+			testProviderContext(nativeDb),
+		);
+		await provider.onMigrate(client);
+
+		expect(nativeDb.transactionTimeouts).toEqual([300_000]);
+		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",
+		]);
+	});
+
+	test("routes transaction work through the transaction handle", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		const client = await db().createClient(testProviderContext(nativeDb));
+		await client.transaction(
+			async (tx) => {
+				await tx.execute(
+					"INSERT INTO items(value) VALUES (?)",
+					"inside",
+				);
+			},
+			{ timeout: 120_000 },
+		);
+
+		expect(nativeDb.transactionTimeouts).toEqual([120_000]);
+		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
+			"BEGIN",
+			"INSERT INTO items(value) VALUES (?)",
+			"COMMIT",
+		]);
+	});
+
+	test("validates zero, negative, and non-finite 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");
+		}
+	});
+});