rivetkit 0.0.0-sqlite-uds.2015d80 → 0.0.0-sqlite-profiling-compat.87d61c0

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();
 }
 

package/src/db/drizzle.ts

@@ -28,19 +28,8 @@ export {
 } from "drizzle-orm/sqlite-core";
 
 type DrizzleSchema = Record<string, unknown>;
-type DrizzleDatabase<TSchema extends DrizzleSchema> = Omit<
-	SqliteRemoteDatabase<TSchema>,
-	"transaction"
-> &
-	Omit<RawAccess, "transaction"> & {
-		transaction: <T>(
-			callback: (
-				tx: DrizzleDatabase<TSchema>,
-			) => Promise<T> | T,
-		) => Promise<T>;
-	};
-
-const MIGRATION_LEASE_TIMEOUT_MS = 60_000;
+type DrizzleDatabase<TSchema extends DrizzleSchema> =
+	SqliteRemoteDatabase<TSchema> & RawAccess;
 
 interface DrizzleMigrationJournalEntry {
 	idx: number;
@@ -101,7 +90,6 @@ 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(
@@ -110,192 +98,107 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 				}
 			};
 
-			const createDrizzleClient = (lease?: {
-				key: string;
-				timeoutMs?: number;
-			}): DrizzleDatabase<TSchema> => {
-				const runSql = async (
-					query: string,
-					params: unknown[],
-					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;
-					const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
-					try {
-						const { rows } = await nativeDb.execute(
-							query,
-							toSqliteBindings(params),
-							lease,
-						);
-						if (method === "run") {
-							return { rows: [] };
-						}
-						if (method === "get") {
-							return { rows: rows[0] };
-						}
-						return { rows };
-					} finally {
-						const durationMs = performance.now() - start;
-						ctx.metrics?.trackSql(query, durationMs);
-						if (ctx.metrics) {
-							ctx.log?.debug({
-								msg: "sql query",
-								query: query.slice(0, 120),
-								durationMs,
-								kvReads:
-									ctx.metrics.totalKvReads - kvReadsBefore,
-								kvWrites:
-									ctx.metrics.totalKvWrites - kvWritesBefore,
-							});
-						}
-					}
-				};
+			const runSql = async (
+				query: string,
+				params: unknown[],
+				method: "run" | "all" | "values" | "get",
+			) => {
+				ensureOpen();
 
-				const callback: RemoteCallback = async (
-					query,
-					params,
-					method,
-				) => {
-					return await runSql(query, params, method);
-				};
-
-				const drizzleDb = drizzle(callback, {
-					schema,
-				}) as unknown as DrizzleDatabase<TSchema>;
-				drizzleDb.execute = async <
-					TRow extends Record<string, unknown> = Record<
-						string,
-						unknown
-					>,
-				>(
-					query: string,
-					...args: unknown[]
-				): Promise<TRow[]> => {
-					return await executeRaw<TRow>(
-						nativeDb,
-						ctx,
-						ensureOpen,
+				const start = performance.now();
+				const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
+				const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
+				try {
+					const { rows } = await nativeDb.execute(
 						query,
-						args,
-						lease,
+						toSqliteBindings(params),
 					);
-				};
-				drizzleDb.transaction = async <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.",
-						);
+					if (method === "run") {
+						return { rows: [] };
 					}
-					const transactionLease = { key: crypto.randomUUID() };
-					const tx = createDrizzleClient(transactionLease);
-					await nativeDb.execute(
-						"BEGIN",
-						undefined,
-						transactionLease,
-					);
-					transactionActive = true;
-					try {
-						const result = await transactionCallback(tx);
-						await nativeDb.execute(
-							"COMMIT",
-							undefined,
-							transactionLease,
-						);
-						return result;
-					} catch (error) {
-						try {
-							await nativeDb.execute(
-								"ROLLBACK",
-								undefined,
-								transactionLease,
-							);
-						} catch {
-							// Preserve the callback or commit error after lease poisoning.
-						}
-						throw error;
-					} finally {
-						transactionActive = false;
+					if (method === "get") {
+						return { rows: rows[0] };
 					}
-				};
-				drizzleDb.__withSqliteLease = async <T>(
-					timeoutMs: number,
-					leaseCallback: (leased: RawAccess) => Promise<T> | T,
-				): Promise<T> => {
-					const internalLease = {
-						key: crypto.randomUUID(),
-						timeoutMs,
-					};
-					return await leaseCallback(
-						createDrizzleClient(internalLease),
-					);
-				};
-				drizzleDb.close = async () => {
-					if (!closed) {
-						closed = true;
-						await nativeDb.close();
+					return { rows };
+				} finally {
+					const durationMs = performance.now() - start;
+					ctx.metrics?.trackSql(query, durationMs);
+					if (ctx.metrics) {
+						ctx.log?.debug({
+							msg: "sql query",
+							query: query.slice(0, 120),
+							durationMs,
+							kvReads: ctx.metrics.totalKvReads - kvReadsBefore,
+							kvWrites:
+								ctx.metrics.totalKvWrites - kvWritesBefore,
+						});
 					}
-				};
+				}
+			};
 
-				return drizzleDb;
+			const callback: RemoteCallback = async (query, params, method) => {
+				return await runSql(query, params, method);
 			};
 
-			return createDrizzleClient();
+			const drizzleDb = drizzle(callback, {
+				schema,
+			}) as DrizzleDatabase<TSchema>;
+			drizzleDb.execute = async <
+				TRow extends Record<string, unknown> = Record<string, unknown>,
+			>(
+				query: string,
+				...args: unknown[]
+			): Promise<TRow[]> => {
+				return await executeRaw<TRow>(
+					nativeDb,
+					ctx,
+					ensureOpen,
+					query,
+					args,
+				);
+			};
+			drizzleDb.close = async () => {
+				if (!closed) {
+					closed = true;
+					await nativeDb.close();
+				}
+			};
+
+			return drizzleDb;
 		},
 		onMigrate: async (client) => {
 			if (!migrations && !onMigrate) {
 				return;
 			}
-			await withMigrationSavepoint(client, async (leased) => {
+			await withMigrationSavepoint(client, async () => {
 				if (migrations) {
-					await runMigrations(leased, migrations);
+					await runMigrations(client, migrations);
 				}
 				if (onMigrate) {
-					await onMigrate(leased);
+					await onMigrate(client);
 				}
 			});
 		},
 	};
 }
 
-async function withMigrationSavepoint<TSchema extends DrizzleSchema, T>(
-	client: DrizzleDatabase<TSchema>,
-	callback: (leased: DrizzleDatabase<TSchema>) => Promise<T> | T,
+async function withMigrationSavepoint<T>(
+	client: RawAccess,
+	callback: () => Promise<T> | T,
 ): Promise<T> {
-	return await client.__withSqliteLease(
-		MIGRATION_LEASE_TIMEOUT_MS,
-		async (leasedAccess) => {
-			const leased = leasedAccess as DrizzleDatabase<TSchema>;
-			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;
+	}
 }
 
 async function runMigrations<TSchema extends DrizzleSchema>(
@@ -383,7 +286,6 @@ async function executeRaw<TRow extends Record<string, unknown>>(
 	ensureOpen: () => void,
 	query: string,
 	args: unknown[],
-	lease?: { key: string; timeoutMs?: number },
 ): Promise<TRow[]> {
 	ensureOpen();
 
@@ -395,28 +297,23 @@ async function executeRaw<TRow extends Record<string, unknown>>(
 			const { rows, columns } = await db.execute(
 				query,
 				toSqliteBindings(args),
-				lease,
 			);
 			return rows.map((row) => rowToObject<TRow>(row, columns));
 		}
 
 		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));
 		}
 
 		const results: Record<string, unknown>[] = [];
 		let columnNames: string[] | null = null;
-		await db.exec(
-			query,
-			(row, columns) => {
-				if (!columnNames) {
-					columnNames = columns;
-				}
-				results.push(rowToObject(row, columnNames));
-			},
-			lease,
-		);
+		await db.exec(query, (row, columns) => {
+			if (!columnNames) {
+				columnNames = columns;
+			}
+			results.push(rowToObject(row, columnNames));
+		});
 		return results as TRow[];
 	} finally {
 		const durationMs = performance.now() - start;

package/src/db/mod.ts

@@ -8,6 +8,7 @@ export type {
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteNativeMetrics,
+	SqliteProfilingOptions,
 	SqliteQueryResult,
 } from "@/common/database/config";
 export { db } from "@/common/database/mod";