@lobomfz/db 0.4.0 → 0.5.0

package/src/migration/introspect.ts

@@ -1,6 +1,15 @@
 import type { Database } from "bun:sqlite";
 
-import type { IntrospectedColumn, IntrospectedIndex, IntrospectedTable } from "./types.js";
+import { FTS_TRIGGER_PREFIX } from "../fts/sql.js";
+import { VEC_TRIGGER_PREFIX } from "../vector-index/sql.js";
+import {
+	MIGRATIONS_TABLE,
+	type IntrospectedColumn,
+	type IntrospectedFts,
+	type IntrospectedIndex,
+	type IntrospectedTable,
+	type IntrospectedVectorIndex,
+} from "./types.js";
 
 type TableListRow = {
 	name: string;
@@ -12,6 +21,7 @@ type TableXInfoRow = {
 	type: string;
 	notnull: number;
 	dflt_value: string | null;
+	hidden: number;
 };
 
 type IndexListRow = {
@@ -31,18 +41,34 @@ type ForeignKeyListRow = {
 	on_delete: string;
 };
 
+type MasterRow = {
+	name: string;
+	sql: string | null;
+};
+
 export class Introspector {
 	constructor(private db: Database) {}
 
 	introspect() {
 		const tables = new Map<string, IntrospectedTable>();
 		const tableRows = this.db.prepare("PRAGMA table_list").all() as TableListRow[];
+		const virtualPrefixes = tableRows
+			.filter((r) => r.type === "virtual")
+			.map((r) => `${r.name}_`);
 
 		for (const row of tableRows) {
 			if (row.type !== "table" || row.name.startsWith("sqlite_")) {
 				continue;
 			}
 
+			if (row.name === MIGRATIONS_TABLE) {
+				continue;
+			}
+
+			if (virtualPrefixes.some((prefix) => row.name.startsWith(prefix))) {
+				continue;
+			}
+
 			const indexRows = this.db.prepare(`PRAGMA index_list("${row.name}")`).all() as IndexListRow[];
 			const uniqueCols = this.uniqueColumns(indexRows);
 			const fkMap = this.foreignKeys(row.name);
@@ -56,6 +82,117 @@ export class Introspector {
 		return tables;
 	}
 
+	hasVec0Tables() {
+		const row = this.db
+			.prepare(
+				"SELECT 1 FROM sqlite_master WHERE type = 'table' AND sql LIKE '%USING vec0%' LIMIT 1",
+			)
+			.get();
+
+		return row !== null;
+	}
+
+	introspectFts() {
+		const fts = new Map<string, IntrospectedFts>();
+		const tableRows = this.db.prepare("PRAGMA table_list").all() as TableListRow[];
+		const libraryTriggers = this.triggersByPrefix(FTS_TRIGGER_PREFIX);
+
+		for (const row of tableRows) {
+			if (row.type !== "virtual" || row.name.startsWith("sqlite_")) {
+				continue;
+			}
+
+			const master = this.db
+				.prepare(`SELECT name, sql FROM sqlite_master WHERE name = ? AND type = 'table'`)
+				.get(row.name) as MasterRow | null;
+
+			if (!master?.sql || !/USING\s+fts5/i.test(master.sql)) {
+				continue;
+			}
+
+			const colRows = this.db.prepare(`PRAGMA table_xinfo("${row.name}")`).all() as TableXInfoRow[];
+			const columns = colRows.filter((c) => c.hidden !== 1).map((c) => c.name);
+
+			fts.set(row.name, {
+				name: row.name,
+				columns,
+				sql: master.sql,
+				triggers: new Map(),
+			});
+		}
+
+		this.assignTriggersByPrefix(libraryTriggers, FTS_TRIGGER_PREFIX, fts);
+
+		return fts;
+	}
+
+	introspectVectorIndexes() {
+		const indexes = new Map<string, IntrospectedVectorIndex>();
+		const tableRows = this.db.prepare("PRAGMA table_list").all() as TableListRow[];
+		const libraryTriggers = this.triggersByPrefix(VEC_TRIGGER_PREFIX);
+
+		for (const row of tableRows) {
+			if (row.type !== "virtual" || row.name.startsWith("sqlite_")) {
+				continue;
+			}
+
+			const master = this.db
+				.prepare(`SELECT name, sql FROM sqlite_master WHERE name = ? AND type = 'table'`)
+				.get(row.name) as MasterRow | null;
+
+			if (!master?.sql || !/USING\s+vec0/i.test(master.sql)) {
+				continue;
+			}
+
+			indexes.set(row.name, {
+				name: row.name,
+				sql: master.sql,
+				triggers: new Map(),
+			});
+		}
+
+		this.assignTriggersByPrefix(libraryTriggers, VEC_TRIGGER_PREFIX, indexes);
+
+		return indexes;
+	}
+
+	private assignTriggersByPrefix(
+		triggers: Map<string, string>,
+		prefix: string,
+		target: Map<string, { triggers: Map<string, string> }>,
+	) {
+		const sortedNames = [...target.keys()].sort((a, b) => b.length - a.length);
+
+		for (const [triggerName, triggerSql] of triggers) {
+			const withoutPrefix = triggerName.slice(prefix.length);
+
+			for (const name of sortedNames) {
+				if (withoutPrefix.startsWith(`${name}_`)) {
+					target.get(name)!.triggers.set(triggerName, triggerSql);
+					break;
+				}
+			}
+		}
+	}
+
+	private triggersByPrefix(prefix: string) {
+		const rows = this.db
+			.prepare(
+				`SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND name LIKE '${prefix}%'`,
+			)
+			.all() as MasterRow[];
+
+		const map = new Map<string, string>();
+
+		for (const row of rows) {
+			if (row.sql !== null) {
+				map.set(row.name, row.sql);
+			}
+		}
+
+		return map;
+	}
+
 	private uniqueColumns(indexRows: IndexListRow[]) {
 		const unique = new Set<string>();
 

package/src/migration/registry.ts

@@ -0,0 +1,43 @@
+import type { Database } from "bun:sqlite";
+
+import { MigrationDriftError } from "./drift-error.js";
+import { MIGRATIONS_TABLE } from "./types.js";
+
+export class MigrationRegistry {
+	constructor(private db: Database) {}
+
+	ensureTable() {
+		this.db.run(
+			`CREATE TABLE IF NOT EXISTS "${MIGRATIONS_TABLE}" ` +
+				`("id" TEXT PRIMARY KEY, "checksum" TEXT NOT NULL, "applied_at" INTEGER NOT NULL)`,
+		);
+	}
+
+	appliedChecksum(id: string) {
+		const row = this.db
+			.prepare(`SELECT "checksum" FROM "${MIGRATIONS_TABLE}" WHERE "id" = ?`)
+			.get(id) as { checksum: string } | null;
+
+		return row?.checksum ?? null;
+	}
+
+	assertNoDrift(id: string, checksum: string) {
+		const applied = this.appliedChecksum(id);
+
+		if (applied !== null && applied !== checksum) {
+			throw new MigrationDriftError(id);
+		}
+	}
+
+	record(id: string, checksum: string) {
+		this.db
+			.prepare(
+				`INSERT INTO "${MIGRATIONS_TABLE}" ("id", "checksum", "applied_at") VALUES (?, ?, ?)`,
+			)
+			.run(id, checksum, Date.now());
+	}
+
+	static checksum(run: (...args: any[]) => unknown) {
+		return new Bun.CryptoHasher("sha256").update(run.toString()).digest("hex");
+	}
+}

package/src/migration/types.ts

@@ -1,3 +1,6 @@
+import type { ResolvedFtsPlan } from "../fts/resolve.js";
+import type { ResolvedVectorIndex } from "../vector-index/resolve.js";
+
 export type CreateTableOp = {
 	type: "CreateTable";
 	table: string;
@@ -39,13 +42,75 @@ 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 const MIGRATIONS_TABLE = "__migrations";
+
+export type DestructiveChange = {
+	table: string;
+	kind: "DropTable" | "DropColumn" | "TypeChange";
+	columns?: string[];
+};
 
 export type ColumnSchema = {
 	name: string;
@@ -72,3 +137,47 @@ 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>;
+};
+
+export interface DesiredColumn extends ColumnSchema {
+	addable: boolean;
+	columnDef: string;
+	nullableAddDef: string;
+}
+
+export type DesiredIndex = {
+	name: string;
+	columns: string[];
+	sql: string;
+};
+
+export type DesiredTable = {
+	name: string;
+	sql: string;
+	columns: DesiredColumn[];
+	indexes?: DesiredIndex[];
+};
+
+export type DesiredSchema = {
+	desiredTables: DesiredTable[];
+	desiredFts?: ResolvedFtsPlan[];
+	desiredVectorIndexes?: ResolvedVectorIndex[];
+};
+
+export type ExistingSchema = {
+	existing: Map<string, IntrospectedTable>;
+	existingFts?: Map<string, IntrospectedFts>;
+	existingVectorIndexes?: Map<string, IntrospectedVectorIndex>;
+};

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;