rivetkit 0.0.0-sqlite-uds.4e59a38 → 0.0.0-sqlite-uds.2015d80

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

@@ -9,6 +9,7 @@ import type {
 import { db } from "./mod";
 
 class FakeSqliteDatabase implements SqliteDatabase {
+	failSql = new Map<string, Error>();
 	executeCalls: {
 		sql: string;
 		params?: SqliteBindings;
@@ -27,6 +28,10 @@ class FakeSqliteDatabase implements SqliteDatabase {
 			params,
 			lease,
 		});
+		const error = this.failSql.get(sql);
+		if (error) {
+			throw error;
+		}
 		return {
 			columns: [],
 			rows: [],
@@ -169,4 +174,53 @@ describe("db", () => {
 			true,
 		);
 	});
+
+	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((call) => call.sql)).toEqual([
+			"BEGIN",
+			"INSERT INTO items(value) VALUES (?)",
+			"ROLLBACK",
+		]);
+	});
+
+	test("rejects use of the outer client inside a transaction callback", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		const client = await db().createClient(testProviderContext(nativeDb));
+
+		await expect(
+			client.transaction(async () => {
+				await client.execute("INSERT INTO items(value) VALUES ('outside')");
+			}),
+		).rejects.toThrow("Use the tx argument");
+
+		expect(nativeDb.executeCalls.map((call) => call.sql)).toEqual([
+			"BEGIN",
+			"ROLLBACK",
+		]);
+	});
+
+	test("surfaces a lease-expired commit error", async () => {
+		const nativeDb = new FakeSqliteDatabase();
+		nativeDb.failSql.set("COMMIT", new Error("LeaseExpired"));
+		const client = await db().createClient(testProviderContext(nativeDb));
+
+		await expect(client.transaction(async () => 42)).rejects.toThrow(
+			"LeaseExpired",
+		);
+		expect(nativeDb.executeCalls.map((call) => call.sql)).toEqual([
+			"BEGIN",
+			"COMMIT",
+			"ROLLBACK",
+		]);
+	});
 });

package/src/common/database/mod.ts

@@ -28,6 +28,7 @@ export function db({
 
 			const db = await nativeDatabaseProvider.open(ctx.actorId);
 			let closed = false;
+			let transactionActive = false;
 			const ensureOpen = () => {
 				if (closed) {
 					throw new Error(
@@ -50,6 +51,11 @@ export function db({
 					...args: unknown[]
 				): Promise<TRow[]> => {
 					ensureOpen();
+					if (!lease && transactionActive) {
+						throw new Error(
+							"Use the tx argument inside db.transaction(); the outer database client cannot run work from the transaction callback.",
+						);
+					}
 
 					const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
 					const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
@@ -105,9 +111,15 @@ export function db({
 				transaction: async <T>(
 					callback: (tx: RawAccess) => Promise<T> | T,
 				): Promise<T> => {
+					if (lease || transactionActive) {
+						throw new Error(
+							"Nested or concurrent db.transaction() calls are not supported.",
+						);
+					}
 					const transactionLease = { key: crypto.randomUUID() };
 					const tx = createClient(transactionLease);
 					await db.execute("BEGIN", undefined, transactionLease);
+					transactionActive = true;
 					try {
 						const result = await callback(tx);
 						await db.execute("COMMIT", undefined, transactionLease);
@@ -124,6 +136,8 @@ export function db({
 							// expected to fail. Preserve the callback or commit error.
 						}
 						throw error;
+					} finally {
+						transactionActive = false;
 					}
 				},
 				__withSqliteLease: async <T>(

package/src/db/drizzle.ts

@@ -32,7 +32,13 @@ type DrizzleDatabase<TSchema extends DrizzleSchema> = Omit<
 	SqliteRemoteDatabase<TSchema>,
 	"transaction"
 > &
-	RawAccess;
+	Omit<RawAccess, "transaction"> & {
+		transaction: <T>(
+			callback: (
+				tx: DrizzleDatabase<TSchema>,
+			) => Promise<T> | T,
+		) => Promise<T>;
+	};
 
 const MIGRATION_LEASE_TIMEOUT_MS = 60_000;
 
@@ -95,6 +101,7 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 
 			const nativeDb = await nativeDatabaseProvider.open(ctx.actorId);
 			let closed = false;
+			let transactionActive = false;
 			const ensureOpen = () => {
 				if (closed) {
 					throw new Error(
@@ -113,6 +120,11 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 					method: "run" | "all" | "values" | "get",
 				) => {
 					ensureOpen();
+					if (!lease && transactionActive) {
+						throw new Error(
+							"Use the tx argument inside db.transaction(); the outer database client cannot run work from the transaction callback.",
+						);
+					}
 
 					const start = performance.now();
 					const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
@@ -177,8 +189,15 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 					);
 				};
 				drizzleDb.transaction = async <T>(
-					transactionCallback: (tx: RawAccess) => Promise<T> | T,
+					transactionCallback: (
+						tx: DrizzleDatabase<TSchema>,
+					) => Promise<T> | T,
 				): Promise<T> => {
+					if (lease || transactionActive) {
+						throw new Error(
+							"Nested or concurrent db.transaction() calls are not supported.",
+						);
+					}
 					const transactionLease = { key: crypto.randomUUID() };
 					const tx = createDrizzleClient(transactionLease);
 					await nativeDb.execute(
@@ -186,6 +205,7 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 						undefined,
 						transactionLease,
 					);
+					transactionActive = true;
 					try {
 						const result = await transactionCallback(tx);
 						await nativeDb.execute(
@@ -205,6 +225,8 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 							// Preserve the callback or commit error after lease poisoning.
 						}
 						throw error;
+					} finally {
+						transactionActive = false;
 					}
 				};
 				drizzleDb.__withSqliteLease = async <T>(