-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
community[minor]: Added
SQLiteRecordManager (#4321)
* created sqlite record manager and integration tests * Updated tests and implementation * Updated tests and implementation Updated implementation to make it cleaner * some tiny refactors * cr * update peer deps, add to config file * format --------- Co-authored-by: bracesproul <braceasproul@gmail.com>
- Loading branch information
1 parent
b52852b
commit 3429110
Showing
6 changed files
with
461 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,234 @@ | ||
| // eslint-disable-next-line import/no-extraneous-dependencies | ||
| import Database, { Database as DatabaseType, Statement } from "better-sqlite3"; | ||
| import { | ||
| ListKeyOptions, | ||
| RecordManagerInterface, | ||
| UpdateOptions, | ||
| } from "./base.js"; | ||
|
|
||
| interface TimeRow { | ||
| epoch: number; | ||
| } | ||
|
|
||
| interface KeyRecord { | ||
| key: string; | ||
| } | ||
|
|
||
| /** | ||
| * Options for configuring the SQLiteRecordManager class. | ||
| */ | ||
| export type SQLiteRecordManagerOptions = { | ||
| /** | ||
| * The file path of the SQLite database. | ||
| * One of either `localPath` or `connectionString` is required. | ||
| */ | ||
| localPath?: string; | ||
| /** | ||
| * The connection string of the SQLite database. | ||
| * One of either `localPath` or `connectionString` is required. | ||
| */ | ||
| connectionString?: string; | ||
| /** | ||
| * The name of the table in the SQLite database. | ||
| */ | ||
| tableName: string; | ||
| }; | ||
|
|
||
| export class SQLiteRecordManager implements RecordManagerInterface { | ||
| lc_namespace = ["langchain", "recordmanagers", "sqlite"]; | ||
|
|
||
| tableName: string; | ||
|
|
||
| db: DatabaseType; | ||
|
|
||
| namespace: string; | ||
|
|
||
| constructor(namespace: string, config: SQLiteRecordManagerOptions) { | ||
| const { localPath, connectionString, tableName } = config; | ||
| if (!connectionString && !localPath) { | ||
| throw new Error( | ||
| "One of either `localPath` or `connectionString` is required." | ||
| ); | ||
| } | ||
| if (connectionString && localPath) { | ||
| throw new Error( | ||
| "Only one of either `localPath` or `connectionString` is allowed." | ||
| ); | ||
| } | ||
| this.namespace = namespace; | ||
| this.tableName = tableName; | ||
| this.db = new Database(connectionString ?? localPath); | ||
| } | ||
|
|
||
| async createSchema(): Promise<void> { | ||
| try { | ||
| this.db.exec(` | ||
| CREATE TABLE IF NOT EXISTS "${this.tableName}" ( | ||
| uuid TEXT PRIMARY KEY DEFAULT (lower(hex(randomblob(16)))), | ||
| key TEXT NOT NULL, | ||
| namespace TEXT NOT NULL, | ||
| updated_at REAL NOT NULL, | ||
| group_id TEXT, | ||
| UNIQUE (key, namespace) | ||
| ); | ||
| CREATE INDEX IF NOT EXISTS updated_at_index ON "${this.tableName}" (updated_at); | ||
| CREATE INDEX IF NOT EXISTS key_index ON "${this.tableName}" (key); | ||
| CREATE INDEX IF NOT EXISTS namespace_index ON "${this.tableName}" (namespace); | ||
| CREATE INDEX IF NOT EXISTS group_id_index ON "${this.tableName}" (group_id);`); | ||
| } catch (error) { | ||
| console.error("Error creating schema"); | ||
| throw error; // Re-throw the error to let the caller handle it | ||
| } | ||
| } | ||
|
|
||
| async getTime(): Promise<number> { | ||
| try { | ||
| const statement: Statement<[]> = this.db.prepare( | ||
| "SELECT strftime('%s', 'now') AS epoch" | ||
| ); | ||
| const { epoch } = statement.get() as TimeRow; | ||
| return Number(epoch); | ||
| } catch (error) { | ||
| console.error("Error getting time in SQLiteRecordManager:"); | ||
| throw error; | ||
| } | ||
| } | ||
|
|
||
| async update(keys: string[], updateOptions?: UpdateOptions): Promise<void> { | ||
| if (keys.length === 0) { | ||
| return; | ||
| } | ||
|
|
||
| const updatedAt = await this.getTime(); | ||
| const { timeAtLeast, groupIds: _groupIds } = updateOptions ?? {}; | ||
|
|
||
| if (timeAtLeast && updatedAt < timeAtLeast) { | ||
| throw new Error( | ||
| `Time sync issue with database ${updatedAt} < ${timeAtLeast}` | ||
| ); | ||
| } | ||
|
|
||
| const groupIds = _groupIds ?? keys.map(() => null); | ||
|
|
||
| if (groupIds.length !== keys.length) { | ||
| throw new Error( | ||
| `Number of keys (${keys.length}) does not match number of group_ids (${groupIds.length})` | ||
| ); | ||
| } | ||
|
|
||
| const recordsToUpsert = keys.map((key, i) => [ | ||
| key, | ||
| this.namespace, | ||
| updatedAt, | ||
| groupIds[i] ?? null, // Ensure groupIds[i] is null if undefined | ||
| ]); | ||
|
|
||
| // Consider using a transaction for batch operations | ||
| const updateTransaction = this.db.transaction(() => { | ||
| for (const row of recordsToUpsert) { | ||
| this.db | ||
| .prepare( | ||
| ` | ||
| INSERT INTO "${this.tableName}" (key, namespace, updated_at, group_id) | ||
| VALUES (?, ?, ?, ?) | ||
| ON CONFLICT (key, namespace) DO UPDATE SET updated_at = excluded.updated_at` | ||
| ) | ||
| .run(...row); | ||
| } | ||
| }); | ||
| updateTransaction(); | ||
| } | ||
|
|
||
| async exists(keys: string[]): Promise<boolean[]> { | ||
| if (keys.length === 0) { | ||
| return []; | ||
| } | ||
|
|
||
| // Prepare the placeholders and the query | ||
| const placeholders = keys.map(() => `?`).join(", "); | ||
| const sql = ` | ||
| SELECT key | ||
| FROM "${this.tableName}" | ||
| WHERE namespace = ? AND key IN (${placeholders})`; | ||
|
|
||
| // Initialize an array to fill with the existence checks | ||
| const existsArray = new Array(keys.length).fill(false); | ||
|
|
||
| try { | ||
| // Execute the query | ||
| const rows = this.db | ||
| .prepare(sql) | ||
| .all(this.namespace, ...keys) as KeyRecord[]; | ||
| // Create a set of existing keys for faster lookup | ||
| const existingKeysSet = new Set(rows.map((row) => row.key)); | ||
| // Map the input keys to booleans indicating if they exist | ||
| keys.forEach((key, index) => { | ||
| existsArray[index] = existingKeysSet.has(key); | ||
| }); | ||
| return existsArray; | ||
| } catch (error) { | ||
| console.error("Error checking existence of keys"); | ||
| throw error; // Allow the caller to handle the error | ||
| } | ||
| } | ||
|
|
||
| async listKeys(options?: ListKeyOptions): Promise<string[]> { | ||
| const { before, after, limit, groupIds } = options ?? {}; | ||
| let query = `SELECT key FROM "${this.tableName}" WHERE namespace = ?`; | ||
| const values: (string | number | string[])[] = [this.namespace]; | ||
|
|
||
| if (before) { | ||
| query += ` AND updated_at < ?`; | ||
| values.push(before); | ||
| } | ||
|
|
||
| if (after) { | ||
| query += ` AND updated_at > ?`; | ||
| values.push(after); | ||
| } | ||
|
|
||
| if (limit) { | ||
| query += ` LIMIT ?`; | ||
| values.push(limit); | ||
| } | ||
|
|
||
| if (groupIds && Array.isArray(groupIds)) { | ||
| query += ` AND group_id IN (${groupIds | ||
| .filter((gid) => gid !== null) | ||
| .map(() => "?") | ||
| .join(", ")})`; | ||
| values.push(...groupIds.filter((gid): gid is string => gid !== null)); | ||
| } | ||
|
|
||
| query += ";"; | ||
|
|
||
| // Directly using try/catch with async/await for cleaner flow | ||
| try { | ||
| const result = this.db.prepare(query).all(...values) as { key: string }[]; | ||
| return result.map((row) => row.key); | ||
| } catch (error) { | ||
| console.error("Error listing keys."); | ||
| throw error; // Re-throw the error to be handled by the caller | ||
| } | ||
| } | ||
|
|
||
| async deleteKeys(keys: string[]): Promise<void> { | ||
| if (keys.length === 0) { | ||
| return; | ||
| } | ||
|
|
||
| const placeholders = keys.map(() => "?").join(", "); | ||
| const query = `DELETE FROM "${this.tableName}" WHERE namespace = ? AND key IN (${placeholders});`; | ||
| const values = [this.namespace, ...keys].map((v) => | ||
| typeof v !== "string" ? `${v}` : v | ||
| ); | ||
|
|
||
| // Directly using try/catch with async/await for cleaner flow | ||
| try { | ||
| this.db.prepare(query).run(...values); | ||
| } catch (error) { | ||
| console.error("Error deleting keys"); | ||
| throw error; // Re-throw the error to be handled by the caller | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.