@lobomfz/db 0.4.0 → 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,7 +13,9 @@ import {
 	ColumnNode,
 	IdentifierNode,
 	ParensNode,
+	RawNode,
 	ReferenceNode,
+	SelectionNode,
 	SelectQueryNode,
 	TableNode,
 } from "kysely";
@@ -21,7 +24,11 @@ import { JsonParseError } from "./errors.js";
 import type { StructureProp } from "./types.js";
 import { ValidationError } from "./validation-error.js";
 
-type ColumnCoercion = "boolean" | "date" | { type: "json"; schema: Type };
+type ColumnCoercion =
+	| "boolean"
+	| "date"
+	| { type: "json"; schema: Type }
+	| { type: "vector"; dim: number };
 
 type ResolvedCoercion = { table: string; coercion: ColumnCoercion };
 
@@ -121,7 +128,27 @@ export class ResultHydrationPlugin implements KyselyPlugin {
 		}
 	}
 
+	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",
 		);
@@ -150,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":
@@ -225,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;
 
@@ -241,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,4 +1,4 @@
-import type { Generated } from "kysely";
+import type { ColumnType, Generated } from "kysely";
 import type { Type } from "arktype";
 
 export type ArkBranch = {
@@ -25,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>;
@@ -48,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;
@@ -70,18 +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;
 };
 
-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/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`);
+	};
+}

package/src/vector-index/resolve.ts

@@ -0,0 +1,78 @@
+import type { VectorIndexEntry, VectorMetric } from "../types.js";
+
+export type VectorAnchorInfo = {
+	columns: Set<string>;
+	pkColumn: string | null;
+	pkIsInteger: boolean;
+	vectorColumns: Map<string, number>;
+};
+
+export type ResolvedVectorIndex = {
+	name: string;
+	anchor: string;
+	anchorPkColumn: string;
+	column: string;
+	dim: number;
+	metric: VectorMetric;
+};
+
+const DEFAULT_METRIC: VectorMetric = "cosine";
+
+const VALID_METRICS = new Set<VectorMetric>(["cosine", "l2", "l1"]);
+
+export function resolveVectorIndex(
+	name: string,
+	entry: VectorIndexEntry<any>,
+	anchors: Map<string, VectorAnchorInfo>,
+) {
+	const anchor = anchors.get(entry.anchor);
+
+	if (!anchor) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" is not declared in schema`,
+		);
+	}
+
+	if (!anchor.pkColumn) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" has no primary key; anchor PK must be an integer`,
+		);
+	}
+
+	if (!anchor.pkIsInteger) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" primary key "${anchor.pkColumn}" must be an integer`,
+		);
+	}
+
+	if (!anchor.columns.has(entry.column)) {
+		throw new Error(
+			`VectorIndex "${name}": column "${entry.column}" does not exist on anchor "${entry.anchor}"`,
+		);
+	}
+
+	const dim = anchor.vectorColumns.get(entry.column);
+
+	if (dim === undefined) {
+		throw new Error(
+			`VectorIndex "${name}": column "${entry.anchor}.${entry.column}" is not a vector column`,
+		);
+	}
+
+	const metric = entry.metric ?? DEFAULT_METRIC;
+
+	if (!VALID_METRICS.has(metric)) {
+		throw new Error(
+			`VectorIndex "${name}": unknown metric "${metric}", expected one of cosine, l2, l1`,
+		);
+	}
+
+	return {
+		name,
+		anchor: entry.anchor,
+		anchorPkColumn: anchor.pkColumn,
+		column: entry.column,
+		dim,
+		metric,
+	};
+}