186 lines
6.9 KiB
JavaScript
186 lines
6.9 KiB
JavaScript
|
"use strict";
|
||
|
var __importDefault = (this && this.__importDefault) || function (mod) {
|
||
|
return (mod && mod.__esModule) ? mod : { "default": mod };
|
||
|
};
|
||
|
Object.defineProperty(exports, "__esModule", { value: true });
|
||
|
exports.PostgresRecordManager = void 0;
|
||
|
const pg_1 = __importDefault(require("pg"));
|
||
|
class PostgresRecordManager {
|
||
|
constructor(namespace, config) {
|
||
|
Object.defineProperty(this, "lc_namespace", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: ["langchain", "recordmanagers", "postgres"]
|
||
|
});
|
||
|
Object.defineProperty(this, "pool", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "tableName", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "namespace", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "finalTableName", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
const { postgresConnectionOptions, tableName, pool } = config;
|
||
|
this.namespace = namespace;
|
||
|
this.pool = pool || new pg_1.default.Pool(postgresConnectionOptions);
|
||
|
this.tableName = tableName || "upsertion_records";
|
||
|
this.finalTableName = config.schema
|
||
|
? `"${config.schema}"."${this.tableName}"`
|
||
|
: `"${this.tableName}"`;
|
||
|
}
|
||
|
async createSchema() {
|
||
|
try {
|
||
|
await this.pool.query(`
|
||
|
CREATE TABLE IF NOT EXISTS ${this.finalTableName} (
|
||
|
uuid UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||
|
key TEXT NOT NULL,
|
||
|
namespace TEXT NOT NULL,
|
||
|
updated_at Double PRECISION NOT NULL,
|
||
|
group_id TEXT,
|
||
|
UNIQUE (key, namespace)
|
||
|
);
|
||
|
CREATE INDEX IF NOT EXISTS updated_at_index ON ${this.finalTableName} (updated_at);
|
||
|
CREATE INDEX IF NOT EXISTS key_index ON ${this.finalTableName} (key);
|
||
|
CREATE INDEX IF NOT EXISTS namespace_index ON ${this.finalTableName} (namespace);
|
||
|
CREATE INDEX IF NOT EXISTS group_id_index ON ${this.finalTableName} (group_id);`);
|
||
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||
|
}
|
||
|
catch (e) {
|
||
|
// This error indicates that the table already exists
|
||
|
// Due to asynchronous nature of the code, it is possible that
|
||
|
// the table is created between the time we check if it exists
|
||
|
// and the time we try to create it. It can be safely ignored.
|
||
|
if ("code" in e && e.code === "23505") {
|
||
|
return;
|
||
|
}
|
||
|
throw e;
|
||
|
}
|
||
|
}
|
||
|
async getTime() {
|
||
|
const res = await this.pool.query("SELECT EXTRACT(EPOCH FROM CURRENT_TIMESTAMP)");
|
||
|
return Number.parseFloat(res.rows[0].extract);
|
||
|
}
|
||
|
/**
|
||
|
* Generates the SQL placeholders for a specific row at the provided index.
|
||
|
*
|
||
|
* @param index - The index of the row for which placeholders need to be generated.
|
||
|
* @param numOfColumns - The number of columns we are inserting data into.
|
||
|
* @returns The SQL placeholders for the row values.
|
||
|
*/
|
||
|
generatePlaceholderForRowAt(index, numOfColumns) {
|
||
|
const placeholders = [];
|
||
|
for (let i = 0; i < numOfColumns; i += 1) {
|
||
|
placeholders.push(`$${index * numOfColumns + i + 1}`);
|
||
|
}
|
||
|
return `(${placeholders.join(", ")})`;
|
||
|
}
|
||
|
async update(keys, updateOptions) {
|
||
|
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],
|
||
|
]);
|
||
|
const valuesPlaceholders = recordsToUpsert
|
||
|
.map((_, j) => this.generatePlaceholderForRowAt(j, recordsToUpsert[0].length))
|
||
|
.join(", ");
|
||
|
const query = `INSERT INTO ${this.finalTableName} (key, namespace, updated_at, group_id) VALUES ${valuesPlaceholders} ON CONFLICT (key, namespace) DO UPDATE SET updated_at = EXCLUDED.updated_at;`;
|
||
|
await this.pool.query(query, recordsToUpsert.flat());
|
||
|
}
|
||
|
async exists(keys) {
|
||
|
if (keys.length === 0) {
|
||
|
return [];
|
||
|
}
|
||
|
const startIndex = 2;
|
||
|
const arrayPlaceholders = keys
|
||
|
.map((_, i) => `$${i + startIndex}`)
|
||
|
.join(", ");
|
||
|
const query = `
|
||
|
WITH ordered_keys AS (
|
||
|
SELECT * FROM unnest(ARRAY[${arrayPlaceholders}]) WITH ORDINALITY as t(key, o)
|
||
|
)
|
||
|
SELECT ok.key, (r.key IS NOT NULL) ex
|
||
|
FROM ordered_keys ok
|
||
|
LEFT JOIN ${this.finalTableName} r
|
||
|
ON r.key = ok.key
|
||
|
AND namespace = $1
|
||
|
ORDER BY ok.o;
|
||
|
`;
|
||
|
const res = await this.pool.query(query, [this.namespace, ...keys.flat()]);
|
||
|
return res.rows.map((row) => row.ex);
|
||
|
}
|
||
|
async listKeys(options) {
|
||
|
const { before, after, limit, groupIds } = options ?? {};
|
||
|
let query = `SELECT key FROM ${this.finalTableName} WHERE namespace = $1`;
|
||
|
const values = [this.namespace];
|
||
|
let index = 2;
|
||
|
if (before) {
|
||
|
values.push(before);
|
||
|
query += ` AND updated_at < $${index}`;
|
||
|
index += 1;
|
||
|
}
|
||
|
if (after) {
|
||
|
values.push(after);
|
||
|
query += ` AND updated_at > $${index}`;
|
||
|
index += 1;
|
||
|
}
|
||
|
if (limit) {
|
||
|
values.push(limit);
|
||
|
query += ` LIMIT $${index}`;
|
||
|
index += 1;
|
||
|
}
|
||
|
if (groupIds) {
|
||
|
values.push(groupIds);
|
||
|
query += ` AND group_id = ANY($${index})`;
|
||
|
index += 1;
|
||
|
}
|
||
|
query += ";";
|
||
|
const res = await this.pool.query(query, values);
|
||
|
return res.rows.map((row) => row.key);
|
||
|
}
|
||
|
async deleteKeys(keys) {
|
||
|
if (keys.length === 0) {
|
||
|
return;
|
||
|
}
|
||
|
const query = `DELETE FROM ${this.finalTableName} WHERE namespace = $1 AND key = ANY($2);`;
|
||
|
await this.pool.query(query, [this.namespace, keys]);
|
||
|
}
|
||
|
/**
|
||
|
* Terminates the connection pool.
|
||
|
* @returns {Promise<void>}
|
||
|
*/
|
||
|
async end() {
|
||
|
await this.pool.end();
|
||
|
}
|
||
|
}
|
||
|
exports.PostgresRecordManager = PostgresRecordManager;
|