@firtoz/drizzle-sqlite-wasm 0.1.0

package/src/collections/websocket-collection.ts

@@ -0,0 +1,271 @@
+import type {
+	CollectionConfig,
+	SyncConfig,
+	InsertMutationFnParams,
+	UpdateMutationFnParams,
+	DeleteMutationFnParams,
+	UtilsRecord,
+	BaseCollectionConfig,
+	InferSchemaOutput,
+	InferSchemaInput,
+} from "@tanstack/db";
+import type { StandardSchemaV1 } from "@standard-schema/spec";
+
+interface WebSocketMessage<T> {
+	type: "insert" | "update" | "delete" | "sync" | "transaction" | "ack";
+	data?: T | T[];
+	mutations?: Array<{
+		type: "insert" | "update" | "delete";
+		data: T;
+		id?: string;
+	}>;
+	transactionId?: string;
+	id?: string;
+}
+
+interface WebSocketCollectionConfig<TSchema extends StandardSchemaV1>
+	extends Omit<
+		BaseCollectionConfig<InferSchemaOutput<TSchema>, string | number, TSchema>,
+		"onInsert" | "onUpdate" | "onDelete" | "sync" | "schema"
+	> {
+	url: string;
+	reconnectInterval?: number;
+	schema: TSchema;
+
+	// Note: onInsert/onUpdate/onDelete are handled by the WebSocket connection
+	// Users don't provide these handlers
+}
+
+interface WebSocketUtils extends UtilsRecord {
+	reconnect: () => void;
+	getConnectionState: () => "connected" | "disconnected" | "connecting";
+}
+
+export function webSocketCollectionOptions<TSchema extends StandardSchemaV1>(
+	config: WebSocketCollectionConfig<TSchema>,
+): CollectionConfig<InferSchemaOutput<TSchema>, string | number, TSchema> & {
+	utils: WebSocketUtils;
+	schema: TSchema;
+} {
+	type TItem = InferSchemaOutput<TSchema>;
+
+	let ws: WebSocket | null = null;
+	let reconnectTimer: NodeJS.Timeout | null = null;
+	let connectionState: "connected" | "disconnected" | "connecting" =
+		"disconnected";
+
+	// Track pending transactions awaiting acknowledgment
+	const pendingTransactions = new Map<
+		string,
+		{
+			resolve: () => void;
+			reject: (error: Error) => void;
+			timeout: NodeJS.Timeout;
+		}
+	>();
+
+	const handlers: {
+		onOpen?: () => void;
+		onMessage?: (event: MessageEvent) => void;
+		onError?: (error: Event) => void;
+		onClose?: () => void;
+	} = {};
+
+	const connect = () => {
+		connectionState = "connecting";
+		ws = new WebSocket(config.url);
+
+		ws.onopen = () => {
+			handlers.onOpen?.();
+		};
+		ws.onmessage = (event) => {
+			handlers.onMessage?.(event);
+		};
+		ws.onerror = (error) => {
+			handlers.onError?.(error);
+		};
+		ws.onclose = () => {
+			handlers.onClose?.();
+		};
+	};
+
+	const sync: SyncConfig<TItem>["sync"] = (params) => {
+		const { begin, write, commit, markReady } = params;
+
+		handlers.onOpen = () => {
+			if (!ws) return;
+
+			connectionState = "connected";
+			// Request initial sync
+			ws.send(JSON.stringify({ type: "sync" }));
+		};
+
+		handlers.onMessage = (event) => {
+			if (!ws) return;
+
+			const message: WebSocketMessage<TItem> = JSON.parse(event.data);
+
+			switch (message.type) {
+				case "sync":
+					// Initial sync with array of items
+					begin();
+					if (Array.isArray(message.data)) {
+						for (const item of message.data) {
+							write({ type: "insert", value: item });
+						}
+					}
+					commit();
+					markReady();
+					break;
+
+				case "insert":
+				case "update":
+				case "delete":
+					// Real-time updates from other clients
+					begin();
+					write({
+						type: message.type,
+						value: message.data as InferSchemaInput<TSchema>,
+					});
+					commit();
+					break;
+
+				case "ack":
+					// Server acknowledged our transaction
+					if (message.transactionId) {
+						const pending = pendingTransactions.get(message.transactionId);
+						if (pending) {
+							clearTimeout(pending.timeout);
+							pendingTransactions.delete(message.transactionId);
+							pending.resolve();
+						}
+					}
+					break;
+
+				case "transaction":
+					// Server sending back the actual data after processing our transaction
+					if (message.mutations) {
+						begin();
+						for (const mutation of message.mutations) {
+							write({
+								type: mutation.type,
+								value: mutation.data,
+							});
+						}
+						commit();
+					}
+					break;
+			}
+		};
+
+		handlers.onError = (error) => {
+			console.error("WebSocket error:", error);
+			connectionState = "disconnected";
+		};
+
+		handlers.onClose = () => {
+			connectionState = "disconnected";
+			// Auto-reconnect
+			if (!reconnectTimer) {
+				reconnectTimer = setTimeout(() => {
+					reconnectTimer = null;
+					connect();
+				}, config.reconnectInterval || 5000);
+			}
+		};
+
+		// Start connection
+		connect();
+
+		// Return cleanup function
+		return () => {
+			if (reconnectTimer) {
+				clearTimeout(reconnectTimer);
+				reconnectTimer = null;
+			}
+			if (ws) {
+				ws.close();
+				ws = null;
+			}
+		};
+	};
+
+	// Helper function to send transaction and wait for server acknowledgment
+	const sendTransaction = async (
+		params:
+			| InsertMutationFnParams<TItem>
+			| UpdateMutationFnParams<TItem>
+			| DeleteMutationFnParams<TItem>,
+	): Promise<void> => {
+		if (ws?.readyState !== WebSocket.OPEN) {
+			throw new Error("WebSocket not connected");
+		}
+
+		const transactionId = crypto.randomUUID();
+
+		// Convert all mutations in the transaction to the wire format
+		const mutations = params.transaction.mutations.map((mutation) => ({
+			type: mutation.type,
+			id: mutation.key,
+			data:
+				mutation.type === "delete"
+					? undefined
+					: mutation.type === "update"
+						? mutation.changes
+						: mutation.modified,
+		}));
+
+		// Send the entire transaction at once
+		ws.send(
+			JSON.stringify({
+				type: "transaction",
+				transactionId,
+				mutations,
+			}),
+		);
+
+		// Wait for server acknowledgment
+		return new Promise<void>((resolve, reject) => {
+			const timeout = setTimeout(() => {
+				pendingTransactions.delete(transactionId);
+				reject(new Error(`Transaction ${transactionId} timed out`));
+			}, 10000); // 10 second timeout
+
+			pendingTransactions.set(transactionId, {
+				resolve,
+				reject,
+				timeout,
+			});
+		});
+	};
+
+	// All mutation handlers use the same transaction sender
+	const onInsert = async (params: InsertMutationFnParams<TItem>) => {
+		await sendTransaction(params);
+	};
+
+	const onUpdate = async (params: UpdateMutationFnParams<TItem>) => {
+		await sendTransaction(params);
+	};
+
+	const onDelete = async (params: DeleteMutationFnParams<TItem>) => {
+		await sendTransaction(params);
+	};
+
+	return {
+		id: config.id,
+		schema: config.schema,
+		getKey: config.getKey,
+		sync: { sync },
+		onInsert,
+		onUpdate,
+		onDelete,
+		utils: {
+			reconnect: () => {
+				if (ws) ws.close();
+				connect();
+			},
+			getConnectionState: () => connectionState,
+		},
+	};
+}

package/src/context/useDrizzleSqlite.ts

@@ -0,0 +1,35 @@
+import { useContext } from "react";
+import type { DrizzleSqliteContextValue } from "./DrizzleSqliteProvider";
+import {
+	DrizzleSqliteContext,
+	useSqliteCollection,
+} from "./DrizzleSqliteProvider";
+import type { ValidTableNames } from "../collections/sqlite-collection";
+
+export type UseDrizzleSqliteReturn<TSchema extends Record<string, unknown>> = {
+	drizzle: DrizzleSqliteContextValue<TSchema>["drizzle"];
+	useCollection: <TTableName extends string & ValidTableNames<TSchema>>(
+		tableName: TTableName,
+	) => ReturnType<typeof useSqliteCollection<TSchema, TTableName>>;
+};
+
+export function useDrizzleSqlite<
+	TSchema extends Record<string, unknown>,
+>(): UseDrizzleSqliteReturn<TSchema> {
+	const context = useContext(
+		DrizzleSqliteContext,
+	) as DrizzleSqliteContextValue<TSchema> | null;
+
+	if (!context) {
+		throw new Error(
+			"useDrizzleSqlite must be used within a DrizzleSqliteProvider",
+		);
+	}
+
+	return {
+		drizzle: context.drizzle,
+		useCollection: <TTableName extends string & ValidTableNames<TSchema>>(
+			tableName: TTableName,
+		) => useSqliteCollection(context, tableName),
+	};
+}

package/src/drizzle/direct.ts

@@ -0,0 +1,27 @@
+import type { Database } from "@sqlite.org/sqlite-wasm";
+import type { DrizzleConfig } from "drizzle-orm";
+import { drizzle as drizzleSqliteProxy } from "drizzle-orm/sqlite-proxy";
+import { handleRemoteCallback } from "./handle-callback";
+
+export const drizzleSqliteWasm = <
+	TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+	sqliteDb: Database,
+	config: DrizzleConfig<TSchema> = {},
+	debug?: boolean,
+) => {
+	return drizzleSqliteProxy<TSchema>(async (sql, params, method) => {
+		const result = await handleRemoteCallback({
+			sqliteDb,
+			sql,
+			params,
+			method,
+			debug,
+		});
+		if (result.success) {
+			return result.result;
+		}
+		// If the callback failed, throw an error for drizzle to handle
+		throw new Error(result.error);
+	}, config);
+};

package/src/drizzle/handle-callback.ts

@@ -0,0 +1,113 @@
+import {
+	exhaustiveGuard,
+	fail,
+	success,
+	type MaybeError,
+} from "@firtoz/maybe-error";
+import type { Database } from "@sqlite.org/sqlite-wasm";
+
+export const handleRemoteCallback = async ({
+	sqliteDb,
+	sql,
+	params,
+	method,
+	debug: _debug = false,
+}: {
+	sqliteDb: Database;
+	sql: string;
+	// biome-ignore lint/suspicious/noExplicitAny: This is what drizzle-orm expects.
+	params: any[];
+	method: "run" | "all" | "values" | "get";
+	debug?: boolean;
+}): Promise<MaybeError<{ rows: unknown[] }, string>> => {
+	switch (method) {
+		case "run": {
+			// For INSERT, UPDATE, DELETE operations
+			try {
+				sqliteDb.exec({
+					sql,
+					bind: params,
+					callback: () => {},
+				});
+
+				return success({ rows: [] });
+			} catch (e: unknown) {
+				const errorMsg = e instanceof Error ? e.message : String(e);
+				console.error("Error executing run query:", errorMsg);
+				return fail(errorMsg);
+			}
+		}
+		case "get": {
+			// For getting a single row
+			const columnNames: string[] = [];
+			let rowData: unknown[] = [];
+			let callbackReceived = false;
+
+			try {
+				// Get column names and data in one go
+				sqliteDb.exec({
+					sql,
+					bind: params,
+					columnNames,
+					callback: (row) => {
+						callbackReceived = true;
+						if (Array.isArray(row)) {
+							// Store the first row's values
+							rowData = row;
+						} else {
+							// Convert object to array if needed
+							rowData = columnNames.map((col) => row[col]);
+						}
+					},
+				});
+			} catch (e: unknown) {
+				const errorMsg = e instanceof Error ? e.message : String(e);
+				console.error("Error getting row data:", errorMsg);
+				return fail(errorMsg);
+			}
+
+			if (!callbackReceived) {
+				const errorMsg = "No callback received for get method";
+				console.error(errorMsg);
+				return fail(errorMsg);
+			}
+
+			// For get method, return a single array of values
+			return success({ rows: rowData });
+		}
+
+		case "all":
+		case "values": {
+			// For getting multiple rows
+			const columnNames: string[] = [];
+			const rowsData: unknown[][] = [];
+
+			try {
+				// Get column names and data in one go
+				sqliteDb.exec({
+					sql,
+					bind: params,
+					columnNames,
+					callback: (row) => {
+						if (Array.isArray(row)) {
+							// Convert all values to strings
+							rowsData.push(row);
+						} else {
+							// Convert object to array if needed
+							rowsData.push(columnNames.map((col) => row[col]));
+						}
+					},
+				});
+			} catch (e: unknown) {
+				const errorMsg = e instanceof Error ? e.message : String(e);
+				console.error("Error getting all/values data:", errorMsg);
+				return fail(errorMsg);
+			}
+
+			// For all/values methods, return an array of arrays
+			return success({ rows: rowsData });
+		}
+	}
+
+	return exhaustiveGuard(method);
+};

package/src/drizzle/worker.ts

@@ -0,0 +1,24 @@
+import type { DrizzleConfig } from "drizzle-orm";
+import { drizzle as drizzleSqliteProxy } from "drizzle-orm/sqlite-proxy";
+import type { ISqliteWorkerClient } from "../worker/client";
+
+export const drizzleSqliteWasmWorker = <
+	TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+	client: ISqliteWorkerClient,
+	config: DrizzleConfig<TSchema> = {},
+) => {
+	return drizzleSqliteProxy<TSchema>(async (sql, params, method) => {
+		return new Promise<{ rows: unknown[] }>((resolve, reject) => {
+			client.performRemoteCallback(
+				{
+					sql,
+					params,
+					method,
+				},
+				resolve,
+				reject,
+			);
+		});
+	}, config);
+};

package/src/hooks/useDrizzleSqliteDb.ts

@@ -0,0 +1,139 @@
+import { useEffect, useMemo, useRef, useState } from "react";
+import {
+	customSqliteMigrate,
+	type DurableSqliteMigrationConfig,
+} from "@firtoz/drizzle-sqlite-wasm/sqlite-wasm-migrator";
+import { drizzleSqliteWasmWorker } from "@firtoz/drizzle-sqlite-wasm/drizzle-sqlite-wasm-worker";
+import type { ISqliteWorkerClient } from "../worker/manager";
+import {
+	initializeSqliteWorker,
+	isSqliteWorkerInitialized,
+} from "../worker/global-manager";
+
+export const useDrizzleSqliteDb = <TSchema extends Record<string, unknown>>(
+	WorkerConstructor: new () => Worker,
+	dbName: string,
+	schema: TSchema,
+	migrations: DurableSqliteMigrationConfig,
+	debug?: boolean,
+) => {
+	const resolveRef = useRef<null | (() => void)>(null);
+	const rejectRef = useRef<null | ((error: unknown) => void)>(null);
+	const [sqliteClient, setSqliteClient] = useState<ISqliteWorkerClient | null>(
+		null,
+	);
+	const sqliteClientRef = useRef<ISqliteWorkerClient | null>(null);
+
+	const readyPromise = useMemo(() => {
+		return new Promise<void>((resolve, reject) => {
+			resolveRef.current = resolve;
+			rejectRef.current = reject;
+		});
+	}, []);
+
+	// Initialize the global manager and get db instance
+	useEffect(() => {
+		if (typeof window === "undefined") {
+			// SSR stub
+			setSqliteClient({
+				performRemoteCallback: () => {},
+				checkpoint: () => Promise.resolve(),
+				onStarted: () => {},
+				terminate: () => {},
+			});
+			return;
+		}
+
+		let mounted = true;
+
+		const init = async () => {
+			// Initialize manager if not already initialized
+			if (!isSqliteWorkerInitialized()) {
+				await initializeSqliteWorker(WorkerConstructor);
+			}
+
+			// Get manager and create db instance
+			const { getSqliteWorkerManager } = await import(
+				"../worker/global-manager"
+			);
+			const manager = getSqliteWorkerManager();
+			const instance = await manager.getDbInstance(dbName);
+
+			if (mounted) {
+				sqliteClientRef.current = instance;
+				setSqliteClient(instance);
+			}
+		};
+
+		init();
+
+		return () => {
+			mounted = false;
+		};
+	}, [dbName, WorkerConstructor]);
+
+	// Create drizzle instance with a callback-based approach that waits for the client
+	const drizzle = useMemo(() => {
+		if (debug) {
+			console.log(`[DEBUG] ${dbName} - creating drizzle proxy wrapper`);
+		}
+		return drizzleSqliteWasmWorker<TSchema>(
+			{
+				performRemoteCallback: (data, resolve, reject) => {
+					const client = sqliteClientRef.current;
+					if (!client) {
+						console.error(
+							`[DEBUG] ${dbName} - performRemoteCallback called but no sqliteClient yet`,
+						);
+						reject(
+							new Error(
+								`Database ${dbName} not ready yet - still initializing`,
+							),
+						);
+						return;
+					}
+					client.performRemoteCallback(data, resolve, reject);
+				},
+				onStarted: (callback) => {
+					const client = sqliteClientRef.current;
+					if (!client) {
+						console.warn(
+							`[DEBUG] ${dbName} - onStarted called but no sqliteClient yet`,
+						);
+						return;
+					}
+					client.onStarted(callback);
+				},
+				terminate: () => {
+					sqliteClientRef.current?.terminate();
+				},
+			},
+			{ schema },
+		);
+	}, [schema, dbName]); // Using ref for sqliteClient to avoid recreating drizzle
+
+	useEffect(() => {
+		if (!sqliteClient) {
+			if (debug) {
+				console.log(`[DEBUG] ${dbName} - waiting for sqliteClient...`);
+			}
+			return;
+		}
+
+		sqliteClient.onStarted(async () => {
+			try {
+				await customSqliteMigrate(drizzle, migrations);
+				resolveRef.current?.();
+			} catch (error) {
+				console.error(`Migration error for ${dbName}:`, error);
+				rejectRef.current?.(error);
+			}
+		});
+
+		return () => {
+			sqliteClient.terminate();
+		};
+	}, [sqliteClient, drizzle, migrations, dbName]);
+
+	return { drizzle, readyPromise, sqliteClient };
+};

package/src/index.ts

@@ -0,0 +1,32 @@
+export { drizzleSqliteWasm } from "./drizzle/direct";
+export { sqliteCollectionOptions as drizzleCollectionOptions } from "./collections/sqlite-collection";
+export { syncableTable } from "@firtoz/drizzle-utils";
+export { makeId } from "@firtoz/drizzle-utils";
+export type {
+	IdOf,
+	TableId,
+	Branded,
+	SelectSchema,
+	InsertSchema,
+} from "@firtoz/drizzle-utils";
+export { useDrizzleSqliteDb } from "./hooks/useDrizzleSqliteDb";
+// SQLite WASM Provider
+export {
+	DrizzleSqliteProvider,
+	DrizzleSqliteContext,
+	useSqliteCollection,
+} from "./context/DrizzleSqliteProvider";
+export type { DrizzleSqliteContextValue } from "./context/DrizzleSqliteProvider";
+export { useDrizzleSqlite } from "./context/useDrizzleSqlite";
+export type { UseDrizzleSqliteReturn } from "./context/useDrizzleSqlite";
+
+export {
+	initializeSqliteWorker,
+	getSqliteWorkerManager,
+	isSqliteWorkerInitialized,
+	resetSqliteWorkerManager,
+} from "./worker/global-manager";
+export { SqliteWorkerManager, DbInstance } from "./worker/manager";
+export type { ISqliteWorkerClient } from "./worker/manager";
+export { customSqliteMigrate } from "./migration/migrator";
+export type { DurableSqliteMigrationConfig } from "./migration/migrator";