@mastra/pg 0.2.10-alpha.4 → 0.2.10-alpha.6

package/src/vector/index.ts

@@ -64,12 +64,20 @@ export class PgVector extends MastraVector {
   private describeIndexCache: Map<string, PGIndexStats> = new Map();
   private createdIndexes = new Map<string, number>();
   private mutexesByName = new Map<string, Mutex>();
+  private schema?: string;
+  private setupSchemaPromise: Promise<void> | null = null;
   private installVectorExtensionPromise: Promise<void> | null = null;
   private vectorExtensionInstalled: boolean | undefined = undefined;
+  private schemaSetupComplete: boolean | undefined = undefined;
 
-  constructor(connectionString: string) {
+  constructor(connectionString: string);
+  constructor(config: { connectionString: string; schemaName?: string });
+  constructor(config: string | { connectionString: string; schemaName?: string }) {
     super();
 
+    const connectionString = typeof config === 'string' ? config : config.connectionString;
+    this.schema = typeof config === 'string' ? undefined : config.schemaName;
+
     const basePool = new pg.Pool({
       connectionString,
       max: 20, // Maximum number of clients in the pool
@@ -108,6 +116,10 @@ export class PgVector extends MastraVector {
     return this.mutexesByName.get(indexName)!;
   }
 
+  private getTableName(indexName: string) {
+    return this.schema ? `${this.schema}.${indexName}` : indexName;
+  }
+
   transformFilter(filter?: VectorFilter) {
     const translator = new PGFilterTranslator();
     return translator.translate(filter);
@@ -149,6 +161,8 @@ export class PgVector extends MastraVector {
         await client.query(`SET LOCAL ivfflat.probes = ${probes}`);
       }
 
+      const tableName = this.getTableName(indexName);
+
       const query = `
         WITH vector_scores AS (
           SELECT
@@ -156,7 +170,7 @@ export class PgVector extends MastraVector {
             1 - (embedding <=> '${vectorStr}'::vector) as score,
             metadata
             ${includeVector ? ', embedding' : ''}
-          FROM ${indexName}
+          FROM ${tableName}
           ${filterQuery}
         )
         SELECT *
@@ -181,6 +195,7 @@ export class PgVector extends MastraVector {
     const params = this.normalizeArgs<UpsertVectorParams>('upsert', args);
 
     const { indexName, vectors, metadata, ids } = params;
+    const tableName = this.getTableName(indexName);
 
     // Start a transaction
     const client = await this.pool.connect();
@@ -190,7 +205,7 @@ export class PgVector extends MastraVector {
 
       for (let i = 0; i < vectors.length; i++) {
         const query = `
-          INSERT INTO ${indexName} (vector_id, embedding, metadata)
+          INSERT INTO ${tableName} (vector_id, embedding, metadata)
           VALUES ($1, $2::vector, $3::jsonb)
           ON CONFLICT (vector_id)
           DO UPDATE SET
@@ -231,6 +246,57 @@ export class PgVector extends MastraVector {
     const existingIndexCacheKey = this.createdIndexes.get(indexName);
     return existingIndexCacheKey && existingIndexCacheKey === newKey;
   }
+  private async setupSchema(client: pg.PoolClient) {
+    if (!this.schema || this.schemaSetupComplete) {
+      return;
+    }
+
+    if (!this.setupSchemaPromise) {
+      this.setupSchemaPromise = (async () => {
+        try {
+          // First check if schema exists and we have usage permission
+          const schemaCheck = await client.query(
+            `
+            SELECT EXISTS (
+              SELECT 1 FROM information_schema.schemata 
+              WHERE schema_name = $1
+            )
+          `,
+            [this.schema],
+          );
+
+          const schemaExists = schemaCheck.rows[0].exists;
+
+          if (!schemaExists) {
+            try {
+              await client.query(`CREATE SCHEMA IF NOT EXISTS ${this.schema}`);
+              this.logger.info(`Schema "${this.schema}" created successfully`);
+            } catch (error) {
+              this.logger.error(`Failed to create schema "${this.schema}"`, { error });
+              throw new Error(
+                `Unable to create schema "${this.schema}". This requires CREATE privilege on the database. ` +
+                  `Either create the schema manually or grant CREATE privilege to the user.`,
+              );
+            }
+          }
+
+          // If we got here, schema exists and we can use it
+          this.schemaSetupComplete = true;
+          this.logger.debug(`Schema "${this.schema}" is ready for use`);
+        } catch (error) {
+          // Reset flags so we can retry
+          this.schemaSetupComplete = undefined;
+          this.setupSchemaPromise = null;
+          throw error;
+        } finally {
+          this.setupSchemaPromise = null;
+        }
+      })();
+    }
+
+    await this.setupSchemaPromise;
+  }
+
   async createIndex(...args: ParamsToArgs<PgCreateIndexParams> | PgCreateIndexArgs): Promise<void> {
     const params = this.normalizeArgs<PgCreateIndexParams, PgCreateIndexArgs>('createIndex', args, [
       'indexConfig',
@@ -238,6 +304,7 @@ export class PgVector extends MastraVector {
     ]);
 
     const { indexName, dimension, metric = 'cosine', indexConfig = {}, buildIndex = true } = params;
+    const tableName = this.getTableName(indexName);
 
     // Validate inputs
     if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
@@ -262,17 +329,21 @@ export class PgVector extends MastraVector {
       }
 
       const client = await this.pool.connect();
+
       try {
-        // install vector extension
+        // Setup schema if needed
+        await this.setupSchema(client);
+
+        // Install vector extension first (needs to be in public schema)
         await this.installVectorExtension(client);
         await client.query(`
-          CREATE TABLE IF NOT EXISTS ${indexName} (
+          CREATE TABLE IF NOT EXISTS ${tableName} (
             id SERIAL PRIMARY KEY,
             vector_id TEXT UNIQUE NOT NULL,
             embedding vector(${dimension}),
             metadata JSONB DEFAULT '{}'::jsonb
           );
-      `);
+        `);
         this.createdIndexes.set(indexName, indexCacheKey);
 
         if (buildIndex) {
@@ -280,7 +351,6 @@ export class PgVector extends MastraVector {
         }
       } catch (error: any) {
         this.createdIndexes.delete(indexName);
-        console.error('Failed to create vector table:', error);
         throw error;
       } finally {
         client.release();
@@ -319,8 +389,10 @@ export class PgVector extends MastraVector {
     const mutex = this.getMutexByName(`build-${indexName}`);
     // Use async-mutex instead of advisory lock for perf (over 2x as fast)
     await mutex.runExclusive(async () => {
+      const tableName = this.getTableName(indexName);
+
       if (this.createdIndexes.has(indexName)) {
-        await client.query(`DROP INDEX IF EXISTS ${indexName}_vector_idx`);
+        await client.query(`DROP INDEX IF EXISTS ${tableName}_vector_idx`);
       }
 
       if (indexConfig.type === 'flat') {
@@ -338,7 +410,7 @@ export class PgVector extends MastraVector {
 
         indexSQL = `
           CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx 
-          ON ${indexName} 
+          ON ${tableName} 
           USING hnsw (embedding ${metricOp})
           WITH (
             m = ${m},
@@ -350,12 +422,12 @@ export class PgVector extends MastraVector {
         if (indexConfig.ivf?.lists) {
           lists = indexConfig.ivf.lists;
         } else {
-          const size = (await client.query(`SELECT COUNT(*) FROM ${indexName}`)).rows[0].count;
+          const size = (await client.query(`SELECT COUNT(*) FROM ${tableName}`)).rows[0].count;
           lists = Math.max(100, Math.min(4000, Math.floor(Math.sqrt(size) * 2)));
         }
         indexSQL = `
           CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx
-          ON ${indexName}
+          ON ${tableName}
           USING ivfflat (embedding ${metricOp})
           WITH (lists = ${lists});
         `;
@@ -423,10 +495,10 @@ export class PgVector extends MastraVector {
       const vectorTablesQuery = `
             SELECT DISTINCT table_name
             FROM information_schema.columns
-            WHERE table_schema = 'public'
+            WHERE table_schema = $1
             AND udt_name = 'vector';
         `;
-      const vectorTables = await client.query(vectorTablesQuery);
+      const vectorTables = await client.query(vectorTablesQuery, [this.schema || 'public']);
       return vectorTables.rows.map(row => row.table_name);
     } finally {
       client.release();
@@ -436,6 +508,8 @@ export class PgVector extends MastraVector {
   async describeIndex(indexName: string): Promise<PGIndexStats> {
     const client = await this.pool.connect();
     try {
+      const tableName = this.getTableName(indexName);
+
       // Get vector dimension
       const dimensionQuery = `
                 SELECT atttypmod as dimension
@@ -445,8 +519,9 @@ export class PgVector extends MastraVector {
             `;
 
       // Get row count
-      const countQuery = `                SELECT COUNT(*) as count
-                FROM ${indexName};
+      const countQuery = `
+                SELECT COUNT(*) as count
+                FROM ${tableName};
             `;
 
       // Get index metric type
@@ -459,11 +534,11 @@ export class PgVector extends MastraVector {
             JOIN pg_class c ON i.indexrelid = c.oid
             JOIN pg_am am ON c.relam = am.oid
             JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
-            WHERE c.relname = '${indexName}_vector_idx';
+            WHERE c.relname = '${tableName}_vector_idx';
             `;
 
       const [dimResult, countResult, indexResult] = await Promise.all([
-        client.query(dimensionQuery, [indexName]),
+        client.query(dimensionQuery, [tableName]),
         client.query(countQuery),
         client.query(indexQuery),
       ]);
@@ -512,8 +587,9 @@ export class PgVector extends MastraVector {
   async deleteIndex(indexName: string): Promise<void> {
     const client = await this.pool.connect();
     try {
+      const tableName = this.getTableName(indexName);
       // Drop the table
-      await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
+      await client.query(`DROP TABLE IF EXISTS ${tableName} CASCADE`);
       this.createdIndexes.delete(indexName);
     } catch (error: any) {
       await client.query('ROLLBACK');
@@ -526,7 +602,8 @@ export class PgVector extends MastraVector {
   async truncateIndex(indexName: string) {
     const client = await this.pool.connect();
     try {
-      await client.query(`TRUNCATE ${indexName}`);
+      const tableName = this.getTableName(indexName);
+      await client.query(`TRUNCATE ${tableName}`);
     } catch (e: any) {
       await client.query('ROLLBACK');
       throw new Error(`Failed to truncate vector table: ${e.message}`);
@@ -572,10 +649,12 @@ export class PgVector extends MastraVector {
         return;
       }
 
+      const tableName = this.getTableName(indexName);
+
       // query looks like this:
       // UPDATE table SET embedding = $2::vector, metadata = $3::jsonb WHERE id = $1
       const query = `
-        UPDATE ${indexName}
+        UPDATE ${tableName}
         SET ${updateParts.join(', ')}
         WHERE vector_id = $1
       `;
@@ -589,8 +668,9 @@ export class PgVector extends MastraVector {
   async deleteIndexById(indexName: string, id: string): Promise<void> {
     const client = await this.pool.connect();
     try {
+      const tableName = this.getTableName(indexName);
       const query = `
-        DELETE FROM ${indexName}
+        DELETE FROM ${tableName}
         WHERE vector_id = $1
       `;
       await client.query(query, [id]);