engine.test.ts•9.06 kB
/**
* Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import {
AuthTypes,
Connector,
IpAddressTypes,
} from '@google-cloud/cloud-sql-connector';
import { afterAll, beforeAll, describe, expect, test } from '@jest/globals';
import * as dotenv from 'dotenv';
import knex from 'knex';
import {
Column,
PostgresEngine,
PostgresEngineArgs,
VectorStoreTableArgs,
} from '../src/engine';
dotenv.config();
const USER_AGENT = 'genkit-google-cloud-sql-pg-js';
const CUSTOM_TABLE = 'test_table_custom';
const VECTOR_SIZE = 768;
const ID_COLUMN = 'uuid';
const CONTENT_COLUMN = 'my_content';
const EMBEDDING_COLUMN = 'my_embedding';
const METADATA_COLUMNS = [
new Column('page', 'TEXT'),
new Column('source', 'TEXT'),
];
const STORE_METADATA = true;
const REQUIRED_ENV_VARS = [
'PROJECT_ID',
'REGION',
'INSTANCE_ID',
'DATABASE_ID',
'DB_USER',
'DB_PASSWORD',
'IP_ADDRESS',
'EMAIL', // Add EMAIL here if it's consistently required for some tests
];
let envVarsFound = true;
let missingEnvVars: string[] = [];
function validateEnvVars() {
missingEnvVars = REQUIRED_ENV_VARS.filter((varName) => !process.env[varName]);
if (missingEnvVars.length > 0) {
envVarsFound = false;
console.warn(
`Skipping tests due to missing environment variables: ${missingEnvVars.join(', ')}`
);
}
}
// Validate environment variables once at the top level
validateEnvVars();
describe('PostgresEngine Instance creation', () => {
// Conditionally skip the entire describe block
if (!envVarsFound) {
describe.skip('PostgresEngine Instance creation (skipped due to missing env vars)', () => {
test('this test will be skipped', () => {
// This block won't execute
});
});
return; // Exit the describe block early
}
let PEInstance: PostgresEngine;
const poolConfig: knex.Knex.PoolConfig = {
min: 0,
max: 5,
};
test('should throw an error if only user or password are passed', async () => {
const pgArgs: PostgresEngineArgs = {
user: process.env.DB_USER ?? '',
};
async function createInstance() {
PEInstance = await PostgresEngine.fromInstance(
process.env.PROJECT_ID!,
process.env.REGION!,
process.env.INSTANCE_ID!,
process.env.DATABASE_ID!,
pgArgs
);
}
await expect(createInstance).rejects.toBe(
"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.'
);
});
test('should create a PostgresEngine Instance using user and password', async () => {
const pgArgs: PostgresEngineArgs = {
user: process.env.DB_USER!,
password: process.env.DB_PASSWORD!,
};
PEInstance = await PostgresEngine.fromInstance(
process.env.PROJECT_ID!,
process.env.REGION!,
process.env.INSTANCE_ID!,
process.env.DATABASE_ID!,
pgArgs
);
const result = await PEInstance.testConnection();
const currentTimestamp = result[0].currentTimestamp;
expect(currentTimestamp).toBeDefined();
try {
await PEInstance.closeConnection();
} catch (error) {
throw new Error(`Error on closing connection: ${error}`);
}
});
// Example of conditionally skipping a single test if a specific env var is missing
(process.env.EMAIL ? test : test.skip)(
'should create a PostgresEngine Instance with IAM email',
async () => {
const pgArgs: PostgresEngineArgs = {
ipType: IpAddressTypes.PUBLIC,
iamAccountEmail: process.env.EMAIL!, // Use ! as we checked for its existence
};
PEInstance = await PostgresEngine.fromInstance(
process.env.PROJECT_ID!,
process.env.REGION!,
process.env.INSTANCE_ID!,
process.env.DATABASE_ID!,
pgArgs
);
const result = await PEInstance.testConnection();
const currentTimestamp = result[0].currentTimestamp;
expect(currentTimestamp).toBeDefined();
try {
await PEInstance.closeConnection();
} catch (error) {
throw new Error(`Error on closing connection: ${error}`);
}
}
);
test('should create a PostgresEngine Instance through from_engine method', async () => {
PostgresEngine.connector = new Connector({ userAgent: USER_AGENT });
const clientOpts = await PostgresEngine.connector.getOptions({
instanceConnectionName: `${process.env.PROJECT_ID}:${process.env.REGION}:${process.env.INSTANCE_ID}`,
ipType: IpAddressTypes.PUBLIC,
authType: AuthTypes.PASSWORD,
});
const dbConfig: knex.Knex.Config<any> = {
client: 'pg',
connection: {
...clientOpts,
password: process.env.DB_PASSWORD,
user: process.env.DB_USER,
database: process.env.DATABASE_ID,
},
};
const engine = knex(dbConfig);
PEInstance = await PostgresEngine.fromEngine(engine);
const result = await PEInstance.testConnection();
const currentTimestamp = result[0].currentTimestamp;
expect(currentTimestamp).toBeDefined();
try {
await PEInstance.closeConnection();
} catch (error) {
throw new Error(`Error on closing connection: ${error}`);
}
});
test('should throw an error if the URL passed to from_engine_args does not have the driver', async () => {
const url = '';
async function createInstance() {
PEInstance = await PostgresEngine.fromEngineArgs(url);
}
await expect(createInstance).rejects.toBe(
"Driver must be type 'postgresql+asyncpg'"
);
});
test('should create a PostgresEngine Instance through from_engine_args using a URL', async () => {
const url = `postgresql+asyncpg://${process.env.DB_USER}:${process.env.DB_PASSWORD}@${process.env.IP_ADDRESS}:5432/${process.env.DATABASE_ID}`;
PEInstance = await PostgresEngine.fromEngineArgs(url, poolConfig);
const result = await PEInstance.testConnection();
const currentTimestamp = result[0].currentTimestamp;
expect(currentTimestamp).toBeDefined();
try {
await PEInstance.closeConnection();
} catch (error) {
throw new Error(`Error on closing connection: ${error}`);
}
});
});
describe('PostgresEngine - table initialization', () => {
// Conditionally skip the entire describe block
if (!envVarsFound) {
describe.skip('PostgresEngine - table initialization (skipped due to missing env vars)', () => {
test('this test will be skipped', () => {
// This block won't execute
});
});
return; // Exit the describe block early
}
let PEInstance: PostgresEngine;
beforeAll(async () => {
const pgArgs: PostgresEngineArgs = {
user: process.env.DB_USER!,
password: process.env.DB_PASSWORD!,
};
PEInstance = await PostgresEngine.fromInstance(
process.env.PROJECT_ID!,
process.env.REGION!,
process.env.INSTANCE_ID!,
process.env.DATABASE_ID!,
pgArgs
);
});
test('should create the vectorstore table', async () => {
const vsTableArgs: VectorStoreTableArgs = {
contentColumn: CONTENT_COLUMN,
embeddingColumn: EMBEDDING_COLUMN,
idColumn: ID_COLUMN,
metadataColumns: METADATA_COLUMNS,
storeMetadata: STORE_METADATA,
overwriteExisting: true,
};
await PEInstance.initVectorstoreTable(
CUSTOM_TABLE,
VECTOR_SIZE,
vsTableArgs
);
const query = `SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '${CUSTOM_TABLE}';`;
const expected = [
{ column_name: 'uuid', data_type: 'uuid' },
{ column_name: 'my_embedding', data_type: 'USER-DEFINED' },
{ column_name: 'json_metadata', data_type: 'json' },
{ column_name: 'my_content', data_type: 'text' },
{ column_name: 'page', data_type: 'text' },
{ column_name: 'source', data_type: 'text' },
];
const { rows } = await PEInstance.pool.raw(query);
rows.forEach((row: any, index: number) => {
expect(row).toMatchObject(expected[index]);
});
});
afterAll(async () => {
// Only attempt to drop the table if the instance was successfully created
if (PEInstance && PEInstance.pool) {
await PEInstance.pool.raw(`DROP TABLE "${CUSTOM_TABLE}"`);
}
try {
if (PEInstance) {
await PEInstance.closeConnection();
}
} catch (error) {
throw new Error(`Error on closing connection: ${error}`);
}
});
});