@mastra/pg 0.16.0 → 0.16.1-alpha.1

package/dist/index.js

@@ -363,7 +363,9 @@ var PgVector = class extends MastraVector {
   setupSchemaPromise = null;
   installVectorExtensionPromise = null;
   vectorExtensionInstalled = void 0;
+  vectorExtensionSchema = null;
   schemaSetupComplete = void 0;
+  cacheWarmupPromise = null;
   constructor({
     connectionString,
     schemaName,
@@ -394,18 +396,24 @@ var PgVector = class extends MastraVector {
           "vector.type": "postgres"
         }
       }) ?? basePool;
-      void (async () => {
-        const existingIndexes = await this.listIndexes();
-        void existingIndexes.map(async (indexName) => {
-          const info = await this.getIndexInfo({ indexName });
-          const key = await this.getIndexCacheKey({
-            indexName,
-            metric: info.metric,
-            dimension: info.dimension,
-            type: info.type
-          });
-          this.createdIndexes.set(indexName, key);
-        });
+      this.cacheWarmupPromise = (async () => {
+        try {
+          const existingIndexes = await this.listIndexes();
+          await Promise.all(
+            existingIndexes.map(async (indexName) => {
+              const info = await this.getIndexInfo({ indexName });
+              const key = await this.getIndexCacheKey({
+                indexName,
+                metric: info.metric,
+                dimension: info.dimension,
+                type: info.type
+              });
+              this.createdIndexes.set(indexName, key);
+            })
+          );
+        } catch (error) {
+          this.logger?.debug("Cache warming skipped or failed", { error });
+        }
       })();
     } catch (error) {
       throw new MastraError(
@@ -425,6 +433,45 @@ var PgVector = class extends MastraVector {
     if (!this.mutexesByName.has(indexName)) this.mutexesByName.set(indexName, new Mutex());
     return this.mutexesByName.get(indexName);
   }
+  /**
+   * Detects which schema contains the vector extension
+   */
+  async detectVectorExtensionSchema(client) {
+    try {
+      const result = await client.query(`
+        SELECT n.nspname as schema_name
+        FROM pg_extension e
+        JOIN pg_namespace n ON e.extnamespace = n.oid
+        WHERE e.extname = 'vector'
+        LIMIT 1;
+      `);
+      if (result.rows.length > 0) {
+        this.vectorExtensionSchema = result.rows[0].schema_name;
+        this.logger.debug("Vector extension found in schema", { schema: this.vectorExtensionSchema });
+        return this.vectorExtensionSchema;
+      }
+      return null;
+    } catch (error) {
+      this.logger.debug("Could not detect vector extension schema", { error });
+      return null;
+    }
+  }
+  /**
+   * Gets the properly qualified vector type name
+   */
+  getVectorTypeName() {
+    if (this.vectorExtensionSchema) {
+      if (this.vectorExtensionSchema === "pg_catalog") {
+        return "vector";
+      }
+      if (this.vectorExtensionSchema === (this.schema || "public")) {
+        return "vector";
+      }
+      const validatedSchema = parseSqlIdentifier(this.vectorExtensionSchema, "vector extension schema");
+      return `${validatedSchema}.vector`;
+    }
+    return "vector";
+  }
   getTableName(indexName) {
     const parsedIndexName = parseSqlIdentifier(indexName, "index name");
     const quotedIndexName = `"${parsedIndexName}"`;
@@ -496,11 +543,12 @@ var PgVector = class extends MastraVector {
         await client.query(`SET LOCAL ivfflat.probes = ${probes}`);
       }
       const { tableName } = this.getTableName(indexName);
+      const vectorType = this.getVectorTypeName();
       const query = `
         WITH vector_scores AS (
           SELECT
             vector_id as id,
-            1 - (embedding <=> '${vectorStr}'::vector) as score,
+            1 - (embedding <=> '${vectorStr}'::${vectorType}) as score,
             metadata
             ${includeVector ? ", embedding" : ""}
           FROM ${tableName}
@@ -544,13 +592,14 @@ var PgVector = class extends MastraVector {
     try {
       await client.query("BEGIN");
       const vectorIds = ids || vectors.map(() => crypto.randomUUID());
+      const vectorType = this.getVectorTypeName();
       for (let i = 0; i < vectors.length; i++) {
         const query = `
           INSERT INTO ${tableName} (vector_id, embedding, metadata)
-          VALUES ($1, $2::vector, $3::jsonb)
+          VALUES ($1, $2::${vectorType}, $3::jsonb)
           ON CONFLICT (vector_id)
           DO UPDATE SET
-            embedding = $2::vector,
+            embedding = $2::${vectorType},
             metadata = $3::jsonb
           RETURNING embedding::text
         `;
@@ -697,11 +746,15 @@ var PgVector = class extends MastraVector {
       try {
         await this.setupSchema(client);
         await this.installVectorExtension(client);
+        if (this.schema && this.vectorExtensionSchema && this.schema !== this.vectorExtensionSchema && this.vectorExtensionSchema !== "pg_catalog") {
+          await client.query(`SET search_path TO ${this.getSchemaName()}, "${this.vectorExtensionSchema}"`);
+        }
+        const vectorType = this.getVectorTypeName();
         await client.query(`
           CREATE TABLE IF NOT EXISTS ${tableName} (
             id SERIAL PRIMARY KEY,
             vector_id TEXT UNIQUE NOT NULL,
-            embedding vector(${dimension}),
+            embedding ${vectorType}(${dimension}),
             metadata JSONB DEFAULT '{}'::jsonb
           );
         `);
@@ -756,17 +809,63 @@ var PgVector = class extends MastraVector {
   async setupIndex({ indexName, metric, indexConfig }, client) {
     const mutex = this.getMutexByName(`build-${indexName}`);
     await mutex.runExclusive(async () => {
+      const isConfigEmpty = !indexConfig || Object.keys(indexConfig).length === 0 || !indexConfig.type && !indexConfig.ivf && !indexConfig.hnsw;
+      const indexType = isConfigEmpty ? "ivfflat" : indexConfig.type || "ivfflat";
       const { tableName, vectorIndexName } = this.getTableName(indexName);
-      if (this.createdIndexes.has(indexName)) {
+      let existingIndexInfo = null;
+      let dimension = 0;
+      try {
+        existingIndexInfo = await this.getIndexInfo({ indexName });
+        dimension = existingIndexInfo.dimension;
+        if (isConfigEmpty && existingIndexInfo.metric === metric) {
+          if (existingIndexInfo.type === "flat") {
+            this.logger?.debug(`No index exists for ${vectorIndexName}, will create default ivfflat index`);
+          } else {
+            this.logger?.debug(
+              `Index ${vectorIndexName} already exists (type: ${existingIndexInfo.type}, metric: ${existingIndexInfo.metric}), preserving existing configuration`
+            );
+            const cacheKey = await this.getIndexCacheKey({
+              indexName,
+              dimension,
+              type: existingIndexInfo.type,
+              metric: existingIndexInfo.metric
+            });
+            this.createdIndexes.set(indexName, cacheKey);
+            return;
+          }
+        }
+        let configMatches = existingIndexInfo.metric === metric && existingIndexInfo.type === indexType;
+        if (indexType === "hnsw") {
+          configMatches = configMatches && existingIndexInfo.config.m === (indexConfig.hnsw?.m ?? 8) && existingIndexInfo.config.efConstruction === (indexConfig.hnsw?.efConstruction ?? 32);
+        } else if (indexType === "flat") {
+          configMatches = configMatches && existingIndexInfo.type === "flat";
+        } else if (indexType === "ivfflat" && indexConfig.ivf?.lists) {
+          configMatches = configMatches && existingIndexInfo.config.lists === indexConfig.ivf?.lists;
+        }
+        if (configMatches) {
+          this.logger?.debug(`Index ${vectorIndexName} already exists with same configuration, skipping recreation`);
+          const cacheKey = await this.getIndexCacheKey({
+            indexName,
+            dimension,
+            type: existingIndexInfo.type,
+            metric: existingIndexInfo.metric
+          });
+          this.createdIndexes.set(indexName, cacheKey);
+          return;
+        }
+        this.logger?.info(`Index ${vectorIndexName} configuration changed, rebuilding index`);
         await client.query(`DROP INDEX IF EXISTS ${vectorIndexName}`);
+        this.describeIndexCache.delete(indexName);
+      } catch {
+        this.logger?.debug(`Index ${indexName} doesn't exist yet, will create it`);
       }
-      if (indexConfig.type === "flat") {
+      if (indexType === "flat") {
         this.describeIndexCache.delete(indexName);
         return;
       }
       const metricOp = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
       let indexSQL;
-      if (indexConfig.type === "hnsw") {
+      if (indexType === "hnsw") {
         const m = indexConfig.hnsw?.m ?? 8;
         const efConstruction = indexConfig.hnsw?.efConstruction ?? 32;
         indexSQL = `
@@ -803,27 +902,48 @@ var PgVector = class extends MastraVector {
     if (!this.installVectorExtensionPromise) {
       this.installVectorExtensionPromise = (async () => {
         try {
-          const extensionCheck = await client.query(`
-            SELECT EXISTS (
-              SELECT 1 FROM pg_extension WHERE extname = 'vector'
+          const existingSchema = await this.detectVectorExtensionSchema(client);
+          if (existingSchema) {
+            this.vectorExtensionInstalled = true;
+            this.vectorExtensionSchema = existingSchema;
+            this.logger.info(`Vector extension already installed in schema: ${existingSchema}`);
+            return;
+          }
+          try {
+            if (this.schema && this.schema !== "public") {
+              try {
+                await client.query(`CREATE EXTENSION IF NOT EXISTS vector SCHEMA ${this.getSchemaName()}`);
+                this.vectorExtensionInstalled = true;
+                this.vectorExtensionSchema = this.schema;
+                this.logger.info(`Vector extension installed in schema: ${this.schema}`);
+                return;
+              } catch (schemaError) {
+                this.logger.debug(`Could not install vector extension in schema ${this.schema}, trying public schema`, {
+                  error: schemaError
+                });
+              }
+            }
+            await client.query("CREATE EXTENSION IF NOT EXISTS vector");
+            const installedSchema = await this.detectVectorExtensionSchema(client);
+            if (installedSchema) {
+              this.vectorExtensionInstalled = true;
+              this.vectorExtensionSchema = installedSchema;
+              this.logger.info(`Vector extension installed in schema: ${installedSchema}`);
+            }
+          } catch (error) {
+            this.logger.warn(
+              "Could not install vector extension. This requires superuser privileges. If the extension is already installed, you can ignore this warning.",
+              { error }
             );
-          `);
-          this.vectorExtensionInstalled = extensionCheck.rows[0].exists;
-          if (!this.vectorExtensionInstalled) {
-            try {
-              await client.query("CREATE EXTENSION IF NOT EXISTS vector");
+            const existingSchema2 = await this.detectVectorExtensionSchema(client);
+            if (existingSchema2) {
               this.vectorExtensionInstalled = true;
-              this.logger.info("Vector extension installed successfully");
-            } catch {
-              this.logger.warn(
-                "Could not install vector extension. This requires superuser privileges. If the extension is already installed globally, you can ignore this warning."
-              );
+              this.vectorExtensionSchema = existingSchema2;
+              this.logger.info(`Vector extension found in schema: ${existingSchema2}`);
             }
-          } else {
-            this.logger.debug("Vector extension already installed, skipping installation");
           }
         } catch (error) {
-          this.logger.error("Error checking vector extension status", { error });
+          this.logger.error("Error setting up vector extension", { error });
           this.vectorExtensionInstalled = void 0;
           this.installVectorExtensionPromise = null;
           throw error;
@@ -1025,6 +1145,12 @@ var PgVector = class extends MastraVector {
     }
   }
   async disconnect() {
+    if (this.cacheWarmupPromise) {
+      try {
+        await this.cacheWarmupPromise;
+      } catch {
+      }
+    }
     await this.pool.end();
   }
   /**
@@ -1047,8 +1173,9 @@ var PgVector = class extends MastraVector {
       let updateParts = [];
       let values = [id];
       let valueIndex = 2;
+      const vectorType = this.getVectorTypeName();
       if (update.vector) {
-        updateParts.push(`embedding = $${valueIndex}::vector`);
+        updateParts.push(`embedding = $${valueIndex}::${vectorType}`);
         values.push(`[${update.vector.join(",")}]`);
         valueIndex++;
       }