genkitx-cloud-sql-pg
Version:
Genkit AI framework plugin for Cloud SQL for PostgreSQL.
307 lines • 11.7 kB
JavaScript
;
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var engine_exports = {};
__export(engine_exports, {
Column: () => Column,
IpAddressTypes: () => import_cloud_sql_connector2.IpAddressTypes,
PostgresEngine: () => PostgresEngine
});
module.exports = __toCommonJS(engine_exports);
var import_cloud_sql_connector = require("@google-cloud/cloud-sql-connector");
var import_google_auth_library = require("google-auth-library");
var import_knex = __toESM(require("knex"));
var import_indexes = require("./indexes.js");
var import_utils = require("./utils");
var import_cloud_sql_connector2 = require("@google-cloud/cloud-sql-connector");
class Column {
/** The name of the column. */
name;
/** The data type of the column (e.g., 'TEXT', 'INT', 'UUID'). */
dataType;
/** Whether the column can be nullable. */
nullable;
/**
* Creates an instance of Column.
* @param {string} name - The name of the column.
* @param {string} dataType - The data type of the column.
* @param {boolean} [nullable=true] - Whether the column can be nullable. Defaults to true.
*/
constructor(name, dataType, nullable = true) {
this.name = name;
this.dataType = dataType;
this.nullable = nullable;
this.postInitilization();
}
postInitilization() {
if (typeof this.name !== "string") {
throw "Column name must be type string";
}
if (typeof this.dataType !== "string") {
throw "Column data_type must be type string";
}
}
}
const USER_AGENT = "genkit-cloud-sql-pg-js";
class PostgresEngine {
static _createKey = Symbol();
pool;
static connector;
constructor(key, pool) {
if (key !== PostgresEngine._createKey) {
throw new Error("Only create class through 'create' method!");
}
this.pool = pool;
}
/**
* @param projectId Required - GCP Project ID
* @param region Required - Postgres Instance Region
* @param instance Required - Postgres Instance name
* @param database Required - Database name
* @param ipType Optional - IP address type. Defaults to IPAddressType.PUBLIC
* @param user Optional - Postgres user name. Defaults to undefined
* @param password Optional - Postgres user password. Defaults to undefined
* @param iamAccountEmail Optional - IAM service account email. Defaults to undefined
* @returns PostgresEngine instance
*/
static async fromInstance(projectId, region, instance, database, {
ipType = import_cloud_sql_connector.IpAddressTypes.PUBLIC,
user,
password,
iamAccountEmail
} = {}) {
let dbUser;
let enableIAMAuth;
if (!user && password || user && !password) {
throw "Only one of 'user' or 'password' were specified. Either both should be specified to use basic user/password authentication or neither for IAM DB authentication.";
}
if (user !== void 0 && password !== void 0) {
enableIAMAuth = false;
dbUser = user;
} else {
enableIAMAuth = true;
if (iamAccountEmail !== void 0) {
dbUser = iamAccountEmail;
} else {
const auth = new import_google_auth_library.GoogleAuth({
scopes: "https://www.googleapis.com/auth/cloud-platform"
});
dbUser = await (0, import_utils.getIAMPrincipalEmail)(auth);
}
}
PostgresEngine.connector = new import_cloud_sql_connector.Connector({ userAgent: USER_AGENT });
const clientOpts = await PostgresEngine.connector.getOptions({
instanceConnectionName: `${projectId}:${region}:${instance}`,
ipType,
authType: enableIAMAuth ? import_cloud_sql_connector.AuthTypes.IAM : import_cloud_sql_connector.AuthTypes.PASSWORD
});
const dbConfig = {
client: "pg",
connection: {
...clientOpts,
...password ? { password } : {},
user: dbUser,
database
}
};
const engine = (0, import_knex.default)(dbConfig);
return new PostgresEngine(PostgresEngine._createKey, engine);
}
/**
* Create a PostgresEngine instance from an Knex instance.
*
* @param engine knex instance
* @returns PostgresEngine instance from a knex instance
*/
static async fromEngine(engine) {
return new PostgresEngine(PostgresEngine._createKey, engine);
}
/**
* Create a PostgresEngine instance from arguments.
*
* @param url URL use to connect to a database
* @param poolConfig Optional - Configuration pool to use in the Knex configuration
* @returns PostgresEngine instance
*/
static async fromEngineArgs(url, poolConfig) {
const driver = "postgresql+asyncpg";
if (typeof url === "string" && !url.startsWith(driver)) {
throw "Driver must be type 'postgresql+asyncpg'";
}
const dbConfig = {
client: "pg",
connection: url,
acquireConnectionTimeout: 1e6,
pool: {
...poolConfig,
acquireTimeoutMillis: 6e5
}
};
const engine = (0, import_knex.default)(dbConfig);
return new PostgresEngine(PostgresEngine._createKey, engine);
}
/**
* Create a table for saving of vectors to be used with PostgresVectorStore.
*
* @param tableName Postgres database table name
* @param vectorSize Vector size for the embedding model to be used.
* @param schemaName The schema name to store Postgres database table. Default: "public".
* @param contentColumn Name of the column to store document content. Default: "content".
* @param embeddingColumn Name of the column to store vector embeddings. Default: "embedding".
* @param metadataColumns Optional - A list of Columns to create for custom metadata. Default: [].
* @param metadataJsonColumn Optional - The column to store extra metadata in JSON format. Default: "json_metadata".
* @param idColumn Optional - Column to store ids. Default: "id" column name with data type UUID.
* @param overwriteExisting Whether to drop existing table. Default: False.
* @param storeMetadata Whether to store metadata in the table. Default: True.
*/
async initVectorstoreTable(tableName, vectorSize, {
schemaName = "public",
contentColumn = "content",
embeddingColumn = "embedding",
metadataColumns = [],
metadataJsonColumn = "json_metadata",
idColumn = "id",
overwriteExisting = false,
storeMetadata = true
} = {}) {
await this.pool.raw("CREATE EXTENSION IF NOT EXISTS vector");
if (overwriteExisting) {
await this.pool.schema.withSchema(schemaName).dropTableIfExists(tableName);
}
const idDataType = typeof idColumn === "string" ? "UUID" : idColumn.dataType;
const idColumnName = typeof idColumn === "string" ? idColumn : idColumn.name;
let query = `CREATE TABLE "${schemaName}"."${tableName}"(
${idColumnName} ${idDataType} PRIMARY KEY,
${contentColumn} TEXT NOT NULL,
${embeddingColumn} vector(${vectorSize}) NOT NULL`;
for (const column of metadataColumns) {
const nullable = !column.nullable ? "NOT NULL" : "";
query += `,
${column.name} ${column.dataType} ${nullable}`;
}
if (storeMetadata) {
query += `,
${metadataJsonColumn} JSON`;
}
query += `
);`;
await this.pool.raw(query);
}
/**
* Dispose of connection pool
*/
async closeConnection() {
await this.pool.destroy();
if (PostgresEngine.connector !== void 0) {
PostgresEngine.connector.close();
}
}
/**
* Just to test the connection to the database.
* @returns A Promise that resolves to a row containing the current timestamp.
*/
testConnection() {
return this.pool.raw("SELECT NOW() as currentTimestamp").then((result) => result.entries[0].currentTimestamp);
}
/**
* Create an index on the vector store table
* @param {string} tableName
* @param {BaseIndex} index
* @param {VectorStoreTableArgs}
*/
async applyVectorIndex(tableName, index, {
schemaName = "public",
embeddingColumn = "embedding",
concurrently = false
} = {}) {
if (index instanceof import_indexes.ExactNearestNeighbor) {
await this.dropVectorIndex({ tableName });
return;
}
const filter = index.partialIndexes ? `WHERE (${index.partialIndexes})` : "";
const indexOptions = `WITH ${index.indexOptions()}`;
const funct = index.distanceStrategy.indexFunction;
const name = index.name ? index.name : tableName + import_indexes.DEFAULT_INDEX_NAME_SUFFIX;
const stmt = `CREATE INDEX ${concurrently ? "CONCURRENTLY" : ""} ${name} ON "${schemaName}"."${tableName}" USING ${index.indexType} (${embeddingColumn} ${funct}) ${indexOptions} ${filter};`;
await this.pool.raw(stmt);
}
/**
* Check if index exists in the table.
* @param {string} tableName
* @param VectorStoreTableArgs Optional
*/
async isValidIndex(tableName, { indexName, schemaName = "public" } = {}) {
const idxName = indexName || tableName + import_indexes.DEFAULT_INDEX_NAME_SUFFIX;
const stmt = `SELECT tablename, indexname
FROM pg_indexes
WHERE tablename = '${tableName}' AND schemaname = '${schemaName}' AND indexname = '${idxName}';`;
const { rows } = await this.pool.raw(stmt);
return rows.length === 1;
}
/**
* Drop the vector index
* @param {string} tableName
* @param {string} indexName Optional - index name
*/
async dropVectorIndex(params) {
let idxName = "";
if (params.indexName) {
idxName = params.indexName;
} else if (params.tableName) {
idxName = params.tableName + import_indexes.DEFAULT_INDEX_NAME_SUFFIX;
} else {
throw new Error("Index name or Table name are not provided.");
}
const query = `DROP INDEX IF EXISTS ${idxName};`;
await this.pool.raw(query);
}
/**
* Re-index the vector store table
* @param {string} tableName
* @param {string} indexName Optional - index name
*/
async reIndex(params) {
let idxName = "";
if (params.indexName) {
idxName = params.indexName;
} else if (params.tableName) {
idxName = params.tableName + import_indexes.DEFAULT_INDEX_NAME_SUFFIX;
} else {
throw new Error("Index name or Table name are not provided.");
}
const query = `REINDEX INDEX ${idxName};`;
this.pool.raw(query);
}
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
Column,
IpAddressTypes,
PostgresEngine
});
//# sourceMappingURL=engine.js.map