@mastra/pg 1.9.0 → 1.9.1-alpha.1

package/dist/index.js

@@ -709,6 +709,25 @@ var PgVector = class extends MastraVector {
       return null;
     }
   }
+  /**
+   * Sets search_path on the client connection so that vector operators (e.g. <=>, vector_cosine_ops)
+   * are resolvable when the pgvector extension is installed in a non-default schema.
+   *
+   * PostgreSQL's default search_path is ("$user", public). If the extension lives in a custom schema
+   * (e.g. "myapp"), operator classes and distance operators won't resolve without this.
+   */
+  async ensureSearchPath(client) {
+    if (!this.vectorExtensionSchema) {
+      await this.detectVectorExtensionSchema(client);
+    }
+    if (this.vectorExtensionSchema && this.vectorExtensionSchema !== "public" && this.vectorExtensionSchema !== "pg_catalog") {
+      const schemas = /* @__PURE__ */ new Set();
+      schemas.add(this.vectorExtensionSchema);
+      if (this.schema) schemas.add(this.schema);
+      schemas.add("public");
+      await client.query(`SET search_path TO ${[...schemas].map((s) => `"${s}"`).join(", ")}`);
+    }
+  }
   /**
    * Checks if the installed pgvector version supports halfvec type.
    * halfvec was introduced in pgvector 0.7.0.
@@ -922,6 +941,7 @@ var PgVector = class extends MastraVector {
     }
     const client = await this.pool.connect();
     try {
+      await this.ensureSearchPath(client);
       await client.query("BEGIN");
       const translatedFilter = this.transformFilter(filter);
       const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore, topK);
@@ -941,7 +961,23 @@ var PgVector = class extends MastraVector {
       const qualifiedVectorType = this.getVectorTypeName(indexInfo.vectorType, indexInfo.dimension);
       const distanceExpr = `embedding ${ops.distanceOperator} '${vectorStr}'::${qualifiedVectorType}`;
       const scoreExpr = ops.scoreExpr(distanceExpr);
-      const query = `
+      const hasFilter = filterQuery.trim().length > 0;
+      const useIndexedOrder = indexInfo.type === "hnsw" && !hasFilter && minScore <= 0;
+      const query = useIndexedOrder ? `
+        WITH vector_scores AS (
+          SELECT
+            vector_id as id,
+            ${scoreExpr} as score,
+            metadata
+            ${includeVector ? ", embedding" : ""}
+          FROM ${tableName}
+          ORDER BY ${distanceExpr}
+          LIMIT $2
+        )
+        SELECT *
+        FROM vector_scores
+        WHERE score > $1
+        ORDER BY score DESC` : `
         WITH vector_scores AS (
           SELECT
             vector_id as id,
@@ -994,6 +1030,7 @@ var PgVector = class extends MastraVector {
     const { tableName } = this.getTableName(indexName);
     const client = await this.pool.connect();
     try {
+      await this.ensureSearchPath(client);
       await client.query("BEGIN");
       if (deleteFilter) {
         this.logger?.debug(`Deleting vectors matching filter before upsert`, { indexName, deleteFilter });
@@ -1244,9 +1281,6 @@ var PgVector = class extends MastraVector {
             }
           });
         }
-        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 qualifiedVectorType = this.getVectorTypeName(vectorType);
         await client.query(`
           CREATE TABLE IF NOT EXISTS ${tableName} (
@@ -1389,6 +1423,7 @@ var PgVector = class extends MastraVector {
         this.describeIndexCache.delete(indexName);
         return;
       }
+      await this.ensureSearchPath(client);
       const effectiveVectorType = existingIndexInfo?.vectorType ?? vectorType;
       const metricOp = this.getVectorOps(effectiveVectorType, metric).operatorClass;
       let indexSQL;
@@ -1739,6 +1774,7 @@ var PgVector = class extends MastraVector {
         });
       }
       client = await this.pool.connect();
+      await this.ensureSearchPath(client);
       const { tableName } = this.getTableName(indexName);
       const indexInfo = await this.getIndexInfo({ indexName });
       const qualifiedVectorType = this.getVectorTypeName(indexInfo.vectorType, indexInfo.dimension);
@@ -1977,6 +2013,7 @@ var PoolAdapter = class {
   constructor($pool) {
     this.$pool = $pool;
   }
+  $pool;
   connect() {
     return this.$pool.connect();
   }
@@ -2045,6 +2082,7 @@ var TransactionClient = class {
   constructor(client) {
     this.client = client;
   }
+  client;
   async none(query, values) {
     await this.client.query(query, values);
     return null;
@@ -2474,36 +2512,39 @@ var PgDB = class extends MastraBase {
         return getDefaultValue(type);
     }
   }
-  async insert({ tableName, record }) {
-    try {
-      this.addTimestampZColumns(record);
-      const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
-      const schemaName = getSchemaName(this.schemaName);
-      const columns = Object.keys(filteredRecord).map((col) => parseSqlIdentifier(col, "column name"));
-      if (columns.length === 0) return;
-      const values = this.prepareValuesForInsert(filteredRecord, tableName);
-      const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
-      const fullTableName = getTableName({ indexName: tableName, schemaName });
-      const columnList = columns.map((c) => `"${c}"`).join(", ");
-      if (tableName === TABLE_SPANS) {
-        const updateColumns = columns.filter((c) => c !== "traceId" && c !== "spanId");
-        if (updateColumns.length > 0) {
-          const updateClause = updateColumns.map((c) => `"${c}" = EXCLUDED."${c}"`).join(", ");
-          await this.client.none(
-            `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
+  async executeInsert(client, { tableName, record }) {
+    this.addTimestampZColumns(record);
+    const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
+    const schemaName = getSchemaName(this.schemaName);
+    const columns = Object.keys(filteredRecord).map((col) => parseSqlIdentifier(col, "column name"));
+    if (columns.length === 0) return;
+    const values = this.prepareValuesForInsert(filteredRecord, tableName);
+    const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
+    const fullTableName = getTableName({ indexName: tableName, schemaName });
+    const columnList = columns.map((c) => `"${c}"`).join(", ");
+    if (tableName === TABLE_SPANS) {
+      const updateColumns = columns.filter((c) => c !== "traceId" && c !== "spanId");
+      if (updateColumns.length > 0) {
+        const updateClause = updateColumns.map((c) => `"${c}" = EXCLUDED."${c}"`).join(", ");
+        await client.none(
+          `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
              ON CONFLICT ("traceId", "spanId") DO UPDATE SET ${updateClause}`,
-            values
-          );
-        } else {
-          await this.client.none(
-            `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
-             ON CONFLICT ("traceId", "spanId") DO NOTHING`,
-            values
-          );
-        }
+          values
+        );
       } else {
-        await this.client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
+        await client.none(
+          `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
+             ON CONFLICT ("traceId", "spanId") DO NOTHING`,
+          values
+        );
       }
+    } else {
+      await client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      await this.executeInsert(this.client, { tableName, record });
     } catch (error) {
       throw new MastraError(
         {
@@ -2978,13 +3019,12 @@ Note: This migration may take some time for large tables.
   }
   async batchInsert({ tableName, records }) {
     try {
-      await this.client.query("BEGIN");
-      for (const record of records) {
-        await this.insert({ tableName, record });
-      }
-      await this.client.query("COMMIT");
+      await this.client.tx(async (tx) => {
+        for (const record of records) {
+          await this.executeInsert(tx, { tableName, record });
+        }
+      });
     } catch (error) {
-      await this.client.query("ROLLBACK");
       throw new MastraError(
         {
           id: createStorageErrorId("PG", "BATCH_INSERT", "FAILED"),
@@ -3262,30 +3302,7 @@ Note: This migration may take some time for large tables.
     data
   }) {
     try {
-      const filteredData = await this.filterRecordToKnownColumns(tableName, data);
-      if (Object.keys(filteredData).length === 0) return;
-      const setColumns = [];
-      const setValues = [];
-      let paramIndex = 1;
-      Object.entries(filteredData).forEach(([key, value]) => {
-        const parsedKey = parseSqlIdentifier(key, "column name");
-        setColumns.push(`"${parsedKey}" = $${paramIndex++}`);
-        setValues.push(this.prepareValue(value, key, tableName));
-      });
-      const whereConditions = [];
-      const whereValues = [];
-      Object.entries(keys).forEach(([key, value]) => {
-        const parsedKey = parseSqlIdentifier(key, "column name");
-        whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
-        whereValues.push(this.prepareValue(value, key, tableName));
-      });
-      const tableName_ = getTableName({
-        indexName: tableName,
-        schemaName: getSchemaName(this.schemaName)
-      });
-      const sql = `UPDATE ${tableName_} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
-      const values = [...setValues, ...whereValues];
-      await this.client.none(sql, values);
+      await this.executeUpdate(this.client, { tableName, keys, data });
     } catch (error) {
       throw new MastraError(
         {
@@ -3305,13 +3322,12 @@ Note: This migration may take some time for large tables.
     updates
   }) {
     try {
-      await this.client.query("BEGIN");
-      for (const { keys, data } of updates) {
-        await this.update({ tableName, keys, data });
-      }
-      await this.client.query("COMMIT");
+      await this.client.tx(async (tx) => {
+        for (const { keys, data } of updates) {
+          await this.executeUpdate(tx, { tableName, keys, data });
+        }
+      });
     } catch (error) {
-      await this.client.query("ROLLBACK");
       throw new MastraError(
         {
           id: createStorageErrorId("PG", "BATCH_UPDATE", "FAILED"),
@@ -3370,6 +3386,36 @@ Note: This migration may take some time for large tables.
   async deleteData({ tableName }) {
     return this.clearTable({ tableName });
   }
+  async executeUpdate(client, {
+    tableName,
+    keys,
+    data
+  }) {
+    const filteredData = await this.filterRecordToKnownColumns(tableName, data);
+    if (Object.keys(filteredData).length === 0) return;
+    const setColumns = [];
+    const setValues = [];
+    let paramIndex = 1;
+    Object.entries(filteredData).forEach(([key, value]) => {
+      const parsedKey = parseSqlIdentifier(key, "column name");
+      setColumns.push(`"${parsedKey}" = $${paramIndex++}`);
+      setValues.push(this.prepareValue(value, key, tableName));
+    });
+    const whereConditions = [];
+    const whereValues = [];
+    Object.entries(keys).forEach(([key, value]) => {
+      const parsedKey = parseSqlIdentifier(key, "column name");
+      whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
+      whereValues.push(this.prepareValue(value, key, tableName));
+    });
+    const tableName_ = getTableName({
+      indexName: tableName,
+      schemaName: getSchemaName(this.schemaName)
+    });
+    const sql = `UPDATE ${tableName_} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
+    const values = [...setValues, ...whereValues];
+    await client.none(sql, values);
+  }
 };
 function getSchemaName2(schema) {
   return schema ? `"${parseSqlIdentifier(schema, "schema name")}"` : void 0;
@@ -5283,6 +5329,16 @@ var ExperimentsPG = class _ExperimentsPG extends ExperimentsStorage {
   async init() {
     await this.#db.createTable({ tableName: TABLE_EXPERIMENTS, schema: EXPERIMENTS_SCHEMA });
     await this.#db.createTable({ tableName: TABLE_EXPERIMENT_RESULTS, schema: EXPERIMENT_RESULTS_SCHEMA });
+    await this.#db.alterTable({
+      tableName: TABLE_EXPERIMENTS,
+      schema: EXPERIMENTS_SCHEMA,
+      ifNotExists: ["agentVersion"]
+    });
+    await this.#db.alterTable({
+      tableName: TABLE_EXPERIMENT_RESULTS,
+      schema: EXPERIMENT_RESULTS_SCHEMA,
+      ifNotExists: ["status", "tags"]
+    });
     await this.createDefaultIndexes();
     await this.createCustomIndexes();
   }
@@ -5522,6 +5578,22 @@ var ExperimentsPG = class _ExperimentsPG extends ExperimentsStorage {
         conditions.push(`"datasetId" = $${paramIndex++}`);
         queryParams.push(args.datasetId);
       }
+      if (args.targetType) {
+        conditions.push(`"targetType" = $${paramIndex++}`);
+        queryParams.push(args.targetType);
+      }
+      if (args.targetId) {
+        conditions.push(`"targetId" = $${paramIndex++}`);
+        queryParams.push(args.targetId);
+      }
+      if (args.agentVersion) {
+        conditions.push(`"agentVersion" = $${paramIndex++}`);
+        queryParams.push(args.agentVersion);
+      }
+      if (args.status) {
+        conditions.push(`"status" = $${paramIndex++}`);
+        queryParams.push(args.status);
+      }
       const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
       const countResult = await this.#db.client.one(
         `SELECT COUNT(*) as count FROM ${tableName} ${whereClause}`,
@@ -5710,9 +5782,21 @@ var ExperimentsPG = class _ExperimentsPG extends ExperimentsStorage {
     try {
       const { page, perPage: perPageInput } = args.pagination;
       const tableName = getTableName2({ indexName: TABLE_EXPERIMENT_RESULTS, schemaName: getSchemaName2(this.#schema) });
+      const conditions = ['"experimentId" = $1'];
+      const queryParams = [args.experimentId];
+      let paramIndex = 2;
+      if (args.traceId) {
+        conditions.push(`"traceId" = $${paramIndex++}`);
+        queryParams.push(args.traceId);
+      }
+      if (args.status) {
+        conditions.push(`"status" = $${paramIndex++}`);
+        queryParams.push(args.status);
+      }
+      const whereClause = `WHERE ${conditions.join(" AND ")}`;
       const countResult = await this.#db.client.one(
-        `SELECT COUNT(*) as count FROM ${tableName} WHERE "experimentId" = $1`,
-        [args.experimentId]
+        `SELECT COUNT(*) as count FROM ${tableName} ${whereClause}`,
+        queryParams
       );
       const total = parseInt(countResult.count, 10);
       if (total === 0) {
@@ -5722,8 +5806,8 @@ var ExperimentsPG = class _ExperimentsPG extends ExperimentsStorage {
       const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
       const limitValue = perPageInput === false ? total : perPage;
       const rows = await this.#db.client.manyOrNone(
-        `SELECT * FROM ${tableName} WHERE "experimentId" = $1 ORDER BY "startedAt" ASC LIMIT $2 OFFSET $3`,
-        [args.experimentId, limitValue, offset]
+        `SELECT * FROM ${tableName} ${whereClause} ORDER BY "startedAt" ASC LIMIT $${paramIndex} OFFSET $${paramIndex + 1}`,
+        [...queryParams, limitValue, offset]
       );
       return {
         results: (rows || []).map((row) => this.transformExperimentResultRow(row)),