fix(mem0-ts): support pgvector connection strings and ssl (#5789)

mem0-ts/src/oss/src/vector_stores/pgvector.ts

@@ -1,4 +1,4 @@
-import type { Client as ClientType } from "pg";
+import type { Client as ClientType, ClientConfig } from "pg";
 import pkg from "pg";
 const { Client, escapeIdentifier } = pkg;
 import { VectorStore } from "./base";
@@ -157,41 +157,90 @@ export function buildFilterConditions(
 
 interface PGVectorConfig extends VectorStoreConfig {
   dbname?: string;
-  user: string;
-  password: string;
-  host: string;
-  port: number;
+  user?: string;
+  password?: string;
+  host?: string;
+  port?: number;
+  connectionString?: string;
+  ssl?: ClientConfig["ssl"];
   embeddingModelDims: number;
   diskann?: boolean;
   hnsw?: boolean;
 }
 
+function getConnectionString(config: PGVectorConfig): string | undefined {
+  return config.connectionString?.trim() || undefined;
+}
+
+function validateConnectionConfig(config: PGVectorConfig): void {
+  if (getConnectionString(config)) {
+    return;
+  }
+
+  const missingFields = ["user", "password", "host", "port"].filter((field) => {
+    const v = config[field as keyof PGVectorConfig];
+    return v === undefined || v === null || v === "";
+  });
+
+  if (missingFields.length > 0) {
+    throw new Error(
+      `PGVector requires either connectionString or ${missingFields.join(", ")}`,
+    );
+  }
+}
+
+function buildClientConfig(
+  config: PGVectorConfig,
+  database?: string,
+): ClientConfig {
+  const connectionString = getConnectionString(config);
+  if (connectionString) {
+    return {
+      connectionString,
+      ...(config.ssl !== undefined ? { ssl: config.ssl } : {}),
+    };
+  }
+
+  return {
+    database,
+    user: config.user,
+    password: config.password,
+    host: config.host,
+    port: config.port,
+    ...(config.ssl !== undefined ? { ssl: config.ssl } : {}),
+  };
+}
+
 export class PGVector implements VectorStore {
   private client: ClientType;
   private collectionName: string;
   private useDiskann: boolean;
   private useHnsw: boolean;
   private readonly dbName: string;
+  private readonly useDirectConnection: boolean;
   private config: PGVectorConfig;
   private _initPromise?: Promise<void>;
 
   constructor(config: PGVectorConfig) {
+    validateConnectionConfig(config);
     this.collectionName = validateIdentifier(
       config.collectionName || "memories",
       "collectionName",
     );
     this.useDiskann = config.diskann || false;
     this.useHnsw = config.hnsw || false;
-    this.dbName = validateIdentifier(config.dbname || "vector_store", "dbname");
+    this.useDirectConnection = !!getConnectionString(config);
+    this.dbName = this.useDirectConnection
+      ? ""
+      : validateIdentifier(config.dbname || "vector_store", "dbname");
     this.config = config;
 
-    this.client = new Client({
-      database: "postgres", // Initially connect to default postgres database
-      user: config.user,
-      password: config.password,
-      host: config.host,
-      port: config.port,
-    });
+    this.client = new Client(
+      buildClientConfig(
+        config,
+        this.useDirectConnection ? undefined : "postgres",
+      ),
+    );
     this.initialize().catch(console.error);
   }
 
@@ -210,29 +259,20 @@ export class PGVector implements VectorStore {
     try {
       await this.client.connect();
 
-      // Check if database exists
-      const dbExists = await this.checkDatabaseExists(this.dbName);
-      if (!dbExists) {
-        await this.createDatabase(this.dbName);
+      if (!this.useDirectConnection) {
+        const dbExists = await this.checkDatabaseExists(this.dbName);
+        if (!dbExists) {
+          await this.createDatabase(this.dbName);
+        }
+
+        await this.client.end();
+
+        this.client = new Client(buildClientConfig(this.config, this.dbName));
+        await this.client.connect();
       }
 
-      // Disconnect from postgres database
-      await this.client.end();
-
-      // Connect to the target database
-      this.client = new Client({
-        database: this.dbName,
-        user: this.config.user,
-        password: this.config.password,
-        host: this.config.host,
-        port: this.config.port,
-      });
-      await this.client.connect();
-
-      // Create vector extension
       await this.client.query("CREATE EXTENSION IF NOT EXISTS vector");
 
-      // Create memory_migrations table
       await this.client.query(`
         CREATE TABLE IF NOT EXISTS memory_migrations (
           id SERIAL PRIMARY KEY,
@@ -240,7 +280,6 @@ export class PGVector implements VectorStore {
         )
       `);
 
-      // Check if the collection exists
       const collections = await this.listCols();
       if (!collections.includes(this.collectionName)) {
         await this.createCol(this.config.embeddingModelDims);

mem0-ts/src/oss/tests/pgvector.filters.test.ts

@@ -1,5 +1,3 @@
-/// <reference types="jest" />
-
 jest.mock("pg", () => {
   const Client = jest.fn().mockImplementation(() => ({
     connect: jest.fn().mockResolvedValue(undefined),

mem0-ts/src/oss/tests/pgvector.unit.test.ts

@@ -1,5 +1,3 @@
-/// <reference types="jest" />
-
 const searchRows = [
   {
     id: "a",
@@ -23,9 +21,13 @@ const searchRows = [
   },
 ];
 
+const mockState = {
+  databaseExists: true,
+};
+
 function mockPgQuery(sql: string) {
   if (sql.includes("SELECT 1 FROM pg_database")) {
-    return { rows: [{ "?column?": 1 }] };
+    return { rows: mockState.databaseExists ? [{ "?column?": 1 }] : [] };
   }
 
   if (sql.includes("FROM information_schema.tables")) {
@@ -69,13 +71,129 @@ jest.mock("pg", () => {
 
 import { PGVector } from "../src/vector_stores/pgvector";
 
-describe("PGVector - search()", () => {
+function getClientQueries(client: { query: jest.Mock }) {
+  return client.query.mock.calls.map(([sql]) => sql as string);
+}
+
+describe("PGVector", () => {
   beforeEach(() => {
     const pg = require("pg");
+    mockState.databaseExists = true;
     pg.__mock.Client.mockClear();
     pg.__mock.clients.length = 0;
   });
 
+  test("uses one direct client for connectionString mode and skips bootstrap database creation", async () => {
+    mockState.databaseExists = false;
+
+    const ssl = { rejectUnauthorized: false };
+    const store = new PGVector({
+      collectionName: "memories",
+      connectionString:
+        "postgresql://postgres:postgres@db.example.com:5432/neondb",
+      ssl,
+      embeddingModelDims: 3,
+      dimension: 3,
+    } as any);
+
+    await store.initialize();
+
+    const pg = require("pg");
+    expect(pg.__mock.Client).toHaveBeenCalledTimes(1);
+    expect(pg.__mock.Client).toHaveBeenCalledWith({
+      connectionString:
+        "postgresql://postgres:postgres@db.example.com:5432/neondb",
+      ssl,
+    });
+
+    const directClient = pg.__mock.clients[0];
+    const queries = getClientQueries(directClient);
+
+    expect(queries).not.toEqual(
+      expect.arrayContaining([
+        expect.stringContaining("SELECT 1 FROM pg_database"),
+      ]),
+    );
+    expect(queries).not.toEqual(
+      expect.arrayContaining([expect.stringContaining("CREATE DATABASE")]),
+    );
+    expect(queries).toEqual(
+      expect.arrayContaining([
+        "CREATE EXTENSION IF NOT EXISTS vector",
+        expect.stringContaining("FROM information_schema.tables"),
+      ]),
+    );
+  });
+
+  test("keeps the split-field bootstrap flow when connectionString is absent", async () => {
+    mockState.databaseExists = false;
+    const ssl = { rejectUnauthorized: false };
+
+    const store = new PGVector({
+      collectionName: "memories",
+      user: "postgres",
+      password: "postgres",
+      host: "localhost",
+      port: 5432,
+      dbname: "vector_store",
+      ssl,
+      embeddingModelDims: 3,
+      dimension: 3,
+    } as any);
+
+    await store.initialize();
+
+    const pg = require("pg");
+    expect(pg.__mock.Client).toHaveBeenCalledTimes(2);
+    expect(pg.__mock.Client).toHaveBeenNthCalledWith(1, {
+      database: "postgres",
+      user: "postgres",
+      password: "postgres",
+      host: "localhost",
+      port: 5432,
+      ssl,
+    });
+    expect(pg.__mock.Client).toHaveBeenNthCalledWith(2, {
+      database: "vector_store",
+      user: "postgres",
+      password: "postgres",
+      host: "localhost",
+      port: 5432,
+      ssl,
+    });
+
+    const bootstrapClient = pg.__mock.clients[0];
+    const activeClient = pg.__mock.clients[1];
+    const bootstrapQueries = getClientQueries(bootstrapClient);
+
+    expect(bootstrapQueries).toEqual(
+      expect.arrayContaining([
+        "SELECT 1 FROM pg_database WHERE datname = $1",
+        'CREATE DATABASE "vector_store"',
+      ]),
+    );
+    expect(bootstrapClient.end).toHaveBeenCalledTimes(1);
+    expect(getClientQueries(activeClient)).toEqual(
+      expect.arrayContaining([
+        "CREATE EXTENSION IF NOT EXISTS vector",
+        expect.stringContaining("FROM information_schema.tables"),
+      ]),
+    );
+  });
+
+  test("throws when connectionString is absent and split-field params are missing", () => {
+    expect(
+      () =>
+        new PGVector({
+          collectionName: "memories",
+          embeddingModelDims: 3,
+          dimension: 3,
+        } as any),
+    ).toThrow(
+      "PGVector requires either connectionString or user, password, host, port",
+    );
+  });
+
   test("returns similarity score (1 - distance) clamped to [0, 1]", async () => {
     const store = new PGVector({
       collectionName: "memories",