@lobomfz/db 0.4.0 → 0.5.0

package/src/schema-compiler.ts

@@ -0,0 +1,362 @@
+import type { Type } from "arktype";
+
+import type { DbFieldMeta } from "./env.js";
+import type { FtsTableInfo } from "./fts/resolve.js";
+import type { GeneratedPreset } from "./generated.js";
+import type { DesiredColumn, DesiredTable } from "./migration/types.js";
+import type { IndexDefinition, StructureProp } from "./types.js";
+import type { VectorAnchorInfo } from "./vector-index/resolve.js";
+
+type Prop = {
+	key: string;
+	kind: "required" | "optional";
+	domain?: string;
+	nullable?: boolean;
+	isBoolean?: boolean;
+	isInteger?: boolean;
+	isDate?: boolean;
+	isBlob?: boolean;
+	isJson?: boolean;
+	isVector?: boolean;
+	vectorDim?: number;
+	jsonSchema?: Type;
+	meta?: DbFieldMeta;
+	generated?: GeneratedPreset;
+	defaultValue?: unknown;
+};
+
+const typeMap: Record<string, string> = {
+	string: "TEXT",
+	number: "REAL",
+};
+
+export class SchemaCompiler {
+	compileTable(name: string, schema: Type, indexDefs: IndexDefinition[]) {
+		const props = this.parseSchemaProps(schema);
+		const columns = props.map((prop) => this.compileColumn(prop));
+
+		const indexes = indexDefs.map((indexDef) => ({
+			name: this.generateIndexName(name, indexDef.columns, indexDef.unique ?? false),
+			columns: indexDef.columns,
+			sql: this.generateCreateIndexSQL(name, indexDef),
+		}));
+
+		const table: DesiredTable = {
+			name,
+			sql: this.generateCreateTableSQL(name, props),
+			columns,
+			indexes,
+		};
+
+		return {
+			table,
+			tableInfo: this.buildTableInfo(props),
+			vectorAnchor: this.buildVectorAnchorInfo(props),
+		};
+	}
+
+	createWriteSchema(schema: Type) {
+		const autoIncrementColumns = this.parseSchemaProps(schema)
+			.filter((prop) => prop.generated === "autoincrement")
+			.map((prop) => prop.key);
+
+		if (autoIncrementColumns.length === 0) {
+			return schema;
+		}
+
+		return (schema as any).omit(...autoIncrementColumns) as Type;
+	}
+
+	private parseSchemaProps(schema: Type) {
+		const structureProps = (schema as any).structure?.props as StructureProp[] | undefined;
+
+		if (!structureProps) {
+			return [];
+		}
+
+		return structureProps.map((p) => this.normalizeProp(p, schema));
+	}
+
+	private normalizeProp(structureProp: StructureProp, parentSchema: Type) {
+		const { key, value: v, inner } = structureProp;
+		const kind: Prop["kind"] = structureProp.required ? "required" : "optional";
+		const defaultValue = inner.default;
+
+		const concrete = v.branches.filter((b) => b.unit !== null && b.domain !== "undefined");
+		const nullable = concrete.length < v.branches.length;
+
+		const branchMeta = v.branches.find((b) => b.meta && Object.keys(b.meta).length > 0)?.meta;
+		const meta = { ...branchMeta, ...v.meta };
+		const generated = meta._generated === "vector" ? undefined : meta._generated;
+
+		if (meta._generated === "vector") {
+			return {
+				key,
+				kind,
+				nullable,
+				isVector: true,
+				vectorDim: meta._vectorDim,
+				meta,
+				defaultValue,
+			};
+		}
+
+		if (v.proto === Date || concrete.some((b) => b.proto === Date)) {
+			return { key, kind, nullable, isDate: true, generated, defaultValue };
+		}
+
+		if (v.proto === Uint8Array || concrete.some((b) => b.proto === Uint8Array)) {
+			return {
+				key,
+				kind,
+				nullable,
+				isBlob: true,
+				meta,
+				generated,
+				defaultValue,
+			};
+		}
+
+		if (concrete.length > 0 && concrete.every((b) => b.domain === "boolean")) {
+			return { key, kind, nullable, isBoolean: true, generated, defaultValue };
+		}
+
+		if (concrete.some((b) => !!b.structure)) {
+			return {
+				key,
+				kind,
+				nullable,
+				isJson: true,
+				jsonSchema: (parentSchema as any).get(key) as Type,
+				meta,
+				generated,
+				defaultValue,
+			};
+		}
+
+		const branch = concrete[0];
+
+		return {
+			key,
+			kind,
+			nullable,
+			domain: branch?.domain,
+			isInteger: !!branch?.inner?.divisor,
+			meta,
+			generated,
+			defaultValue,
+		};
+	}
+
+	private compileColumn(prop: Prop): DesiredColumn {
+		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),
+			nullableAddDef: this.nullableAddColumnDef(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,
+		};
+	}
+
+	private sqlType(prop: Prop) {
+		if (prop.isJson) {
+			return "TEXT";
+		}
+
+		if (prop.isBlob || prop.isVector) {
+			return "BLOB";
+		}
+
+		if (prop.isDate || prop.isBoolean || prop.isInteger) {
+			return "INTEGER";
+		}
+
+		if (prop.domain) {
+			return typeMap[prop.domain] ?? "TEXT";
+		}
+
+		return "TEXT";
+	}
+
+	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;
+	}
+
+	private defaultClause(prop: Prop) {
+		if (prop.generated === "now") {
+			return "DEFAULT (CAST(unixepoch('subsec') * 1000 AS INTEGER))";
+		}
+
+		if (prop.defaultValue === undefined || prop.generated === "autoincrement") {
+			return null;
+		}
+
+		if (prop.defaultValue === null) {
+			return "DEFAULT NULL";
+		}
+
+		if (typeof prop.defaultValue === "string") {
+			return `DEFAULT '${prop.defaultValue}'`;
+		}
+
+		if (typeof prop.defaultValue === "number" || typeof prop.defaultValue === "boolean") {
+			return `DEFAULT ${prop.defaultValue}`;
+		}
+
+		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 referencesClause(prop: Prop) {
+		const ref = prop.meta?.references;
+
+		if (!ref) {
+			return null;
+		}
+
+		const [table, column] = ref.split(".");
+
+		let clause = ` REFERENCES "${table}"("${column}")`;
+
+		if (prop.meta?.onDelete) {
+			clause += ` ON DELETE ${prop.meta.onDelete.toUpperCase()}`;
+		}
+
+		return clause;
+	}
+
+	private foreignKey(prop: Prop) {
+		const clause = this.referencesClause(prop);
+
+		if (!clause) {
+			return null;
+		}
+
+		return `FOREIGN KEY ("${prop.key}")${clause}`;
+	}
+
+	private addColumnDef(prop: Prop) {
+		return `${this.columnDef(prop)}${this.referencesClause(prop) ?? ""}`;
+	}
+
+	private nullableAddColumnDef(prop: Prop) {
+		return `"${prop.key}" ${this.sqlType(prop)}${this.referencesClause(prop) ?? ""}`;
+	}
+
+	private generateCreateTableSQL(tableName: string, props: Prop[]) {
+		const columns: string[] = [];
+		const fks: string[] = [];
+
+		for (const prop of props) {
+			columns.push(this.columnDef(prop));
+
+			const fk = this.foreignKey(prop);
+
+			if (fk) {
+				fks.push(fk);
+			}
+		}
+
+		return `CREATE TABLE IF NOT EXISTS "${tableName}" (${columns.concat(fks).join(", ")})`;
+	}
+
+	private generateIndexName(tableName: string, columns: string[], unique: boolean) {
+		const prefix = unique ? "ux" : "ix";
+
+		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(", ");
+
+		return `CREATE ${unique}INDEX "${indexName}" ON "${tableName}" (${columns})`;
+	}
+
+	private buildTableInfo(props: Prop[]): FtsTableInfo {
+		const pkProp = props.find(
+			(p) => p.generated === "autoincrement" || p.meta?.primaryKey === true,
+		);
+
+		return {
+			columns: new Set(props.map((p) => p.key)),
+			pkColumn: pkProp?.key ?? null,
+			pkIsInteger: pkProp ? this.sqlType(pkProp) === "INTEGER" : false,
+			foreignKeys: props.flatMap((p) => {
+				const ref = p.meta?.references;
+
+				if (!ref) {
+					return [];
+				}
+
+				const [referencedTable, referencedColumn] = ref.split(".");
+
+				if (!referencedTable || !referencedColumn) {
+					throw new Error(
+						`Invalid references format "${ref}" for column "${p.key}": expected "table.column"`,
+					);
+				}
+
+				return [{ column: p.key, referencedTable, referencedColumn }];
+			}),
+		};
+	}
+
+	private buildVectorAnchorInfo(props: Prop[]): VectorAnchorInfo {
+		const pkProp = props.find(
+			(p) => p.generated === "autoincrement" || p.meta?.primaryKey === true,
+		);
+
+		const vectorColumns = new Map<string, number>();
+
+		for (const p of props) {
+			if (p.isVector && typeof p.vectorDim === "number") {
+				vectorColumns.set(p.key, p.vectorDim);
+			}
+		}
+
+		return {
+			columns: new Set(props.map((p) => p.key)),
+			pkColumn: pkProp?.key ?? null,
+			pkIsInteger: pkProp ? this.sqlType(pkProp) === "INTEGER" : false,
+			vectorColumns,
+		};
+	}
+}

package/src/types.ts

@@ -1,11 +1,15 @@
-import type { Generated } from "kysely";
+import type { ColumnType, Generated, Kysely } from "kysely";
 import type { Type } from "arktype";
 
+import type { DbFieldMeta } from "./env.js";
+
 export type ArkBranch = {
 	domain?: string;
 	proto?: unknown;
 	unit?: unknown;
 	structure?: unknown;
+	inner?: { divisor?: unknown };
+	meta?: DbFieldMeta;
 };
 
 export type StructureProp = {
@@ -15,6 +19,7 @@ export type StructureProp = {
 	value: Type & {
 		branches: ArkBranch[];
 		proto?: unknown;
+		meta: DbFieldMeta;
 	};
 };
 
@@ -25,12 +30,21 @@ type IsOptional<T, K extends keyof T> = undefined extends T[K] ? true : false;
 
 type TransformColumn<T> = T extends (infer U)[] ? U[] : T;
 
+type IsVectorColumn<TOutput, K extends keyof TOutput> =
+	NonNullable<TOutput[K]> extends Float32Array ? true : false;
+
+type VectorInsert = Float32Array | number[];
+
 type TransformField<TOutput, TInput, K extends keyof TOutput & keyof TInput> =
-	IsOptional<TInput, K> extends true
+	IsVectorColumn<TOutput, K> extends true
 		? IsOptional<TOutput, K> extends true
-			? TransformColumn<NonNullable<TOutput[K]>> | null
-			: Generated<TransformColumn<NonNullable<TOutput[K]>>>
-		: TransformColumn<TOutput[K]>;
+			? ColumnType<Float32Array | null, VectorInsert | null, VectorInsert | null>
+			: ColumnType<Float32Array, VectorInsert, VectorInsert>
+		: IsOptional<TInput, K> extends true
+			? IsOptional<TOutput, K> extends true
+				? TransformColumn<NonNullable<TOutput[K]>> | null
+				: Generated<TransformColumn<NonNullable<TOutput[K]>>>
+			: TransformColumn<TOutput[K]>;
 
 type TransformTable<TOutput, TInput> = {
 	[K in keyof TOutput & keyof TInput]: TransformField<TOutput, TInput, K>;
@@ -48,12 +62,12 @@ export type SqliteMasterRow = {
 	sql: string | null;
 };
 
-export type TablesFromSchemas<T extends SchemaRecord> = {
-	[K in keyof T]: InferTableType<T[K]>;
-} & { sqlite_master: SqliteMasterRow };
-
 type TableColumns<T extends SchemaRecord, K extends keyof T> = keyof ExtractOutput<T[K]> & string;
 
+type VectorColumnName<T extends SchemaRecord, K extends keyof T> = {
+	[C in TableColumns<T, K>]: NonNullable<ExtractOutput<T[K]>[C]> extends Float32Array ? C : never;
+}[TableColumns<T, K>];
+
 export type IndexDefinition<Columns extends string = string> = {
 	columns: Columns[];
 	unique?: boolean;
@@ -70,18 +84,102 @@ export type DatabasePragmas = {
 	busy_timeout_ms?: number;
 };
 
-export type DatabaseSchema<T extends SchemaRecord> = {
+export type FtsTokenizer =
+	| { type: "unicode61"; removeDiacritics?: 0 | 1 | 2; categories?: string }
+	| { type: "trigram" };
+
+export type FtsFieldRef = string | { source: string; via?: string };
+
+export type FtsEntry<T extends SchemaRecord> = {
+	anchor: keyof T & string;
+	fields: Record<string, FtsFieldRef>;
+	tokenizer?: FtsTokenizer;
+};
+
+export type FtsConfig<T extends SchemaRecord> = Record<string, FtsEntry<T>>;
+
+export type FtsRowType<E extends FtsEntry<any>> = { rowid: number } & {
+	[K in keyof E["fields"]]: string;
+};
+
+export type VectorMetric = "cosine" | "l2" | "l1";
+
+export type VectorIndexEntry<T extends SchemaRecord> = {
+	[K in keyof T & string]: {
+		anchor: K;
+		column: VectorColumnName<T, K>;
+		metric?: VectorMetric;
+	};
+}[keyof T & string];
+
+export type VectorIndexConfig<T extends SchemaRecord> = Record<string, VectorIndexEntry<T>>;
+
+export type DatabaseExtension = { kind: "sqlite-vec"; path?: string };
+
+export type DatabaseSchema<
+	T extends SchemaRecord,
+	F extends FtsConfig<T> = FtsConfig<T>,
+	V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+> = {
 	tables: T;
 	indexes?: IndexesConfig<T>;
+	fts?: F;
+	vectorIndexes?: V;
+};
+
+type FtsTables<F extends FtsConfig<any>> = string extends keyof F
+	? unknown
+	: { [K in keyof F]: FtsRowType<F[K]> };
+
+type VectorIndexRow = {
+	rowid: ColumnType<number, never, never>;
+	distance: ColumnType<number, never, never>;
 };
 
+type VectorIndexTables<V extends VectorIndexConfig<any>> = string extends keyof V
+	? unknown
+	: { [K in keyof V]: VectorIndexRow };
+
+export type TablesFromSchemas<
+	T extends SchemaRecord,
+	F extends FtsConfig<T> = FtsConfig<T>,
+	V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+> = {
+	[K in keyof T]: InferTableType<T[K]>;
+} & FtsTables<F> &
+	VectorIndexTables<V> & { sqlite_master: SqliteMasterRow };
+
+export type FtsApi<F extends FtsConfig<any>> = string extends keyof F
+	? never
+	: { rebuild(name: keyof F & string): void };
+
+export type VectorIndexApi<V extends VectorIndexConfig<any>> = string extends keyof V
+	? never
+	: { rebuild(name: keyof V & string): void };
+
 export type JsonValidation = {
 	onRead?: boolean;
 };
 
-export type DatabaseOptions<T extends SchemaRecord> = {
+export type DataMigration<
+	T extends SchemaRecord,
+	F extends FtsConfig<T> = FtsConfig<T>,
+	V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+> = {
+	id: string;
+	run: (ctx: { db: Kysely<TablesFromSchemas<T, F, V>> }) => Promise<void> | void;
+};
+
+export type DatabaseOptions<
+	T extends SchemaRecord,
+	F extends FtsConfig<T> = FtsConfig<T>,
+	V extends VectorIndexConfig<T> = VectorIndexConfig<T>,
+> = {
 	path: string;
-	schema: DatabaseSchema<T>;
+	schema: DatabaseSchema<T, F, V>;
 	pragmas?: DatabasePragmas;
 	validation?: JsonValidation;
+	extensions?: DatabaseExtension[];
+	failOnDataLoss?: boolean;
+	migration?: DataMigration<T, F, V>;
 };

package/src/vector-index/load.ts

@@ -0,0 +1,67 @@
+import type { Database as BunDatabase } from "bun:sqlite";
+import { createRequire } from "node:module";
+
+import type { DatabaseExtension } from "../types.js";
+
+export class ExtensionLoadError extends Error {
+	readonly path: string;
+	override readonly cause: unknown;
+
+	constructor(path: string, cause: unknown, hint: string) {
+		const message =
+			cause instanceof Error ? cause.message : typeof cause === "string" ? cause : String(cause);
+
+		super(`Failed to load sqlite-vec extension at "${path}": ${message}. ${hint}`);
+		this.name = "ExtensionLoadError";
+		this.path = path;
+		this.cause = cause;
+	}
+}
+
+const installHint = "Install it with `bun add sqlite-vec` (or npm install sqlite-vec).";
+
+function resolveSqliteVecPath(explicit: string | undefined) {
+	if (explicit) {
+		return explicit;
+	}
+
+	const require = createRequire(import.meta.url);
+
+	let mod: { getLoadablePath?: () => string };
+
+	try {
+		mod = require("sqlite-vec");
+	} catch (cause: any) {
+		throw new ExtensionLoadError("sqlite-vec", cause, installHint);
+	}
+
+	if (typeof mod.getLoadablePath !== "function") {
+		throw new ExtensionLoadError(
+			"sqlite-vec",
+			new Error("sqlite-vec does not export getLoadablePath"),
+			installHint,
+		);
+	}
+
+	try {
+		return mod.getLoadablePath();
+	} catch (cause: any) {
+		throw new ExtensionLoadError("sqlite-vec", cause, installHint);
+	}
+}
+
+export function loadExtensions(sqlite: BunDatabase, extensions: DatabaseExtension[]) {
+	for (const ext of extensions) {
+		if (ext.kind !== "sqlite-vec") {
+			continue;
+		}
+
+		const path = resolveSqliteVecPath(ext.path);
+
+		try {
+			sqlite.loadExtension(path);
+		} catch (cause: any) {
+			throw new ExtensionLoadError(path, cause, installHint);
+		}
+	}
+}

package/src/vector-index/nearest.ts

@@ -0,0 +1,68 @@
+import { sql, type SelectQueryBuilder } from "kysely";
+
+import type { ResolvedVectorIndex } from "./resolve.js";
+
+type IndexMeta = Pick<ResolvedVectorIndex, "anchor" | "anchorPkColumn" | "column" | "dim">;
+
+const registry = new Map<string, IndexMeta>();
+
+export function registerVectorIndex(index: ResolvedVectorIndex) {
+	registry.set(index.name, {
+		anchor: index.anchor,
+		anchorPkColumn: index.anchorPkColumn,
+		column: index.column,
+		dim: index.dim,
+	});
+}
+
+export function unregisterVectorIndex(name: string) {
+	registry.delete(name);
+}
+
+export type NearestOptions<IndexName extends string = string> = {
+	indexName: IndexName;
+	vector: Float32Array | number[];
+	k: number;
+};
+
+type WithDistance<O> = O & { distance: number };
+
+function vectorToBytes(v: Float32Array | number[]) {
+	const f = v instanceof Float32Array ? v : Float32Array.from(v);
+
+	return new Uint8Array(f.buffer, f.byteOffset, f.byteLength);
+}
+
+export function nearest<IndexName extends string>(options: NearestOptions<IndexName>) {
+	return <DB extends { [K in IndexName]: unknown }, TB extends keyof DB & string, O>(
+		qb: SelectQueryBuilder<DB, TB, O>,
+	): SelectQueryBuilder<DB, TB | (IndexName & keyof DB & string), WithDistance<O>> => {
+		const meta = registry.get(options.indexName);
+
+		if (!meta) {
+			throw new Error(
+				`Vector index "${options.indexName}" is not registered. Declare it in the schema's vectorIndexes and construct the Database before calling nearest.`,
+			);
+		}
+
+		const bytes = vectorToBytes(options.vector);
+		const indexName = options.indexName;
+		const column = meta.column;
+		const anchor = meta.anchor;
+		const anchorPk = meta.anchorPkColumn;
+
+		const anyQb = qb as unknown as {
+			innerJoin: (a: any, b: any, c: any) => any;
+			where: (expr: any) => any;
+			select: (expr: any) => any;
+			orderBy: (expr: any) => any;
+		};
+
+		return anyQb
+			.innerJoin(indexName, `${indexName}.rowid`, `${anchor}.${anchorPk}`)
+			.where(sql<boolean>`${sql.ref(`${indexName}.${column}`)} MATCH ${bytes}`)
+			.where(sql<boolean>`${sql.ref(`${indexName}.k`)} = ${options.k}`)
+			.select(sql<number>`${sql.ref(`${indexName}.distance`)}`.as("distance"))
+			.orderBy(`${indexName}.distance`);
+	};
+}