feat: search index factory

packages/storage/sqlite-node/src/sqlite/index.ts

@@ -1,4 +1,5 @@
 export * from "./migrations.ts";
 export * from "./repo.ts";
+export * from "./search-index.ts";
 export * from "./storage/index.ts";
 export * from "./types.ts";

packages/storage/sqlite-node/src/sqlite/repo.ts

@@ -14,7 +14,7 @@ import {
 	toSessionSearchIndexRecord,
 } from "@earendil-works/pi-agent-core";
 import { applyMigrations } from "./migrations.ts";
-import { SqliteSessionSearchIndex } from "./search-index-sink.ts";
+import { SqliteSessionSearchIndex } from "./search-index.ts";
 import { SqliteSessionStorage } from "./storage/index.ts";
 import { decodeEntry, type SessionEntryRow } from "./storage/session-entries.ts";
 import { rowToMetadata, type SessionRow } from "./storage/sessions.ts";
@@ -26,6 +26,7 @@ import type {
 	SqliteSessionMetadata,
 	SqliteSessionRepoApi,
 	SqliteSessionRepoEnv,
+	SqliteSessionSearchIndexFactory,
 } from "./types.ts";
 
 function getParentPath(path: string): string {
@@ -53,12 +54,20 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 	private readonly env: SqliteSessionRepoEnv;
 	private readonly sqlite: SqliteDatabaseFactory;
 	private readonly databasePathInput: string;
+	private readonly searchIndexFactory: SqliteSessionSearchIndexFactory;
 	private databasePath: string | undefined;
 
-	constructor(options: { env: SqliteSessionRepoEnv; sqlite: SqliteDatabaseFactory; databasePath: string }) {
+	constructor(options: {
+		env: SqliteSessionRepoEnv;
+		sqlite: SqliteDatabaseFactory;
+		databasePath: string;
+		searchIndexFactory?: SqliteSessionSearchIndexFactory;
+	}) {
 		this.env = options.env;
 		this.sqlite = options.sqlite;
 		this.databasePathInput = options.databasePath;
+		this.searchIndexFactory =
+			options.searchIndexFactory ?? (({ db, databasePath }) => new SqliteSessionSearchIndex(db, databasePath));
 	}
 
 	private async getDatabasePath(): Promise<string> {
@@ -71,6 +80,10 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 		return this.databasePath;
 	}
 
+	private async createSearchIndex(db: SqliteDatabase) {
+		return this.searchIndexFactory({ db, databasePath: await this.getDatabasePath() });
+	}
+
 	private async ensureDatabaseDir(): Promise<void> {
 		const path = await this.getDatabasePath();
 		const directory = getParentPath(path);
@@ -103,7 +116,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 				parentSessionId: options.parentSessionId,
 				metadata: options.metadata,
 			});
-			return toSession(storage, new SqliteSessionSearchIndex(db, await this.getDatabasePath()));
+			return toSession(storage, await this.createSearchIndex(db));
 		} catch (error) {
 			await db.close();
 			throw error;
@@ -119,7 +132,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 		const db = await this.openDatabase();
 		try {
 			const storage = await SqliteSessionStorage.open(db, metadata);
-			return toSession(storage, new SqliteSessionSearchIndex(db, await this.getDatabasePath()));
+			return toSession(storage, await this.createSearchIndex(db));
 		} catch (error) {
 			await db.close();
 			throw error;
@@ -155,7 +168,7 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 		if (!getFileSystemResultOrThrow(await this.env.exists(path), `Failed to check database ${path}`)) return [];
 		const db = await this.openDatabase();
 		try {
-			return await new SqliteSessionSearchIndex(db, path).search(options);
+			return await (await this.createSearchIndex(db)).search(options);
 		} finally {
 			await db.close();
 		}
@@ -165,14 +178,14 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 		const db = await this.openDatabase();
 		try {
 			await db.transaction(async () => {
-				const searchIndexSink = new SqliteSessionSearchIndex(db, await this.getDatabasePath());
+				const searchIndex = await this.createSearchIndex(db);
 				const entries = await db
 					.prepare(
 						"SELECT rowid, session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ?",
 					)
 					.all<SessionEntryRow & { rowid: number }>(metadata.id);
 				for (const row of entries) {
-					await searchIndexSink.remove(toSessionSearchIndexRecord(metadata, decodeEntry(row)));
+					await searchIndex.remove(toSessionSearchIndexRecord(metadata, decodeEntry(row)));
 				}
 				await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(metadata.id);
 				await db.prepare("DELETE FROM session_entries WHERE session_id = ?").run(metadata.id);
@@ -209,13 +222,13 @@ export class SqliteSessionRepo implements SqliteSessionRepoApi {
 				parentSessionId: options.parentSessionId ?? sourceMetadata.id,
 				metadata: options.metadata ?? sourceMetadata.metadata,
 			});
-			const searchIndexSink = new SqliteSessionSearchIndex(db, await this.getDatabasePath());
+			const searchIndex = await this.createSearchIndex(db);
 			const metadata = await storage.getMetadata();
 			for (const entry of forkedEntries) {
 				await storage.appendEntry(entry);
-				await searchIndexSink.write(toSessionSearchIndexRecord(metadata, entry));
+				await searchIndex.write(toSessionSearchIndexRecord(metadata, entry));
 			}
-			return toSession(storage, searchIndexSink);
+			return toSession(storage, searchIndex);
 		} catch (error) {
 			await db.close();
 			throw error;

packages/storage/sqlite-node/src/sqlite/search-index.ts

@@ -1,15 +1,15 @@
 import type {
 	SessionSearchHit,
+	SessionSearchIndex,
 	SessionSearchIndexRecord,
-	SessionSearchIndexSink,
 	SessionSearchOptions,
 } from "@earendil-works/pi-agent-core";
 import { SessionError } from "@earendil-works/pi-agent-core";
 import { rowToMetadata, type SessionRow } from "./storage/sessions.ts";
 import type { SqliteDatabase, SqliteSessionMetadata } from "./types.ts";
 
-/** SQLite-backed search index and sink for persisted session entries. */
-export class SqliteSessionSearchIndex implements SessionSearchIndexSink {
+/** SQLite-backed search index for persisted session entries. */
+export class SqliteSessionSearchIndex implements SessionSearchIndex<SqliteSessionMetadata> {
 	private readonly db: SqliteDatabase;
 	private readonly databasePath: string;
 
@@ -18,7 +18,7 @@ export class SqliteSessionSearchIndex implements SessionSearchIndexSink {
 		this.databasePath = databasePath;
 	}
 
-	async write(record: SessionSearchIndexRecord): Promise<void> {
+	async write(record: SessionSearchIndexRecord<SqliteSessionMetadata>): Promise<void> {
 		const row = await this.db
 			.prepare("SELECT rowid FROM session_entries WHERE session_id = ? AND id = ?")
 			.get<{ rowid: number }>(record.metadata.id, record.entry.id);
@@ -30,7 +30,7 @@ export class SqliteSessionSearchIndex implements SessionSearchIndexSink {
 			.run(row.rowid, JSON.stringify(record.entry));
 	}
 
-	async remove(record: SessionSearchIndexRecord): Promise<void> {
+	async remove(record: SessionSearchIndexRecord<SqliteSessionMetadata>): Promise<void> {
 		const row = await this.db
 			.prepare("SELECT rowid FROM session_entries WHERE session_id = ? AND id = ?")
 			.get<{ rowid: number }>(record.metadata.id, record.entry.id);

packages/storage/sqlite-node/src/sqlite/types.ts

@@ -1,4 +1,10 @@
-import type { FileSystem, SessionCreateOptions, SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core";
+import type {
+	FileSystem,
+	SessionCreateOptions,
+	SessionMetadata,
+	SessionRepo,
+	SessionSearchIndex,
+} from "@earendil-works/pi-agent-core";
 
 /** Result of a prepared SQLite statement execution. */
 export interface SqliteRunResult {
@@ -53,3 +59,8 @@ export interface SqliteSessionRepoApi
 	extends SessionRepo<SqliteSessionMetadata, SqliteSessionCreateOptions, SqliteSessionListOptions> {}
 
 export type SqliteSessionRepoEnv = Pick<FileSystem, "absolutePath" | "createDir" | "exists">;
+
+export type SqliteSessionSearchIndexFactory = (options: {
+	db: SqliteDatabase;
+	databasePath: string;
+}) => SessionSearchIndex<SqliteSessionMetadata>;