UNPKG

genkitx-cloud-sql-pg

Version:

Genkit AI framework plugin for Cloud SQL for PostgreSQL.

1 lines 18.6 kB
{"version":3,"sources":["../src/engine.ts"],"sourcesContent":["/**\n * Copyright 2025 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\nimport {\n AuthTypes,\n Connector,\n IpAddressTypes,\n} from '@google-cloud/cloud-sql-connector';\nimport { GoogleAuth } from 'google-auth-library';\nimport knex from 'knex';\nimport {\n BaseIndex,\n DEFAULT_INDEX_NAME_SUFFIX,\n ExactNearestNeighbor,\n} from './indexes.js';\n\nimport { getIAMPrincipalEmail } from './utils';\n\nexport { IpAddressTypes } from '@google-cloud/cloud-sql-connector';\n\n/**\n * Defines the arguments for configuring a PostgreSQL engine.\n */\nexport interface PostgresEngineArgs {\n /** The IP address type to use for the connection (e.g., PUBLIC, PRIVATE). */\n ipType?: IpAddressTypes;\n /** The PostgreSQL username for basic authentication. */\n user?: string;\n /** The PostgreSQL password for basic authentication. */\n password?: string;\n /** The IAM service account email for IAM database authentication. */\n iamAccountEmail?: string;\n}\n\n/**\n * Defines the arguments for configuring a vector store table.\n */\nexport interface VectorStoreTableArgs {\n /** The schema name for the table. Defaults to \"public\". */\n schemaName?: string;\n /** The name of the column to store document content. Defaults to \"content\". */\n contentColumn?: string;\n /** The name of the column to store vector embeddings. Defaults to \"embedding\". */\n embeddingColumn?: string;\n /** An optional list of `Column` objects to create for custom metadata. Defaults to []. */\n metadataColumns?: Column[];\n /** The column to store extra metadata in JSON format. Defaults to \"json_metadata\". */\n metadataJsonColumn?: string;\n /**\n * The column to store IDs. Can be a string (column name) or a `Column` object\n * for more detailed configuration. Defaults to \"id\" with data type UUID.\n */\n idColumn?: string | Column;\n /** Whether to drop the existing table if it already exists. Defaults to false. */\n overwriteExisting?: boolean;\n /** Whether to store metadata in the table. Defaults to true. */\n storeMetadata?: boolean;\n /**\n * Whether to build the index concurrently (allowing concurrent operations on the table).\n * Defaults to false.\n */\n concurrently?: boolean;\n /** The name of the index. If not provided, a default name will be generated. */\n indexName?: string;\n}\n\n/**\n * Represents a database table column.\n */\nexport class Column {\n /** The name of the column. */\n name: string;\n /** The data type of the column (e.g., 'TEXT', 'INT', 'UUID'). */\n dataType: string;\n /** Whether the column can be nullable. */\n nullable: boolean;\n\n /**\n * Creates an instance of Column.\n * @param {string} name - The name of the column.\n * @param {string} dataType - The data type of the column.\n * @param {boolean} [nullable=true] - Whether the column can be nullable. Defaults to true.\n */\n constructor(name: string, dataType: string, nullable: boolean = true) {\n this.name = name;\n this.dataType = dataType;\n this.nullable = nullable;\n\n this.postInitilization();\n }\n\n private postInitilization() {\n if (typeof this.name !== 'string') {\n throw 'Column name must be type string';\n }\n\n if (typeof this.dataType !== 'string') {\n throw 'Column data_type must be type string';\n }\n }\n}\n\nconst USER_AGENT = 'genkit-cloud-sql-pg-js';\n\n/**\n * Manages connections and operations for a PostgreSQL database,\n * particularly for vector store functionalities.\n */\nexport class PostgresEngine {\n private static _createKey = Symbol();\n pool: knex.Knex<any, any[]>;\n static connector: Connector;\n\n constructor(key: Symbol, pool: knex.Knex<any, any[]>) {\n if (key !== PostgresEngine._createKey) {\n throw new Error(\"Only create class through 'create' method!\");\n }\n this.pool = pool;\n }\n\n /**\n * @param projectId Required - GCP Project ID\n * @param region Required - Postgres Instance Region\n * @param instance Required - Postgres Instance name\n * @param database Required - Database name\n * @param ipType Optional - IP address type. Defaults to IPAddressType.PUBLIC\n * @param user Optional - Postgres user name. Defaults to undefined\n * @param password Optional - Postgres user password. Defaults to undefined\n * @param iamAccountEmail Optional - IAM service account email. Defaults to undefined\n * @returns PostgresEngine instance\n */\n\n static async fromInstance(\n projectId: string,\n region: string,\n instance: string,\n database: string,\n {\n ipType = IpAddressTypes.PUBLIC,\n user,\n password,\n iamAccountEmail,\n }: PostgresEngineArgs = {}\n ): Promise<PostgresEngine> {\n let dbUser: string;\n let enableIAMAuth: boolean;\n\n if ((!user && password) || (user && !password)) {\n // XOR for strings\n throw (\n \"Only one of 'user' or 'password' were specified. Either \" +\n 'both should be specified to use basic user/password ' +\n 'authentication or neither for IAM DB authentication.'\n );\n }\n\n // User and password are given so we use the basic auth\n if (user !== undefined && password !== undefined) {\n enableIAMAuth = false;\n dbUser = user!;\n } else {\n enableIAMAuth = true;\n if (iamAccountEmail !== undefined) {\n dbUser = iamAccountEmail;\n } else {\n // Get application default credentials\n const auth = new GoogleAuth({\n scopes: 'https://www.googleapis.com/auth/cloud-platform',\n });\n // dbUser should be the iam principal email by passing the credentials obtained\n dbUser = await getIAMPrincipalEmail(auth);\n }\n }\n\n PostgresEngine.connector = new Connector({ userAgent: USER_AGENT });\n const clientOpts = await PostgresEngine.connector.getOptions({\n instanceConnectionName: `${projectId}:${region}:${instance}`,\n ipType: ipType,\n authType: enableIAMAuth ? AuthTypes.IAM : AuthTypes.PASSWORD,\n });\n\n const dbConfig: knex.Knex.Config<any> = {\n client: 'pg',\n connection: {\n ...clientOpts,\n ...(password ? { password: password } : {}),\n user: dbUser,\n database: database,\n },\n };\n\n const engine = knex(dbConfig);\n\n return new PostgresEngine(PostgresEngine._createKey, engine);\n }\n\n /**\n * Create a PostgresEngine instance from an Knex instance.\n *\n * @param engine knex instance\n * @returns PostgresEngine instance from a knex instance\n */\n static async fromEngine(engine: knex.Knex<any, any[]>) {\n return new PostgresEngine(PostgresEngine._createKey, engine);\n }\n\n /**\n * Create a PostgresEngine instance from arguments.\n *\n * @param url URL use to connect to a database\n * @param poolConfig Optional - Configuration pool to use in the Knex configuration\n * @returns PostgresEngine instance\n */\n static async fromEngineArgs(\n url: string | knex.Knex.StaticConnectionConfig,\n poolConfig?: knex.Knex.PoolConfig\n ) {\n const driver = 'postgresql+asyncpg';\n\n if (typeof url === 'string' && !url.startsWith(driver)) {\n throw \"Driver must be type 'postgresql+asyncpg'\";\n }\n\n const dbConfig: knex.Knex.Config<any> = {\n client: 'pg',\n connection: url,\n acquireConnectionTimeout: 1000000,\n pool: {\n ...poolConfig,\n acquireTimeoutMillis: 600000,\n },\n };\n\n const engine = knex(dbConfig);\n\n return new PostgresEngine(PostgresEngine._createKey, engine);\n }\n\n /**\n * Create a table for saving of vectors to be used with PostgresVectorStore.\n *\n * @param tableName Postgres database table name\n * @param vectorSize Vector size for the embedding model to be used.\n * @param schemaName The schema name to store Postgres database table. Default: \"public\".\n * @param contentColumn Name of the column to store document content. Default: \"content\".\n * @param embeddingColumn Name of the column to store vector embeddings. Default: \"embedding\".\n * @param metadataColumns Optional - A list of Columns to create for custom metadata. Default: [].\n * @param metadataJsonColumn Optional - The column to store extra metadata in JSON format. Default: \"json_metadata\".\n * @param idColumn Optional - Column to store ids. Default: \"id\" column name with data type UUID.\n * @param overwriteExisting Whether to drop existing table. Default: False.\n * @param storeMetadata Whether to store metadata in the table. Default: True.\n */\n async initVectorstoreTable(\n tableName: string,\n vectorSize: number,\n {\n schemaName = 'public',\n contentColumn = 'content',\n embeddingColumn = 'embedding',\n metadataColumns = [],\n metadataJsonColumn = 'json_metadata',\n idColumn = 'id',\n overwriteExisting = false,\n storeMetadata = true,\n }: VectorStoreTableArgs = {}\n ): Promise<void> {\n await this.pool.raw('CREATE EXTENSION IF NOT EXISTS vector');\n\n if (overwriteExisting) {\n await this.pool.schema\n .withSchema(schemaName)\n .dropTableIfExists(tableName);\n }\n\n const idDataType =\n typeof idColumn === 'string' ? 'UUID' : idColumn.dataType;\n const idColumnName =\n typeof idColumn === 'string' ? idColumn : idColumn.name;\n\n let query = `CREATE TABLE \"${schemaName}\".\"${tableName}\"(\n ${idColumnName} ${idDataType} PRIMARY KEY,\n ${contentColumn} TEXT NOT NULL,\n ${embeddingColumn} vector(${vectorSize}) NOT NULL`;\n\n for (const column of metadataColumns) {\n const nullable = !column.nullable ? 'NOT NULL' : '';\n query += `,\\n ${column.name} ${column.dataType} ${nullable}`;\n }\n\n if (storeMetadata) {\n query += `,\\n${metadataJsonColumn} JSON`;\n }\n\n query += `\\n);`;\n\n await this.pool.raw(query);\n }\n\n /**\n * Dispose of connection pool\n */\n async closeConnection(): Promise<void> {\n await this.pool.destroy();\n if (PostgresEngine.connector !== undefined) {\n PostgresEngine.connector.close();\n }\n }\n\n /**\n * Just to test the connection to the database.\n * @returns A Promise that resolves to a row containing the current timestamp.\n */\n testConnection(): Promise<{ currentTimestamp: Date }[]> {\n return this.pool\n .raw<{ currentTimestamp: Date }[]>('SELECT NOW() as currentTimestamp')\n .then((result) => result.entries[0].currentTimestamp);\n }\n\n /**\n * Create an index on the vector store table\n * @param {string} tableName\n * @param {BaseIndex} index\n * @param {VectorStoreTableArgs}\n */\n async applyVectorIndex(\n tableName: string,\n index: BaseIndex,\n {\n schemaName = 'public',\n embeddingColumn = 'embedding',\n concurrently = false,\n }: VectorStoreTableArgs = {}\n ): Promise<void> {\n if (index instanceof ExactNearestNeighbor) {\n await this.dropVectorIndex({ tableName: tableName });\n return;\n }\n\n const filter = index.partialIndexes\n ? `WHERE (${index.partialIndexes})`\n : '';\n const indexOptions = `WITH ${index.indexOptions()}`;\n const funct = index.distanceStrategy.indexFunction;\n\n const name = index.name\n ? index.name\n : tableName + DEFAULT_INDEX_NAME_SUFFIX;\n\n const stmt = `CREATE INDEX ${concurrently ? 'CONCURRENTLY' : ''} ${name} ON \"${schemaName}\".\"${tableName}\" USING ${index.indexType} (${embeddingColumn} ${funct}) ${indexOptions} ${filter};`;\n\n await this.pool.raw(stmt);\n }\n\n /**\n * Check if index exists in the table.\n * @param {string} tableName\n * @param VectorStoreTableArgs Optional\n */\n async isValidIndex(\n tableName: string,\n { indexName, schemaName = 'public' }: VectorStoreTableArgs = {}\n ): Promise<boolean> {\n const idxName = indexName || tableName + DEFAULT_INDEX_NAME_SUFFIX;\n const stmt = `SELECT tablename, indexname\n FROM pg_indexes\n WHERE tablename = '${tableName}' AND schemaname = '${schemaName}' AND indexname = '${idxName}';`;\n const { rows } = await this.pool.raw(stmt);\n\n return rows.length === 1;\n }\n\n /**\n * Drop the vector index\n * @param {string} tableName\n * @param {string} indexName Optional - index name\n */\n async dropVectorIndex(params: {\n tableName?: string;\n indexName?: string;\n }): Promise<void> {\n let idxName = '';\n if (params.indexName) {\n idxName = params.indexName;\n } else if (params.tableName) {\n idxName = params.tableName + DEFAULT_INDEX_NAME_SUFFIX;\n } else {\n throw new Error('Index name or Table name are not provided.');\n }\n const query = `DROP INDEX IF EXISTS ${idxName};`;\n await this.pool.raw(query);\n }\n\n /**\n * Re-index the vector store table\n * @param {string} tableName\n * @param {string} indexName Optional - index name\n */\n async reIndex(params: { tableName?: string; indexName?: string }) {\n let idxName = '';\n if (params.indexName) {\n idxName = params.indexName;\n } else if (params.tableName) {\n idxName = params.tableName + DEFAULT_INDEX_NAME_SUFFIX;\n } else {\n throw new Error('Index name or Table name are not provided.');\n }\n const query = `REINDEX INDEX ${idxName};`;\n this.pool.raw(query);\n }\n}\n"],"mappings":"AAeA;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,OACK;AACP,SAAS,kBAAkB;AAC3B,OAAO,UAAU;AACjB;AAAA,EAEE;AAAA,EACA;AAAA,OACK;AAEP,SAAS,4BAA4B;AAErC,SAAS,kBAAAA,uBAAsB;AAmDxB,MAAM,OAAO;AAAA;AAAA,EAElB;AAAA;AAAA,EAEA;AAAA;AAAA,EAEA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAQA,YAAY,MAAc,UAAkB,WAAoB,MAAM;AACpE,SAAK,OAAO;AACZ,SAAK,WAAW;AAChB,SAAK,WAAW;AAEhB,SAAK,kBAAkB;AAAA,EACzB;AAAA,EAEQ,oBAAoB;AAC1B,QAAI,OAAO,KAAK,SAAS,UAAU;AACjC,YAAM;AAAA,IACR;AAEA,QAAI,OAAO,KAAK,aAAa,UAAU;AACrC,YAAM;AAAA,IACR;AAAA,EACF;AACF;AAEA,MAAM,aAAa;AAMZ,MAAM,eAAe;AAAA,EAC1B,OAAe,aAAa,OAAO;AAAA,EACnC;AAAA,EACA,OAAO;AAAA,EAEP,YAAY,KAAa,MAA6B;AACpD,QAAI,QAAQ,eAAe,YAAY;AACrC,YAAM,IAAI,MAAM,4CAA4C;AAAA,IAC9D;AACA,SAAK,OAAO;AAAA,EACd;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAcA,aAAa,aACX,WACA,QACA,UACA,UACA;AAAA,IACE,SAAS,eAAe;AAAA,IACxB;AAAA,IACA;AAAA,IACA;AAAA,EACF,IAAwB,CAAC,GACA;AACzB,QAAI;AACJ,QAAI;AAEJ,QAAK,CAAC,QAAQ,YAAc,QAAQ,CAAC,UAAW;AAE9C,YACE;AAAA,IAIJ;AAGA,QAAI,SAAS,UAAa,aAAa,QAAW;AAChD,sBAAgB;AAChB,eAAS;AAAA,IACX,OAAO;AACL,sBAAgB;AAChB,UAAI,oBAAoB,QAAW;AACjC,iBAAS;AAAA,MACX,OAAO;AAEL,cAAM,OAAO,IAAI,WAAW;AAAA,UAC1B,QAAQ;AAAA,QACV,CAAC;AAED,iBAAS,MAAM,qBAAqB,IAAI;AAAA,MAC1C;AAAA,IACF;AAEA,mBAAe,YAAY,IAAI,UAAU,EAAE,WAAW,WAAW,CAAC;AAClE,UAAM,aAAa,MAAM,eAAe,UAAU,WAAW;AAAA,MAC3D,wBAAwB,GAAG,SAAS,IAAI,MAAM,IAAI,QAAQ;AAAA,MAC1D;AAAA,MACA,UAAU,gBAAgB,UAAU,MAAM,UAAU;AAAA,IACtD,CAAC;AAED,UAAM,WAAkC;AAAA,MACtC,QAAQ;AAAA,MACR,YAAY;AAAA,QACV,GAAG;AAAA,QACH,GAAI,WAAW,EAAE,SAAmB,IAAI,CAAC;AAAA,QACzC,MAAM;AAAA,QACN;AAAA,MACF;AAAA,IACF;AAEA,UAAM,SAAS,KAAK,QAAQ;AAE5B,WAAO,IAAI,eAAe,eAAe,YAAY,MAAM;AAAA,EAC7D;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAQA,aAAa,WAAW,QAA+B;AACrD,WAAO,IAAI,eAAe,eAAe,YAAY,MAAM;AAAA,EAC7D;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EASA,aAAa,eACX,KACA,YACA;AACA,UAAM,SAAS;AAEf,QAAI,OAAO,QAAQ,YAAY,CAAC,IAAI,WAAW,MAAM,GAAG;AACtD,YAAM;AAAA,IACR;AAEA,UAAM,WAAkC;AAAA,MACtC,QAAQ;AAAA,MACR,YAAY;AAAA,MACZ,0BAA0B;AAAA,MAC1B,MAAM;AAAA,QACJ,GAAG;AAAA,QACH,sBAAsB;AAAA,MACxB;AAAA,IACF;AAEA,UAAM,SAAS,KAAK,QAAQ;AAE5B,WAAO,IAAI,eAAe,eAAe,YAAY,MAAM;AAAA,EAC7D;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAgBA,MAAM,qBACJ,WACA,YACA;AAAA,IACE,aAAa;AAAA,IACb,gBAAgB;AAAA,IAChB,kBAAkB;AAAA,IAClB,kBAAkB,CAAC;AAAA,IACnB,qBAAqB;AAAA,IACrB,WAAW;AAAA,IACX,oBAAoB;AAAA,IACpB,gBAAgB;AAAA,EAClB,IAA0B,CAAC,GACZ;AACf,UAAM,KAAK,KAAK,IAAI,uCAAuC;AAE3D,QAAI,mBAAmB;AACrB,YAAM,KAAK,KAAK,OACb,WAAW,UAAU,EACrB,kBAAkB,SAAS;AAAA,IAChC;AAEA,UAAM,aACJ,OAAO,aAAa,WAAW,SAAS,SAAS;AACnD,UAAM,eACJ,OAAO,aAAa,WAAW,WAAW,SAAS;AAErD,QAAI,QAAQ,iBAAiB,UAAU,MAAM,SAAS;AAAA,QAClD,YAAY,IAAI,UAAU;AAAA,QAC1B,aAAa;AAAA,QACb,eAAe,WAAW,UAAU;AAExC,eAAW,UAAU,iBAAiB;AACpC,YAAM,WAAW,CAAC,OAAO,WAAW,aAAa;AACjD,eAAS;AAAA,GAAO,OAAO,IAAI,IAAI,OAAO,QAAQ,IAAI,QAAQ;AAAA,IAC5D;AAEA,QAAI,eAAe;AACjB,eAAS;AAAA,EAAM,kBAAkB;AAAA,IACnC;AAEA,aAAS;AAAA;AAET,UAAM,KAAK,KAAK,IAAI,KAAK;AAAA,EAC3B;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,kBAAiC;AACrC,UAAM,KAAK,KAAK,QAAQ;AACxB,QAAI,eAAe,cAAc,QAAW;AAC1C,qBAAe,UAAU,MAAM;AAAA,IACjC;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,iBAAwD;AACtD,WAAO,KAAK,KACT,IAAkC,kCAAkC,EACpE,KAAK,CAAC,WAAW,OAAO,QAAQ,CAAC,EAAE,gBAAgB;AAAA,EACxD;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAQA,MAAM,iBACJ,WACA,OACA;AAAA,IACE,aAAa;AAAA,IACb,kBAAkB;AAAA,IAClB,eAAe;AAAA,EACjB,IAA0B,CAAC,GACZ;AACf,QAAI,iBAAiB,sBAAsB;AACzC,YAAM,KAAK,gBAAgB,EAAE,UAAqB,CAAC;AACnD;AAAA,IACF;AAEA,UAAM,SAAS,MAAM,iBACjB,UAAU,MAAM,cAAc,MAC9B;AACJ,UAAM,eAAe,QAAQ,MAAM,aAAa,CAAC;AACjD,UAAM,QAAQ,MAAM,iBAAiB;AAErC,UAAM,OAAO,MAAM,OACf,MAAM,OACN,YAAY;AAEhB,UAAM,OAAO,gBAAgB,eAAe,iBAAiB,EAAE,IAAI,IAAI,QAAQ,UAAU,MAAM,SAAS,WAAW,MAAM,SAAS,KAAK,eAAe,IAAI,KAAK,KAAK,YAAY,IAAI,MAAM;AAE1L,UAAM,KAAK,KAAK,IAAI,IAAI;AAAA,EAC1B;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,MAAM,aACJ,WACA,EAAE,WAAW,aAAa,SAAS,IAA0B,CAAC,GAC5C;AAClB,UAAM,UAAU,aAAa,YAAY;AACzC,UAAM,OAAO;AAAA;AAAA,+CAE8B,SAAS,uBAAuB,UAAU,sBAAsB,OAAO;AAClH,UAAM,EAAE,KAAK,IAAI,MAAM,KAAK,KAAK,IAAI,IAAI;AAEzC,WAAO,KAAK,WAAW;AAAA,EACzB;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,MAAM,gBAAgB,QAGJ;AAChB,QAAI,UAAU;AACd,QAAI,OAAO,WAAW;AACpB,gBAAU,OAAO;AAAA,IACnB,WAAW,OAAO,WAAW;AAC3B,gBAAU,OAAO,YAAY;AAAA,IAC/B,OAAO;AACL,YAAM,IAAI,MAAM,4CAA4C;AAAA,IAC9D;AACA,UAAM,QAAQ,wBAAwB,OAAO;AAC7C,UAAM,KAAK,KAAK,IAAI,KAAK;AAAA,EAC3B;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,MAAM,QAAQ,QAAoD;AAChE,QAAI,UAAU;AACd,QAAI,OAAO,WAAW;AACpB,gBAAU,OAAO;AAAA,IACnB,WAAW,OAAO,WAAW;AAC3B,gBAAU,OAAO,YAAY;AAAA,IAC/B,OAAO;AACL,YAAM,IAAI,MAAM,4CAA4C;AAAA,IAC9D;AACA,UAAM,QAAQ,iBAAiB,OAAO;AACtC,SAAK,KAAK,IAAI,KAAK;AAAA,EACrB;AACF;","names":["IpAddressTypes"]}