@mastra/pg 0.2.10-alpha.3 → 0.2.10-alpha.5

package/dist/index.js

@@ -291,10 +291,15 @@ var PgVector = class extends MastraVector {
   describeIndexCache = /* @__PURE__ */ new Map();
   createdIndexes = /* @__PURE__ */ new Map();
   mutexesByName = /* @__PURE__ */ new Map();
+  schema;
+  setupSchemaPromise = null;
   installVectorExtensionPromise = null;
   vectorExtensionInstalled = void 0;
-  constructor(connectionString) {
+  schemaSetupComplete = void 0;
+  constructor(config) {
     super();
+    const connectionString = typeof config === "string" ? config : config.connectionString;
+    this.schema = typeof config === "string" ? void 0 : config.schemaName;
     const basePool = new pg.Pool({
       connectionString,
       max: 20,
@@ -329,6 +334,9 @@ var PgVector = class extends MastraVector {
     if (!this.mutexesByName.has(indexName)) this.mutexesByName.set(indexName, new Mutex());
     return this.mutexesByName.get(indexName);
   }
+  getTableName(indexName) {
+    return this.schema ? `${this.schema}.${indexName}` : indexName;
+  }
   transformFilter(filter) {
     const translator = new PGFilterTranslator();
     return translator.translate(filter);
@@ -360,6 +368,7 @@ var PgVector = class extends MastraVector {
       if (indexInfo.type === "ivfflat" && probes) {
         await client.query(`SET LOCAL ivfflat.probes = ${probes}`);
       }
+      const tableName = this.getTableName(indexName);
       const query = `
         WITH vector_scores AS (
           SELECT
@@ -367,7 +376,7 @@ var PgVector = class extends MastraVector {
             1 - (embedding <=> '${vectorStr}'::vector) as score,
             metadata
             ${includeVector ? ", embedding" : ""}
-          FROM ${indexName}
+          FROM ${tableName}
           ${filterQuery}
         )
         SELECT *
@@ -389,13 +398,14 @@ var PgVector = class extends MastraVector {
   async upsert(...args) {
     const params = this.normalizeArgs("upsert", args);
     const { indexName, vectors, metadata, ids } = params;
+    const tableName = this.getTableName(indexName);
     const client = await this.pool.connect();
     try {
       await client.query("BEGIN");
       const vectorIds = ids || vectors.map(() => crypto.randomUUID());
       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
@@ -432,12 +442,54 @@ var PgVector = class extends MastraVector {
     const existingIndexCacheKey = this.createdIndexes.get(indexName);
     return existingIndexCacheKey && existingIndexCacheKey === newKey;
   }
+  async setupSchema(client) {
+    if (!this.schema || this.schemaSetupComplete) {
+      return;
+    }
+    if (!this.setupSchemaPromise) {
+      this.setupSchemaPromise = (async () => {
+        try {
+          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.`
+              );
+            }
+          }
+          this.schemaSetupComplete = true;
+          this.logger.debug(`Schema "${this.schema}" is ready for use`);
+        } catch (error) {
+          this.schemaSetupComplete = void 0;
+          this.setupSchemaPromise = null;
+          throw error;
+        } finally {
+          this.setupSchemaPromise = null;
+        }
+      })();
+    }
+    await this.setupSchemaPromise;
+  }
   async createIndex(...args) {
     const params = this.normalizeArgs("createIndex", args, [
       "indexConfig",
       "buildIndex"
     ]);
     const { indexName, dimension, metric = "cosine", indexConfig = {}, buildIndex = true } = params;
+    const tableName = this.getTableName(indexName);
     if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
       throw new Error("Invalid index name format");
     }
@@ -455,22 +507,22 @@ var PgVector = class extends MastraVector {
       }
       const client = await this.pool.connect();
       try {
+        await this.setupSchema(client);
         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) {
           await this.setupIndex({ indexName, metric, indexConfig }, client);
         }
       } catch (error) {
         this.createdIndexes.delete(indexName);
-        console.error("Failed to create vector table:", error);
         throw error;
       } finally {
         client.release();
@@ -499,8 +551,9 @@ var PgVector = class extends MastraVector {
   async setupIndex({ indexName, metric, indexConfig }, client) {
     const mutex = this.getMutexByName(`build-${indexName}`);
     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") {
         this.describeIndexCache.delete(indexName);
@@ -513,7 +566,7 @@ var PgVector = class extends MastraVector {
         const efConstruction = indexConfig.hnsw?.efConstruction ?? 32;
         indexSQL = `
           CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx 
-          ON ${indexName} 
+          ON ${tableName} 
           USING hnsw (embedding ${metricOp})
           WITH (
             m = ${m},
@@ -525,12 +578,12 @@ var PgVector = class 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(4e3, 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});
         `;
@@ -582,10 +635,10 @@ var PgVector = class 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();
@@ -594,14 +647,16 @@ var PgVector = class extends MastraVector {
   async describeIndex(indexName) {
     const client = await this.pool.connect();
     try {
+      const tableName = this.getTableName(indexName);
       const dimensionQuery = `
                 SELECT atttypmod as dimension
                 FROM pg_attribute
                 WHERE attrelid = $1::regclass
                 AND attname = 'embedding';
             `;
-      const countQuery = `                SELECT COUNT(*) as count
-                FROM ${indexName};
+      const countQuery = `
+                SELECT COUNT(*) as count
+                FROM ${tableName};
             `;
       const indexQuery = `
             SELECT
@@ -612,10 +667,10 @@ var PgVector = class 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)
       ]);
@@ -652,7 +707,8 @@ var PgVector = class extends MastraVector {
   async deleteIndex(indexName) {
     const client = await this.pool.connect();
     try {
-      await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
+      const tableName = this.getTableName(indexName);
+      await client.query(`DROP TABLE IF EXISTS ${tableName} CASCADE`);
       this.createdIndexes.delete(indexName);
     } catch (error) {
       await client.query("ROLLBACK");
@@ -664,7 +720,8 @@ var PgVector = class extends MastraVector {
   async truncateIndex(indexName) {
     const client = await this.pool.connect();
     try {
-      await client.query(`TRUNCATE ${indexName}`);
+      const tableName = this.getTableName(indexName);
+      await client.query(`TRUNCATE ${tableName}`);
     } catch (e) {
       await client.query("ROLLBACK");
       throw new Error(`Failed to truncate vector table: ${e.message}`);
@@ -696,8 +753,9 @@ var PgVector = class extends MastraVector {
       if (updateParts.length === 0) {
         return;
       }
+      const tableName = this.getTableName(indexName);
       const query = `
-        UPDATE ${indexName}
+        UPDATE ${tableName}
         SET ${updateParts.join(", ")}
         WHERE vector_id = $1
       `;
@@ -709,8 +767,9 @@ var PgVector = class extends MastraVector {
   async deleteIndexById(indexName, id) {
     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]);
@@ -722,9 +781,13 @@ var PgVector = class extends MastraVector {
 var PostgresStore = class extends MastraStorage {
   db;
   pgp;
+  schema;
+  setupSchemaPromise = null;
+  schemaSetupComplete = void 0;
   constructor(config) {
     super({ name: "PostgresStore" });
     this.pgp = pgPromise();
+    this.schema = config.schema;
     this.db = this.pgp(
       `connectionString` in config ? { connectionString: config.connectionString } : {
         host: config.host,
@@ -736,9 +799,12 @@ var PostgresStore = class extends MastraStorage {
       }
     );
   }
+  getTableName(indexName) {
+    return this.schema ? `${this.schema}."${indexName}"` : `"${indexName}"`;
+  }
   async getEvalsByAgentName(agentName, type) {
     try {
-      const baseQuery = `SELECT * FROM ${TABLE_EVALS} WHERE agent_name = $1`;
+      const baseQuery = `SELECT * FROM ${this.getTableName(TABLE_EVALS)} WHERE agent_name = $1`;
       const typeCondition = type === "test" ? " AND test_info IS NOT NULL AND test_info->>'testPath' IS NOT NULL" : type === "live" ? " AND (test_info IS NULL OR test_info->>'testPath' IS NULL)" : "";
       const query = `${baseQuery}${typeCondition} ORDER BY created_at DESC`;
       const rows = await this.db.manyOrNone(query, [agentName]);
@@ -832,7 +898,10 @@ var PostgresStore = class extends MastraStorage {
         args.push(value);
       }
     }
-    const result = await this.db.manyOrNone(`SELECT * FROM ${TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`, args);
+    const result = await this.db.manyOrNone(
+      `SELECT * FROM ${this.getTableName(TABLE_TRACES)} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`,
+      args
+    );
     if (!result) {
       return [];
     }
@@ -853,6 +922,46 @@ var PostgresStore = class extends MastraStorage {
       createdAt: row.createdAt
     }));
   }
+  async setupSchema() {
+    if (!this.schema || this.schemaSetupComplete) {
+      return;
+    }
+    if (!this.setupSchemaPromise) {
+      this.setupSchemaPromise = (async () => {
+        try {
+          const schemaExists = await this.db.oneOrNone(
+            `
+            SELECT EXISTS (
+              SELECT 1 FROM information_schema.schemata 
+              WHERE schema_name = $1
+            )
+          `,
+            [this.schema]
+          );
+          if (!schemaExists?.exists) {
+            try {
+              await this.db.none(`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.`
+              );
+            }
+          }
+          this.schemaSetupComplete = true;
+          this.logger.debug(`Schema "${this.schema}" is ready for use`);
+        } catch (error) {
+          this.schemaSetupComplete = void 0;
+          this.setupSchemaPromise = null;
+          throw error;
+        } finally {
+          this.setupSchemaPromise = null;
+        }
+      })();
+    }
+    await this.setupSchemaPromise;
+  }
   async createTable({
     tableName,
     schema
@@ -864,8 +973,11 @@ var PostgresStore = class extends MastraStorage {
         if (!def.nullable) constraints.push("NOT NULL");
         return `"${name}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
       }).join(",\n");
+      if (this.schema) {
+        await this.setupSchema();
+      }
       const sql = `
-        CREATE TABLE IF NOT EXISTS ${tableName} (
+        CREATE TABLE IF NOT EXISTS ${this.getTableName(tableName)} (
           ${columns}
         );
         ${tableName === TABLE_WORKFLOW_SNAPSHOT ? `
@@ -873,7 +985,7 @@ var PostgresStore = class extends MastraStorage {
           IF NOT EXISTS (
             SELECT 1 FROM pg_constraint WHERE conname = 'mastra_workflow_snapshot_workflow_name_run_id_key'
           ) THEN
-            ALTER TABLE ${tableName}
+            ALTER TABLE ${this.getTableName(tableName)}
             ADD CONSTRAINT mastra_workflow_snapshot_workflow_name_run_id_key
             UNIQUE (workflow_name, run_id);
           END IF;
@@ -888,7 +1000,7 @@ var PostgresStore = class extends MastraStorage {
   }
   async clearTable({ tableName }) {
     try {
-      await this.db.none(`TRUNCATE TABLE ${tableName} CASCADE`);
+      await this.db.none(`TRUNCATE TABLE ${this.getTableName(tableName)} CASCADE`);
     } catch (error) {
       console.error(`Error clearing table ${tableName}:`, error);
       throw error;
@@ -900,7 +1012,7 @@ var PostgresStore = class extends MastraStorage {
       const values = Object.values(record);
       const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
       await this.db.none(
-        `INSERT INTO ${tableName} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
+        `INSERT INTO ${this.getTableName(tableName)} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
         values
       );
     } catch (error) {
@@ -913,7 +1025,10 @@ var PostgresStore = class extends MastraStorage {
       const keyEntries = Object.entries(keys);
       const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
       const values = keyEntries.map(([_, value]) => value);
-      const result = await this.db.oneOrNone(`SELECT * FROM ${tableName} WHERE ${conditions}`, values);
+      const result = await this.db.oneOrNone(
+        `SELECT * FROM ${this.getTableName(tableName)} WHERE ${conditions}`,
+        values
+      );
       if (!result) {
         return null;
       }
@@ -940,7 +1055,7 @@ var PostgresStore = class extends MastraStorage {
           metadata,
           "createdAt",
           "updatedAt"
-        FROM "${TABLE_THREADS}"
+        FROM ${this.getTableName(TABLE_THREADS)}
         WHERE id = $1`,
         [threadId]
       );
@@ -968,7 +1083,7 @@ var PostgresStore = class extends MastraStorage {
           metadata,
           "createdAt",
           "updatedAt"
-        FROM "${TABLE_THREADS}"
+        FROM ${this.getTableName(TABLE_THREADS)}
         WHERE "resourceId" = $1`,
         [resourceId]
       );
@@ -986,7 +1101,7 @@ var PostgresStore = class extends MastraStorage {
   async saveThread({ thread }) {
     try {
       await this.db.none(
-        `INSERT INTO "${TABLE_THREADS}" (
+        `INSERT INTO ${this.getTableName(TABLE_THREADS)} (
           id,
           "resourceId",
           title,
@@ -1030,7 +1145,7 @@ var PostgresStore = class extends MastraStorage {
         ...metadata
       };
       const thread = await this.db.one(
-        `UPDATE "${TABLE_THREADS}"
+        `UPDATE ${this.getTableName(TABLE_THREADS)}
         SET title = $1,
             metadata = $2,
             "updatedAt" = $3
@@ -1052,8 +1167,8 @@ var PostgresStore = class extends MastraStorage {
   async deleteThread({ threadId }) {
     try {
       await this.db.tx(async (t) => {
-        await t.none(`DELETE FROM "${TABLE_MESSAGES}" WHERE thread_id = $1`, [threadId]);
-        await t.none(`DELETE FROM "${TABLE_THREADS}" WHERE id = $1`, [threadId]);
+        await t.none(`DELETE FROM ${this.getTableName(TABLE_MESSAGES)} WHERE thread_id = $1`, [threadId]);
+        await t.none(`DELETE FROM ${this.getTableName(TABLE_THREADS)} WHERE id = $1`, [threadId]);
       });
     } catch (error) {
       console.error("Error deleting thread:", error);
@@ -1072,7 +1187,7 @@ var PostgresStore = class extends MastraStorage {
             SELECT 
               *,
               ROW_NUMBER() OVER (ORDER BY "createdAt" DESC) as row_num
-            FROM "${TABLE_MESSAGES}"
+            FROM ${this.getTableName(TABLE_MESSAGES)}
             WHERE thread_id = $1
           )
           SELECT
@@ -1115,7 +1230,7 @@ var PostgresStore = class extends MastraStorage {
             type,
             "createdAt", 
             thread_id AS "threadId"
-        FROM "${TABLE_MESSAGES}"
+        FROM ${this.getTableName(TABLE_MESSAGES)}
         WHERE thread_id = $1
         AND id != ALL($2)
         ORDER BY "createdAt" DESC
@@ -1153,7 +1268,7 @@ var PostgresStore = class extends MastraStorage {
       await this.db.tx(async (t) => {
         for (const message of messages) {
           await t.none(
-            `INSERT INTO "${TABLE_MESSAGES}" (id, thread_id, content, "createdAt", role, type) 
+            `INSERT INTO ${this.getTableName(TABLE_MESSAGES)} (id, thread_id, content, "createdAt", role, type) 
              VALUES ($1, $2, $3, $4, $5, $6)`,
             [
               message.id,
@@ -1180,7 +1295,7 @@ var PostgresStore = class extends MastraStorage {
     try {
       const now = (/* @__PURE__ */ new Date()).toISOString();
       await this.db.none(
-        `INSERT INTO "${TABLE_WORKFLOW_SNAPSHOT}" (
+        `INSERT INTO ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} (
           workflow_name,
           run_id,
           snapshot,
@@ -1247,13 +1362,13 @@ var PostgresStore = class extends MastraStorage {
     let total = 0;
     if (limit !== void 0 && offset !== void 0) {
       const countResult = await this.db.one(
-        `SELECT COUNT(*) as count FROM ${TABLE_WORKFLOW_SNAPSHOT} ${whereClause}`,
+        `SELECT COUNT(*) as count FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} ${whereClause}`,
         values
       );
       total = Number(countResult.count);
     }
     const query = `
-      SELECT * FROM ${TABLE_WORKFLOW_SNAPSHOT} 
+      SELECT * FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} 
       ${whereClause} 
       ORDER BY "createdAt" DESC
       ${limit !== void 0 && offset !== void 0 ? ` LIMIT $${paramIndex} OFFSET $${paramIndex + 1}` : ""}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/pg",
-  "version": "0.2.10-alpha.3",
+  "version": "0.2.10-alpha.5",
   "description": "Postgres provider for Mastra - includes both vector and db storage capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -24,7 +24,7 @@
     "pg": "^8.13.3",
     "pg-promise": "^11.11.0",
     "xxhash-wasm": "^1.1.0",
-    "@mastra/core": "^0.8.3-alpha.3"
+    "@mastra/core": "^0.8.3-alpha.4"
   },
   "devDependencies": {
     "@microsoft/api-extractor": "^7.52.1",