@lobomfz/db 0.3.9 → 0.4.1

package/src/migration/types.ts

@@ -39,13 +39,67 @@ export type DropIndexOp = {
 	index: string;
 };
 
+export type FtsTriggerSpec = {
+	name: string;
+	sql: string;
+};
+
+export type CreateFtsOp = {
+	type: "CreateFts";
+	name: string;
+	createSql: string;
+	triggers: FtsTriggerSpec[];
+};
+
+export type DropFtsOp = {
+	type: "DropFts";
+	name: string;
+	triggerNames: string[];
+};
+
+export type RebuildFtsOp = {
+	type: "RebuildFts";
+	name: string;
+	rebuildSql: string;
+};
+
+export type VectorTriggerSpec = {
+	name: string;
+	sql: string;
+};
+
+export type CreateVectorIndexOp = {
+	type: "CreateVectorIndex";
+	name: string;
+	createSql: string;
+	triggers: VectorTriggerSpec[];
+};
+
+export type DropVectorIndexOp = {
+	type: "DropVectorIndex";
+	name: string;
+	triggerNames: string[];
+};
+
+export type RebuildVectorIndexOp = {
+	type: "RebuildVectorIndex";
+	name: string;
+	rebuildSql: string;
+};
+
 export type MigrationOp =
 	| CreateTableOp
 	| DropTableOp
 	| AddColumnOp
 	| RebuildTableOp
 	| CreateIndexOp
-	| DropIndexOp;
+	| DropIndexOp
+	| CreateFtsOp
+	| DropFtsOp
+	| RebuildFtsOp
+	| CreateVectorIndexOp
+	| DropVectorIndexOp
+	| RebuildVectorIndexOp;
 
 export type ColumnSchema = {
 	name: string;
@@ -72,3 +126,16 @@ export type IntrospectedTable = {
 	indexes: IntrospectedIndex[];
 	hasData: boolean;
 };
+
+export type IntrospectedFts = {
+	name: string;
+	columns: string[];
+	sql: string;
+	triggers: Map<string, string>;
+};
+
+export type IntrospectedVectorIndex = {
+	name: string;
+	sql: string;
+	triggers: Map<string, string>;
+};

package/src/plugin.ts

@@ -5,6 +5,7 @@ import {
 	type OperationNode,
 	type QueryId,
 	type RootOperationNode,
+	type SelectionNode as KyselySelectionNode,
 	type UnknownRow,
 	AggregateFunctionNode,
 	AliasNode,
@@ -12,16 +13,22 @@ import {
 	ColumnNode,
 	IdentifierNode,
 	ParensNode,
+	RawNode,
 	ReferenceNode,
+	SelectionNode,
 	SelectQueryNode,
 	TableNode,
 } from "kysely";
 
 import { JsonParseError } from "./errors.js";
-import { JsonValidationError } from "./validation-error.js";
+import type { StructureProp } from "./types.js";
+import { ValidationError } from "./validation-error.js";
 
-export type ColumnCoercion = "boolean" | "date" | { type: "json"; schema: Type };
-export type ColumnsMap = Map<string, Map<string, ColumnCoercion>>;
+type ColumnCoercion =
+	| "boolean"
+	| "date"
+	| { type: "json"; schema: Type }
+	| { type: "vector"; dim: number };
 
 type ResolvedCoercion = { table: string; coercion: ColumnCoercion };
 
@@ -79,13 +86,90 @@ const jsonObjectFromFragments = [
 const typePreservingAggregateFunctions = new Set(["max", "min"]);
 
 export class ResultHydrationPlugin implements KyselyPlugin {
+	private columns = new Map<string, Map<string, ColumnCoercion>>();
+	private tableColumns = new Map<string, Set<string>>();
 	private queryPlans = new WeakMap<QueryId, QueryPlan>();
 
 	constructor(
-		private columns: ColumnsMap,
-		private tableColumns: Map<string, Set<string>>,
+		schemas: Map<string, Type>,
 		private validation: { onRead: boolean },
-	) {}
+	) {
+		this.registerSchemas(schemas);
+	}
+
+	private registerSchemas(schemas: Map<string, Type>) {
+		for (const [table, schema] of schemas) {
+			const structureProps = (schema as any).structure?.props as StructureProp[] | undefined;
+
+			if (!structureProps) {
+				continue;
+			}
+
+			this.tableColumns.set(
+				table,
+				new Set(structureProps.map((prop) => prop.key)),
+			);
+
+			const columns = new Map<string, ColumnCoercion>();
+
+			for (const prop of structureProps) {
+				const coercion = this.getColumnCoercion(prop, schema);
+
+				if (!coercion) {
+					continue;
+				}
+
+				columns.set(prop.key, coercion);
+			}
+
+			if (columns.size > 0) {
+				this.columns.set(table, columns);
+			}
+		}
+	}
+
+	private extractFieldMeta(prop: StructureProp) {
+		const rootMeta = (prop.value as any).meta as
+			| { _generated?: string; _vectorDim?: number }
+			| undefined;
+		const branchWithMeta = prop.value.branches.find(
+			(b) => (b as any).meta && Object.keys((b as any).meta).length > 0,
+		);
+		const fromBranch = (branchWithMeta as any)?.meta as
+			| { _generated?: string; _vectorDim?: number }
+			| undefined;
+
+		return { ...fromBranch, ...rootMeta };
+	}
+
+	private getColumnCoercion(prop: StructureProp, parentSchema: Type) {
+		const meta = this.extractFieldMeta(prop);
+
+		if (meta._generated === "vector" && typeof meta._vectorDim === "number") {
+			return { type: "vector", dim: meta._vectorDim } satisfies ColumnCoercion;
+		}
+
+		const concrete = prop.value.branches.filter(
+			(branch) => branch.unit !== null && branch.domain !== "undefined",
+		);
+
+		if (prop.value.proto === Date || concrete.some((branch) => branch.proto === Date)) {
+			return "date" satisfies ColumnCoercion;
+		}
+
+		if (concrete.length > 0 && concrete.every((branch) => branch.domain === "boolean")) {
+			return "boolean" satisfies ColumnCoercion;
+		}
+
+		if (concrete.some((branch) => !!branch.structure)) {
+			return {
+				type: "json",
+				schema: (parentSchema as any).get(prop.key) as Type,
+			} satisfies ColumnCoercion;
+		}
+
+		return null;
+	}
 
 	transformQuery: KyselyPlugin["transformQuery"] = (args) => {
 		this.queryPlans.set(args.queryId, {
@@ -93,9 +177,151 @@ export class ResultHydrationPlugin implements KyselyPlugin {
 			selectionPlans: this.getSelectionPlans(args.node),
 		});
 
-		return args.node;
+		return this.rewriteVectorsInJsonHelpers(args.node);
 	};
 
+	private rewriteVectorsInJsonHelpers(node: RootOperationNode): RootOperationNode {
+		if (!SelectQueryNode.is(node)) {
+			return node;
+		}
+
+		return this.rewriteSelectQueryNode(node) as RootOperationNode;
+	}
+
+	private rewriteSelectQueryNode(node: SelectQueryNode): SelectQueryNode {
+		if (!node.selections) {
+			return node;
+		}
+
+		let changed = false;
+
+		const newSelections = node.selections.map((sel) => {
+			const rewritten = this.rewriteTopLevelSelection(sel);
+
+			if (rewritten !== sel) {
+				changed = true;
+			}
+
+			return rewritten;
+		});
+
+		if (!changed) {
+			return node;
+		}
+
+		return Object.freeze({ ...node, selections: Object.freeze(newSelections) });
+	}
+
+	private rewriteTopLevelSelection(sel: KyselySelectionNode): KyselySelectionNode {
+		const rewritten = this.rewriteNodeForJsonHelpers(sel.selection);
+
+		if (rewritten === sel.selection) {
+			return sel;
+		}
+
+		return SelectionNode.create(rewritten as any);
+	}
+
+	private rewriteNodeForJsonHelpers(node: OperationNode): OperationNode {
+		if (AliasNode.is(node)) {
+			const rewrittenInner = this.rewriteNodeForJsonHelpers(node.node);
+
+			if (rewrittenInner === node.node) {
+				return node;
+			}
+
+			return AliasNode.create(rewrittenInner, node.alias as any);
+		}
+
+		if (!this.isRawNode(node)) {
+			return node;
+		}
+
+		const helper = this.getJsonHelper(node);
+
+		if (!helper) {
+			return node;
+		}
+
+		return this.rewriteJsonHelper(node, helper.query);
+	}
+
+	private rewriteJsonHelper(rawNode: RawOperationNode, subquery: SelectQueryNode) {
+		const scope = this.getTableScope(subquery);
+
+		if (!subquery.selections) {
+			return rawNode;
+		}
+
+		let changed = false;
+
+		const newSubquerySelections = subquery.selections.map((sel) => {
+			const rewrittenNested = this.rewriteNodeForJsonHelpers(sel.selection);
+			const withNested =
+				rewrittenNested === sel.selection
+					? sel
+					: SelectionNode.create(rewrittenNested as any);
+
+			if (withNested !== sel) {
+				changed = true;
+			}
+
+			const wrapped = this.wrapVectorSelectionWithHex(withNested, scope);
+
+			if (wrapped !== withNested) {
+				changed = true;
+			}
+
+			return wrapped;
+		});
+
+		if (!changed) {
+			return rawNode;
+		}
+
+		const newSubquery = Object.freeze({
+			...subquery,
+			selections: Object.freeze(newSubquerySelections),
+		});
+		const newParameters = [...rawNode.parameters];
+		newParameters[1] = newSubquery;
+
+		return Object.freeze({
+			...rawNode,
+			parameters: Object.freeze(newParameters),
+		});
+	}
+
+	private wrapVectorSelectionWithHex(sel: KyselySelectionNode, scope: Map<string, string>) {
+		const selectionNode = sel.selection;
+		const output = this.getSelectionOutputName(selectionNode);
+
+		if (!output) {
+			return sel;
+		}
+
+		const inner = AliasNode.is(selectionNode) ? selectionNode.node : selectionNode;
+
+		if (!ReferenceNode.is(inner) && !ColumnNode.is(inner)) {
+			return sel;
+		}
+
+		const resolved = this.resolveReferenceCoercion(inner, scope);
+
+		if (!resolved) {
+			return sel;
+		}
+
+		if (typeof resolved.coercion !== "object" || resolved.coercion.type !== "vector") {
+			return sel;
+		}
+
+		const hexNode = RawNode.create(["hex(", ")"], [inner]);
+		const aliased = AliasNode.create(hexNode, IdentifierNode.create(output));
+
+		return SelectionNode.create(aliased);
+	}
+
 	private getTableFromNode(node: RootOperationNode) {
 		switch (node.kind) {
 			case "InsertQueryNode":
@@ -141,7 +367,7 @@ export class ResultHydrationPlugin implements KyselyPlugin {
 		const result = schema(value);
 
 		if (result instanceof type.errors) {
-			throw new JsonValidationError(table, col, result.summary);
+			throw new ValidationError(table, result.summary, col);
 		}
 	}
 
@@ -168,12 +394,16 @@ export class ResultHydrationPlugin implements KyselyPlugin {
 
 		if (plan.coercion === "date") {
 			if (typeof value === "number") {
-				return new Date(value * 1000);
+				return new Date(value);
 			}
 
 			return value;
 		}
 
+		if (plan.coercion.type === "vector") {
+			return this.hydrateVector(value);
+		}
+
 		const parsed =
 			typeof value === "string" ? this.parseJson(plan.table, plan.column, value) : value;
 
@@ -184,6 +414,36 @@ export class ResultHydrationPlugin implements KyselyPlugin {
 		return parsed;
 	}
 
+	private hydrateVector(value: unknown) {
+		if (value instanceof Float32Array) {
+			return value;
+		}
+
+		if (value instanceof Uint8Array) {
+			return new Float32Array(
+				value.buffer.slice(value.byteOffset, value.byteOffset + value.byteLength),
+			);
+		}
+
+		if (typeof value === "string") {
+			return this.hexToFloat32Array(value);
+		}
+
+		return value;
+	}
+
+	private hexToFloat32Array(hex: string) {
+		const byteLength = hex.length / 2;
+		const buffer = new ArrayBuffer(byteLength);
+		const bytes = new Uint8Array(buffer);
+
+		for (let i = 0; i < byteLength; i++) {
+			bytes[i] = parseInt(hex.slice(i * 2, i * 2 + 2), 16);
+		}
+
+		return new Float32Array(buffer);
+	}
+
 	private parseStructuredValue(table: string, column: string, value: unknown) {
 		if (value === null || value === undefined) {
 			return value;

package/src/types.ts

@@ -1,6 +1,23 @@
-import type { Generated } from "kysely";
+import type { ColumnType, Generated } from "kysely";
 import type { Type } from "arktype";
 
+export type ArkBranch = {
+	domain?: string;
+	proto?: unknown;
+	unit?: unknown;
+	structure?: unknown;
+};
+
+export type StructureProp = {
+	key: string;
+	required: boolean;
+	inner: { default?: unknown };
+	value: Type & {
+		branches: ArkBranch[];
+		proto?: unknown;
+	};
+};
+
 type ExtractInput<T> = T extends { inferIn: infer I } ? I : never;
 type ExtractOutput<T> = T extends { infer: infer O } ? O : never;
 
@@ -8,12 +25,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>;
@@ -31,12 +57,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;
@@ -53,19 +79,91 @@ 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;
-	onWrite?: boolean;
 };
 
-export type DatabaseOptions<T extends SchemaRecord> = {
+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[];
 };

package/src/validation-error.ts

@@ -1,11 +1,15 @@
-export class JsonValidationError extends Error {
+export class ValidationError extends Error {
 	constructor(
 		readonly table: string,
-		readonly column: string,
 		readonly summary: string,
+		readonly column: string | null = null,
 	) {
-		super(`JSON validation failed for ${table}.${column}: ${summary}`);
+		super(
+			column
+				? `Validation failed for ${table}.${column}: ${summary}`
+				: `Validation failed for ${table}: ${summary}`,
+		);
 
-		this.name = "JsonValidationError";
+		this.name = "ValidationError";
 	}
 }

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