rivetkit 0.0.0-sqlite-uds.37f9e48 → 0.0.0-sqlite-profiling-compat.4a9a623

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

@@ -4,16 +4,13 @@ import type {
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteExecuteResult,
-	SqliteLeaseOptions,
 } from "./config";
 import { db } from "./mod";
 
 class FakeSqliteDatabase implements SqliteDatabase {
-	failSql = new Map<string, Error>();
 	executeCalls: {
 		sql: string;
 		params?: SqliteBindings;
-		lease?: SqliteLeaseOptions;
 	}[] = [];
 
 	async exec(): Promise<void> {}
@@ -21,17 +18,11 @@ class FakeSqliteDatabase implements SqliteDatabase {
 	async execute(
 		sql: string,
 		params?: SqliteBindings,
-		lease?: SqliteLeaseOptions,
 	): Promise<SqliteExecuteResult> {
 		this.executeCalls.push({
 			sql,
 			params,
-			lease,
 		});
-		const error = this.failSql.get(sql);
-		if (error) {
-			throw error;
-		}
 		return {
 			columns: [],
 			rows: [],
@@ -86,9 +77,7 @@ describe("db", () => {
 
 		await provider.onMigrate(client);
 
-		expect(
-			nativeDb.executeCalls.map(({ sql, params }) => ({ sql, params })),
-		).toEqual([
+		expect(nativeDb.executeCalls).toEqual([
 			{
 				sql: "SAVEPOINT __rivet_on_migrate",
 				params: undefined,
@@ -106,9 +95,6 @@ describe("db", () => {
 				params: undefined,
 			},
 		]);
-		const leases = nativeDb.executeCalls.map((call) => call.lease);
-		expect(leases.every((lease) => lease?.timeoutMs === 60_000)).toBe(true);
-		expect(new Set(leases.map((lease) => lease?.key)).size).toBe(1);
 	});
 
 	test("rolls back the migration savepoint when onMigrate fails", async () => {
@@ -129,9 +115,7 @@ describe("db", () => {
 			"migration failed",
 		);
 
-		expect(
-			nativeDb.executeCalls.map(({ sql, params }) => ({ sql, params })),
-		).toEqual([
+		expect(nativeDb.executeCalls).toEqual([
 			{
 				sql: "SAVEPOINT __rivet_on_migrate",
 				params: undefined,
@@ -150,82 +134,4 @@ describe("db", () => {
 			},
 		]);
 	});
-
-	test("runs transaction work under one hidden lease", async () => {
-		const nativeDb = new FakeSqliteDatabase();
-		const provider = db();
-		const client = await provider.createClient(
-			testProviderContext(nativeDb),
-		);
-
-		const value = await client.transaction(async (tx) => {
-			await tx.execute("INSERT INTO items(value) VALUES (?)", "inside");
-			return 42;
-		});
-
-		expect(value).toBe(42);
-		expect(nativeDb.executeCalls.map((call) => call.sql)).toEqual([
-			"BEGIN",
-			"INSERT INTO items(value) VALUES (?)",
-			"COMMIT",
-		]);
-		const keys = nativeDb.executeCalls.map((call) => call.lease?.key);
-		expect(keys.every((key) => key === keys[0] && key !== undefined)).toBe(
-			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

@@ -1,14 +1,21 @@
-import type { DatabaseProvider, RawAccess, SqliteDatabase } from "./config";
+import type {
+	DatabaseProvider,
+	RawAccess,
+	SqliteDatabase,
+	SqliteProfilingOptions,
+} from "./config";
 import { isSqliteBindingObject, toSqliteBindings } from "./shared";
 
 export type { RawAccess } from "./config";
 
-interface DatabaseFactoryConfig {
+export interface DatabaseFactoryConfig {
 	onMigrate?: (db: RawAccess) => Promise<void> | void;
+	/**
+	 * @experimental This entire configuration surface is subject to change.
+	 */
+	profiling?: SqliteProfilingOptions;
 }
 
-const MIGRATION_LEASE_TIMEOUT_MS = 60_000;
-
 function hasMultipleStatements(query: string): boolean {
 	const trimmed = query.trim().replace(/;+$/, "").trimEnd();
 	return trimmed.includes(";");
@@ -16,8 +23,10 @@ function hasMultipleStatements(query: string): boolean {
 
 export function db({
 	onMigrate,
+	profiling,
 }: DatabaseFactoryConfig = {}): DatabaseProvider<RawAccess> {
 	return {
+		sqliteProfiling: profiling,
 		createClient: async (ctx) => {
 			const nativeDatabaseProvider = ctx.nativeDatabaseProvider;
 			if (!nativeDatabaseProvider) {
@@ -28,7 +37,6 @@ export function db({
 
 			const db = await nativeDatabaseProvider.open(ctx.actorId);
 			let closed = false;
-			let transactionActive = false;
 			const ensureOpen = () => {
 				if (closed) {
 					throw new Error(
@@ -37,10 +45,7 @@ export function db({
 				}
 			};
 
-			const createClient = (lease?: {
-				key: string;
-				timeoutMs?: number;
-			}): RawAccess => ({
+			const client: RawAccess = {
 				execute: async <
 					TRow extends Record<string, unknown> = Record<
 						string,
@@ -51,11 +56,6 @@ 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;
@@ -71,7 +71,6 @@ export function db({
 							const { rows, columns } = await db.execute(
 								query,
 								bindings,
-								lease,
 							);
 							return rows.map((row) =>
 								rowToObject<TRow>(row, columns),
@@ -79,17 +78,13 @@ export function db({
 						}
 
 						if (!hasMultipleStatements(query)) {
-							const { rows, columns } = await db.execute(
-								query,
-								undefined,
-								lease,
-							);
+							const { rows, columns } = await db.execute(query);
 							return rows.map((row) =>
 								rowToObject<TRow>(row, columns),
 							);
 						}
 
-						return await execMultiStatement<TRow>(db, query, lease);
+						return await execMultiStatement<TRow>(db, query);
 					} finally {
 						const durationMs = performance.now() - start;
 						ctx.metrics?.trackSql(query, durationMs);
@@ -108,48 +103,6 @@ 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);
-						return result;
-					} catch (error) {
-						try {
-							await db.execute(
-								"ROLLBACK",
-								undefined,
-								transactionLease,
-							);
-						} catch {
-							// A deadline poisons the lease key, so its cleanup rollback is
-							// expected to fail. Preserve the callback or commit error.
-						}
-						throw error;
-					} finally {
-						transactionActive = false;
-					}
-				},
-				__withSqliteLease: async <T>(
-					timeoutMs: number,
-					callback: (leased: RawAccess) => Promise<T> | T,
-				): Promise<T> => {
-					const internalLease = {
-						key: crypto.randomUUID(),
-						timeoutMs,
-					};
-					return await callback(createClient(internalLease));
-				},
 				close: async () => {
 					if (!closed) {
 						closed = true;
@@ -157,15 +110,12 @@ export function db({
 					}
 				},
 				nativeMetrics: () => db.nativeMetrics?.() ?? null,
-			});
-			const client = createClient();
+			};
 			return client;
 		},
 		onMigrate: async (client) => {
 			if (onMigrate) {
-				await withMigrationSavepoint(client, (leased) =>
-					onMigrate(leased),
-				);
+				await withMigrationSavepoint(client, () => onMigrate(client));
 			}
 		},
 	};
@@ -185,47 +135,33 @@ function rowToObject<TRow extends Record<string, unknown>>(
 async function execMultiStatement<TRow extends Record<string, unknown>>(
 	db: SqliteDatabase,
 	query: string,
-	lease?: { key: string; timeoutMs?: number },
 ): Promise<TRow[]> {
 	const results: Record<string, unknown>[] = [];
 	let columnNames: string[] | null = null;
-	await db.exec(
-		query,
-		(row: unknown[], columns: string[]) => {
-			if (!columnNames) {
-				columnNames = columns;
-			}
-			results.push(rowToObject(row, columnNames));
-		},
-		lease,
-	);
+	await db.exec(query, (row: unknown[], columns: string[]) => {
+		if (!columnNames) {
+			columnNames = columns;
+		}
+		results.push(rowToObject(row, columnNames));
+	});
 	return results as TRow[];
 }
 
 async function withMigrationSavepoint<T>(
 	client: RawAccess,
-	callback: (leased: RawAccess) => Promise<T> | T,
+	callback: () => Promise<T> | T,
 ): Promise<T> {
-	return await client.__withSqliteLease(
-		MIGRATION_LEASE_TIMEOUT_MS,
-		async (leased) => {
-			await leased.execute("SAVEPOINT __rivet_on_migrate");
-			try {
-				const result = await callback(leased);
-				await leased.execute("RELEASE SAVEPOINT __rivet_on_migrate");
-				return result;
-			} catch (error) {
-				try {
-					await leased.execute(
-						"ROLLBACK TO SAVEPOINT __rivet_on_migrate",
-					);
-				} finally {
-					await leased.execute(
-						"RELEASE SAVEPOINT __rivet_on_migrate",
-					);
-				}
-				throw error;
-			}
-		},
-	);
+	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;
+	}
 }

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

@@ -3,7 +3,6 @@ import type {
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteExecuteResult,
-	SqliteLeaseOptions,
 	SqliteNativeMetrics,
 } from "./config";
 
@@ -67,16 +66,10 @@ interface NativeExecuteResult {
 }
 
 export interface JsNativeDatabaseLike {
-	exec(
-		sql: string,
-		leaseKey?: string,
-		leaseTimeoutMs?: number,
-	): Promise<NativeExecResult>;
+	exec(sql: string): Promise<NativeExecResult>;
 	execute(
 		sql: string,
 		params?: NativeBindParam[] | null,
-		leaseKey?: string,
-		leaseTimeoutMs?: number,
 	): Promise<NativeExecuteResult>;
 	query(
 		sql: string,
@@ -297,7 +290,6 @@ export function wrapJsNativeDatabase(
 	const executeNative = async (
 		sql: string,
 		params?: SqliteBindings,
-		lease?: SqliteLeaseOptions,
 	): Promise<SqliteExecuteResult> => {
 		const lastInsertRowIdColumn = lastInsertRowIdColumnName(sql);
 		if (lastInsertRowIdColumn) {
@@ -312,12 +304,7 @@ export function wrapJsNativeDatabase(
 		const release = gate.enter();
 		try {
 			const nativeParams = toNativeBindings(sql, params);
-			const result = await database.execute(
-				sql,
-				nativeParams,
-				lease?.key,
-				lease?.timeoutMs,
-			);
+			const result = await database.execute(sql, nativeParams);
 			if (result.lastInsertRowId !== undefined) {
 				lastInsertRowId = result.lastInsertRowId;
 			}
@@ -333,12 +320,11 @@ export function wrapJsNativeDatabase(
 		async exec(
 			sql: string,
 			callback?: (row: unknown[], columns: string[]) => void,
-			lease?: SqliteLeaseOptions,
 		): Promise<void> {
 			const release = gate.enter();
 			let result: NativeExecResult;
 			try {
-				result = await database.exec(sql, lease?.key, lease?.timeoutMs);
+				result = await database.exec(sql);
 			} catch (error) {
 				enrichNativeDatabaseError(database, error);
 			} finally {
@@ -354,9 +340,8 @@ export function wrapJsNativeDatabase(
 		async execute(
 			sql: string,
 			params?: SqliteBindings,
-			lease?: SqliteLeaseOptions,
 		): Promise<SqliteExecuteResult> {
-			return await executeNative(sql, params, lease);
+			return await executeNative(sql, params);
 		},
 		async run(sql: string, params?: SqliteBindings): Promise<void> {
 			await executeNative(sql, params);

package/src/common/log.ts

@@ -62,10 +62,15 @@ export function configureBaseLogger(logger: Logger): void {
 }
 
 /**
- * Create the default logger with optional log level.
+ * Configure the default logger with optional log level.
  */
-export function makeDefaultLogger(logLevel?: LogLevel): Logger {
-	return pino(
+export function configureDefaultLogger(logLevel?: LogLevel) {
+	// Store the configured log level
+	if (logLevel) {
+		configuredLogLevel = logLevel;
+	}
+
+	baseLogger = pino(
 		{
 			level: getPinoLevel(logLevel),
 			messageKey: "msg",
@@ -82,18 +87,7 @@ export function makeDefaultLogger(logLevel?: LogLevel): Logger {
 		},
 		createLogfmtDestination(),
 	);
-}
-
-/**
- * Configure the default logger with optional log level.
- */
-export function configureDefaultLogger(logLevel?: LogLevel) {
-	// Store the configured log level
-	if (logLevel) {
-		configuredLogLevel = logLevel;
-	}
 
-	baseLogger = makeDefaultLogger(logLevel);
 	loggerCache.clear();
 }