genkitx-cloud-sql-pg
Version:
Genkit AI framework plugin for Cloud SQL for PostgreSQL.
278 lines • 9.81 kB
JavaScript
import {
AuthTypes,
Connector,
IpAddressTypes
} from "@google-cloud/cloud-sql-connector";
import { GoogleAuth } from "google-auth-library";
import knex from "knex";
import {
DEFAULT_INDEX_NAME_SUFFIX,
ExactNearestNeighbor
} from "./indexes.js";
import { getIAMPrincipalEmail } from "./utils";
import { IpAddressTypes as IpAddressTypes2 } from "@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 = 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 GoogleAuth({
scopes: "https://www.googleapis.com/auth/cloud-platform"
});
dbUser = await getIAMPrincipalEmail(auth);
}
}
PostgresEngine.connector = new Connector({ userAgent: USER_AGENT });
const clientOpts = await PostgresEngine.connector.getOptions({
instanceConnectionName: `${projectId}:${region}:${instance}`,
ipType,
authType: enableIAMAuth ? AuthTypes.IAM : AuthTypes.PASSWORD
});
const dbConfig = {
client: "pg",
connection: {
...clientOpts,
...password ? { password } : {},
user: dbUser,
database
}
};
const engine = knex(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 = knex(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 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 + 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 + 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 + 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 + 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);
}
}
export {
Column,
IpAddressTypes2 as IpAddressTypes,
PostgresEngine
};
//# sourceMappingURL=engine.mjs.map