@lobomfz/db 0.4.0 → 0.5.0

package/src/database.ts

@@ -4,387 +4,300 @@ import type { Type } from "arktype";
 import { Kysely } from "kysely";
 
 import { BunSqliteDialect } from "./dialect/dialect.js";
-import type { DbFieldMeta } from "./env.js";
-import type { GeneratedPreset } from "./generated.js";
-import { Differ, type DesiredTable } from "./migration/diff.js";
+import { resolveFtsPlan, type FtsTableInfo, type ResolvedFtsPlan } from "./fts/resolve.js";
+import { generateRebuildSelect } from "./fts/sql.js";
+import { DataLossError } from "./migration/data-loss-error.js";
+import { Differ } from "./migration/diff.js";
 import { Executor } from "./migration/execute.js";
+import { GuidedMigrationRunner } from "./migration/guided.js";
 import { Introspector } from "./migration/introspect.js";
+import { MigrationRegistry } from "./migration/registry.js";
+import type { DesiredSchema, DesiredTable, ExistingSchema } from "./migration/types.js";
 import { ResultHydrationPlugin } from "./plugin.js";
+import { SchemaCompiler } from "./schema-compiler.js";
 import type {
+	DataMigration,
 	DatabaseOptions,
-	IndexDefinition,
+	DatabasePragmas,
+	FtsApi,
+	FtsConfig,
 	SchemaRecord,
 	TablesFromSchemas,
-	DatabasePragmas,
+	VectorIndexApi,
+	VectorIndexConfig,
 } from "./types.js";
+import { loadExtensions } from "./vector-index/load.js";
+import { registerVectorIndex, unregisterVectorIndex } from "./vector-index/nearest.js";
+import {
+	resolveVectorIndex,
+	type ResolvedVectorIndex,
+	type VectorAnchorInfo,
+} from "./vector-index/resolve.js";
+import { generateRebuildVectorSql } from "./vector-index/sql.js";
 import { WriteValidationPlugin } from "./write-validation-plugin.js";
 
-type ArkBranch = {
-	domain?: string;
-	proto?: unknown;
-	unit?: unknown;
-	structure?: unknown;
-	inner?: { divisor?: unknown };
-	meta?: DbFieldMeta & { _generated?: GeneratedPreset };
+const defaultPragmas: DatabasePragmas = {
+	foreign_keys: true,
 };
 
-type StructureProp = {
-	key: string;
-	required: boolean;
-	value: Type & {
-		branches: ArkBranch[];
-		proto?: unknown;
-		meta: DbFieldMeta & { _generated?: GeneratedPreset };
-	};
-	inner: { default?: unknown };
-};
+export class Database<
+	T extends SchemaRecord,
+	F extends FtsConfig<T> = FtsConfig<T>,
+	V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+> {
+	private sqlite: BunDatabase;
+	private compiler = new SchemaCompiler();
+	private ftsPlans = new Map<string, ResolvedFtsPlan>();
+	private vectorIndexPlans = new Map<string, ResolvedVectorIndex>();
 
-type Prop = {
-	key: string;
-	kind: "required" | "optional";
-	domain?: string;
-	nullable?: boolean;
-	isBoolean?: boolean;
-	isInteger?: boolean;
-	isDate?: boolean;
-	isBlob?: boolean;
-	isJson?: boolean;
-	jsonSchema?: Type;
-	meta?: DbFieldMeta;
-	generated?: GeneratedPreset;
-	defaultValue?: unknown;
-};
+	readonly infer: TablesFromSchemas<T, F, V> = undefined as any;
 
-const typeMap: Record<string, string> = {
-	string: "TEXT",
-	number: "REAL",
-};
+	readonly kysely: Kysely<TablesFromSchemas<T, F, V>>;
 
-const defaultPragmas: DatabasePragmas = {
-	foreign_keys: true,
-};
+	readonly fts: FtsApi<F>;
 
-export class Database<T extends SchemaRecord> {
-	private sqlite: BunDatabase;
+	readonly vectorIndex: VectorIndexApi<V>;
+
+	static async open<
+		T extends SchemaRecord,
+		F extends FtsConfig<T> = FtsConfig<T>,
+		V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+	>(options: DatabaseOptions<T, F, V>) {
+		const db = new Database(options);
 
-	readonly infer: TablesFromSchemas<T> = undefined as any;
+		await db.runMigrations();
 
-	readonly kysely: Kysely<TablesFromSchemas<T>>;
+		return db;
+	}
 
-	constructor(private options: DatabaseOptions<T>) {
+	private constructor(private options: DatabaseOptions<T, F, V>) {
 		const tableSchemas = new Map<string, Type>(Object.entries(options.schema.tables));
 		const writeSchemas = new Map<string, Type>(
 			Object.entries(options.schema.tables).map(([table, schema]) => [
 				table,
-				this.createWriteSchema(schema),
+				this.compiler.createWriteSchema(schema),
 			]),
 		);
 
 		this.sqlite = new BunDatabase(options.path);
 
 		this.applyPragmas();
-
-		this.migrate();
+		this.assertVectorIndexesHaveExtension();
+		this.loadConfiguredExtensions();
 
 		const validation = {
 			onRead: options.validation?.onRead ?? false,
 		};
 
-		this.kysely = new Kysely<TablesFromSchemas<T>>({
+		this.kysely = new Kysely<TablesFromSchemas<T, F, V>>({
 			dialect: new BunSqliteDialect({ database: this.sqlite }),
 			plugins: [
 				new WriteValidationPlugin(writeSchemas),
 				new ResultHydrationPlugin(tableSchemas, validation),
 			],
 		});
-	}
-
-	private applyPragmas() {
-		const pragmas = { ...defaultPragmas, ...this.options.pragmas };
-
-		if (pragmas.journal_mode) {
-			this.sqlite.run(`PRAGMA journal_mode = ${pragmas.journal_mode.toUpperCase()}`);
-		}
-
-		if (pragmas.synchronous) {
-			this.sqlite.run(`PRAGMA synchronous = ${pragmas.synchronous.toUpperCase()}`);
-		}
 
-		this.sqlite.run(`PRAGMA foreign_keys = ${pragmas.foreign_keys ? "ON" : "OFF"}`);
+		this.fts = {
+			rebuild: (name: string) => this.rebuildFts(name),
+		} as FtsApi<F>;
 
-		if (pragmas.busy_timeout_ms !== undefined) {
-			this.sqlite.run(`PRAGMA busy_timeout = ${pragmas.busy_timeout_ms}`);
-		}
+		this.vectorIndex = {
+			rebuild: (name: string) => this.rebuildVectorIndex(name),
+		} as VectorIndexApi<V>;
 	}
 
-	private normalizeProp(structureProp: StructureProp, parentSchema: Type) {
-		const { key, value: v, inner } = structureProp;
-		const kind: Prop["kind"] = structureProp.required ? "required" : "optional";
-		const defaultValue = inner.default;
+	private loadConfiguredExtensions() {
+		const extensions = this.options.extensions ?? [];
 
-		const concrete = v.branches.filter((b) => b.unit !== null && b.domain !== "undefined");
-		const nullable = concrete.length < v.branches.length;
+		loadExtensions(this.sqlite, extensions);
 
-		const branchMeta = v.branches.find((b) => b.meta && Object.keys(b.meta).length > 0)?.meta;
-		const meta = { ...branchMeta, ...v.meta };
-		const generated = meta._generated;
-
-		if (v.proto === Date || concrete.some((b) => b.proto === Date)) {
-			return { key, kind, nullable, isDate: true, generated, defaultValue };
-		}
+		const hasSqliteVec = extensions.some((e) => e.kind === "sqlite-vec");
 
-		if (v.proto === Uint8Array || concrete.some((b) => b.proto === Uint8Array)) {
-			return {
-				key,
-				kind,
-				nullable,
-				isBlob: true,
-				meta,
-				generated,
-				defaultValue,
-			};
+		if (hasSqliteVec) {
+			return;
 		}
 
-		if (concrete.length > 0 && concrete.every((b) => b.domain === "boolean")) {
-			return { key, kind, nullable, isBoolean: true, generated, defaultValue };
-		}
+		const probe = new Introspector(this.sqlite);
 
-		if (concrete.some((b) => !!b.structure)) {
-			return {
-				key,
-				kind,
-				nullable,
-				isJson: true,
-				jsonSchema: (parentSchema as any).get(key) as Type,
-				meta,
-				generated,
-				defaultValue,
-			};
+		if (probe.hasVec0Tables()) {
+			loadExtensions(this.sqlite, [{ kind: "sqlite-vec" }]);
 		}
-
-		const branch = concrete[0];
-
-		return {
-			key,
-			kind,
-			nullable,
-			domain: branch?.domain,
-			isInteger: !!branch?.inner?.divisor,
-			meta,
-			generated,
-			defaultValue,
-		};
 	}
 
-	private sqlType(prop: Prop) {
-		if (prop.isJson) {
-			return "TEXT";
-		}
+	private assertVectorIndexesHaveExtension() {
+		const vectorIndexes = this.options.schema.vectorIndexes;
 
-		if (prop.isBlob) {
-			return "BLOB";
+		if (!vectorIndexes || Object.keys(vectorIndexes).length === 0) {
+			return;
 		}
 
-		if (prop.isDate || prop.isBoolean || prop.isInteger) {
-			return "INTEGER";
-		}
+		const hasSqliteVec = (this.options.extensions ?? []).some((e) => e.kind === "sqlite-vec");
 
-		if (prop.domain) {
-			return typeMap[prop.domain] ?? "TEXT";
+		if (hasSqliteVec) {
+			return;
 		}
 
-		return "TEXT";
-	}
+		const [first] = Object.keys(vectorIndexes);
 
-	private columnConstraint(prop: Prop) {
-		if (prop.generated === "autoincrement") {
-			return "PRIMARY KEY AUTOINCREMENT";
-		}
-
-		if (prop.meta?.primaryKey) {
-			return "PRIMARY KEY";
-		}
-
-		if (prop.kind === "required" && !prop.nullable) {
-			return "NOT NULL";
-		}
-
-		return null;
+		throw new Error(
+			`vectorIndex "${first}" requires the sqlite-vec extension. Add { kind: "sqlite-vec" } to DatabaseOptions.extensions.`,
+		);
 	}
 
-	private defaultClause(prop: Prop) {
-		if (prop.generated === "now") {
-			return "DEFAULT (unixepoch())";
-		}
+	private applyPragmas() {
+		const pragmas = { ...defaultPragmas, ...this.options.pragmas };
 
-		if (prop.defaultValue === undefined || prop.generated === "autoincrement") {
-			return null;
+		if (pragmas.journal_mode) {
+			this.sqlite.run(`PRAGMA journal_mode = ${pragmas.journal_mode.toUpperCase()}`);
 		}
 
-		if (prop.defaultValue === null) {
-			return "DEFAULT NULL";
+		if (pragmas.synchronous) {
+			this.sqlite.run(`PRAGMA synchronous = ${pragmas.synchronous.toUpperCase()}`);
 		}
 
-		if (typeof prop.defaultValue === "string") {
-			return `DEFAULT '${prop.defaultValue}'`;
-		}
+		this.sqlite.run(`PRAGMA foreign_keys = ${pragmas.foreign_keys ? "ON" : "OFF"}`);
 
-		if (typeof prop.defaultValue === "number" || typeof prop.defaultValue === "boolean") {
-			return `DEFAULT ${prop.defaultValue}`;
+		if (pragmas.busy_timeout_ms !== undefined) {
+			this.sqlite.run(`PRAGMA busy_timeout = ${pragmas.busy_timeout_ms}`);
 		}
-
-		throw new Error(
-			`Unsupported default value type: ${typeof prop.defaultValue} ${JSON.stringify(prop)}`,
-		);
 	}
 
-	private columnDef(prop: Prop) {
-		return [
-			`"${prop.key}"`,
-			this.sqlType(prop),
-			this.columnConstraint(prop),
-			prop.meta?.unique ? "UNIQUE" : null,
-			this.defaultClause(prop),
-		]
-			.filter(Boolean)
-			.join(" ");
-	}
+	private buildDesiredSchema(): DesiredSchema {
+		const desiredTables: DesiredTable[] = [];
+		const tableInfos = new Map<string, FtsTableInfo>();
+		const vectorAnchors = new Map<string, VectorAnchorInfo>();
+		const schemaIndexes = this.options.schema.indexes;
 
-	private foreignKey(prop: Prop) {
-		const ref = prop.meta?.references;
+		for (const [name, schema] of Object.entries(this.options.schema.tables)) {
+			const compiled = this.compiler.compileTable(name, schema, schemaIndexes?.[name] ?? []);
 
-		if (!ref) {
-			return null;
+			desiredTables.push(compiled.table);
+			tableInfos.set(name, compiled.tableInfo);
+			vectorAnchors.set(name, compiled.vectorAnchor);
 		}
 
-		const [table, column] = ref.split(".");
+		const desiredFts: ResolvedFtsPlan[] = [];
 
-		let fk = `FOREIGN KEY ("${prop.key}") REFERENCES "${table}"("${column}")`;
+		if (this.options.schema.fts) {
+			for (const [name, entry] of Object.entries(this.options.schema.fts)) {
+				const plan = resolveFtsPlan(name, entry, tableInfos);
+				desiredFts.push(plan);
+				this.ftsPlans.set(name, plan);
+			}
+		}
 
-		const onDelete = prop.meta?.onDelete;
+		const desiredVectorIndexes: ResolvedVectorIndex[] = [];
 
-		if (onDelete) {
-			fk += ` ON DELETE ${onDelete.toUpperCase()}`;
+		if (this.options.schema.vectorIndexes) {
+			for (const [name, entry] of Object.entries(this.options.schema.vectorIndexes)) {
+				const plan = resolveVectorIndex(name, entry as any, vectorAnchors);
+				desiredVectorIndexes.push(plan);
+				this.vectorIndexPlans.set(name, plan);
+				registerVectorIndex(plan);
+			}
 		}
 
-		return fk;
+		return { desiredTables, desiredFts, desiredVectorIndexes };
 	}
 
-	private addColumnDef(prop: Prop) {
-		let def = this.columnDef(prop);
+	private async runMigrations() {
+		const desired = this.buildDesiredSchema();
+
+		const introspector = new Introspector(this.sqlite);
+		const existingVectorIndexes = introspector.introspectVectorIndexes();
+		const existingState: ExistingSchema = {
+			existing: introspector.introspect(),
+			existingFts: introspector.introspectFts(),
+			existingVectorIndexes,
+		};
 
-		const ref = prop.meta?.references;
+		const migration = this.options.migration;
 
-		if (ref) {
-			const [table, column] = ref.split(".");
-			def += ` REFERENCES "${table}"("${column}")`;
+		if (migration) {
+			await this.applyWithMigration(desired, existingState, migration);
+		} else {
+			this.applyDiff(desired, existingState);
+		}
 
-			if (prop.meta?.onDelete) {
-				def += ` ON DELETE ${prop.meta.onDelete.toUpperCase()}`;
+		for (const name of existingVectorIndexes.keys()) {
+			if (!this.vectorIndexPlans.has(name)) {
+				unregisterVectorIndex(name);
 			}
 		}
-
-		return def;
 	}
 
-	private parseSchemaProps(schema: Type) {
-		const structureProps = (schema as any).structure?.props as StructureProp[] | undefined;
+	private applyDiff(desired: DesiredSchema, existingState: ExistingSchema) {
+		const differ = new Differ(desired, existingState);
 
-		if (!structureProps) {
-			return [];
+		const ops = differ.diff();
+
+		if (this.options.failOnDataLoss && differ.destructive.length > 0) {
+			throw new DataLossError(differ.destructive);
 		}
 
-		return structureProps.map((p) => this.normalizeProp(p, schema));
+		new Executor(this.sqlite, ops).execute();
 	}
 
-	private createWriteSchema(schema: Type) {
-		const autoIncrementColumns = this.parseSchemaProps(schema)
-			.filter((prop) => prop.generated === "autoincrement")
-			.map((prop) => prop.key);
+	private async applyWithMigration(
+		desired: DesiredSchema,
+		existingState: ExistingSchema,
+		migration: DataMigration<T, F, V>,
+	) {
+		const registry = new MigrationRegistry(this.sqlite);
+		registry.ensureTable();
 
-		if (autoIncrementColumns.length === 0) {
-			return schema;
-		}
-
-		return (schema as any).omit(...autoIncrementColumns) as Type;
-	}
+		const checksum = MigrationRegistry.checksum(migration.run);
 
-	private generateCreateTableSQL(tableName: string, props: Prop[]) {
-		const columns: string[] = [];
-		const fks: string[] = [];
+		registry.assertNoDrift(migration.id, checksum);
 
-		for (const prop of props) {
-			columns.push(this.columnDef(prop));
+		if (registry.appliedChecksum(migration.id) !== null) {
+			this.applyDiff(desired, existingState);
+			return;
+		}
 
-			const fk = this.foreignKey(prop);
+		const fresh = !desired.desiredTables.some((table) => existingState.existing.has(table.name));
 
-			if (fk) {
-				fks.push(fk);
-			}
+		if (fresh) {
+			this.applyDiff(desired, existingState);
+			registry.record(migration.id, checksum);
+			return;
 		}
 
-		return `CREATE TABLE IF NOT EXISTS "${tableName}" (${columns.concat(fks).join(", ")})`;
+		await new GuidedMigrationRunner({
+			sqlite: this.sqlite,
+			kysely: this.kysely,
+			desired,
+			existing: existingState.existing,
+			migration,
+			checksum,
+			registry,
+		}).run();
 	}
 
-	private migrate() {
-		const desiredTables: DesiredTable[] = [];
-		const schemaIndexes = this.options.schema.indexes;
+	private rebuildFts(name: string) {
+		const plan = this.ftsPlans.get(name);
 
-		for (const [name, schema] of Object.entries(this.options.schema.tables)) {
-			const props = this.parseSchemaProps(schema);
-
-			const columns = props.map((prop) => {
-				const isNotNull = this.columnConstraint(prop) === "NOT NULL";
-				const defaultClause = this.defaultClause(prop);
-				const hasLiteralDefault = prop.generated !== "now" && defaultClause !== null;
-
-				return {
-					name: prop.key,
-					addable: !isNotNull || hasLiteralDefault,
-					columnDef: this.addColumnDef(prop),
-					type: this.sqlType(prop),
-					notnull: isNotNull,
-					defaultValue: defaultClause
-						? defaultClause.replace("DEFAULT ", "").replace(/^\((.+)\)$/, "$1")
-						: null,
-					unique: !!prop.meta?.unique,
-					references: prop.meta?.references ?? null,
-					onDelete: prop.meta?.onDelete?.toUpperCase() ?? null,
-				};
-			});
-
-			const indexes = (schemaIndexes?.[name] ?? []).map((indexDef) => ({
-				name: this.generateIndexName(name, indexDef.columns, indexDef.unique ?? false),
-				columns: indexDef.columns,
-				sql: this.generateCreateIndexSQL(name, indexDef),
-			}));
-
-			desiredTables.push({
-				name,
-				sql: this.generateCreateTableSQL(name, props),
-				columns,
-				indexes,
-			});
+		if (!plan) {
+			throw new Error(`FTS "${name}" is not declared in the schema`);
 		}
 
-		const existing = new Introspector(this.sqlite).introspect();
-		const ops = new Differ(desiredTables, existing).diff();
-
-		new Executor(this.sqlite, ops).execute();
+		this.sqlite.transaction(() => {
+			this.sqlite.run(`DELETE FROM "${name}"`);
+			this.sqlite.run(generateRebuildSelect(plan));
+		})();
 	}
 
-	private generateIndexName(tableName: string, columns: string[], unique: boolean) {
-		const prefix = unique ? "ux" : "ix";
+	private rebuildVectorIndex(name: string) {
+		const plan = this.vectorIndexPlans.get(name);
 
-		return `${prefix}_${tableName}_${columns.join("_")}`;
-	}
-
-	private generateCreateIndexSQL(tableName: string, indexDef: IndexDefinition) {
-		const indexName = this.generateIndexName(tableName, indexDef.columns, indexDef.unique ?? false);
-		const unique = indexDef.unique ? "UNIQUE " : "";
-		const columns = indexDef.columns.map((c) => `"${c}"`).join(", ");
+		if (!plan) {
+			throw new Error(`vector index "${name}" is not declared in the schema`);
+		}
 
-		return `CREATE ${unique}INDEX "${indexName}" ON "${tableName}" (${columns})`;
+		this.sqlite.transaction(() => {
+			this.sqlite.run(`DELETE FROM "${name}"`);
+			this.sqlite.run(generateRebuildVectorSql(plan));
+		})();
 	}
 
 	reset(table?: keyof T & string): void {

package/src/dialect/serialize.ts

@@ -1,6 +1,6 @@
 export function serializeParam(p: unknown): unknown {
 	if (p instanceof Date) {
-		return Math.floor(p.getTime() / 1000);
+		return p.getTime();
 	}
 
 	if (typeof p === "object" && p !== null && !ArrayBuffer.isView(p)) {

package/src/env.ts

@@ -3,7 +3,8 @@ export type DbFieldMeta = {
 	unique?: boolean;
 	references?: `${string}.${string}`;
 	onDelete?: "cascade" | "set null" | "restrict";
-	_generated?: "autoincrement" | "now";
+	_generated?: "autoincrement" | "now" | "vector";
+	_vectorDim?: number;
 };
 
 declare global {