@mastra/pg 0.3.4 → 0.3.5-alpha.0

package/dist/index.cjs

@@ -1,5 +1,6 @@
 'use strict';
 
+var utils = require('@mastra/core/utils');
 var vector = require('@mastra/core/vector');
 var asyncMutex = require('async-mutex');
 var pg = require('pg');
@@ -81,22 +82,26 @@ var PGFilterTranslator = class extends filter.BaseFilterTranslator {
     return { $regex: flags ? `(?${flags})${pattern}` : pattern };
   }
 };
-
-// src/vector/sql-builder.ts
 var createBasicOperator = (symbol) => {
-  return (key, paramIndex) => ({
-    sql: `CASE 
-      WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${handleKey(key)}}' IS ${symbol === "=" ? "" : "NOT"} NULL
-      ELSE metadata#>>'{${handleKey(key)}}' ${symbol} $${paramIndex}::text
-    END`,
-    needsValue: true
-  });
+  return (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `CASE 
+        WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${jsonPathKey}}' IS ${symbol === "=" ? "" : "NOT"} NULL
+        ELSE metadata#>>'{${jsonPathKey}}' ${symbol} $${paramIndex}::text
+      END`,
+      needsValue: true
+    };
+  };
 };
 var createNumericOperator = (symbol) => {
-  return (key, paramIndex) => ({
-    sql: `(metadata#>>'{${handleKey(key)}}')::numeric ${symbol} $${paramIndex}`,
-    needsValue: true
-  });
+  return (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `(metadata#>>'{${jsonPathKey}}')::numeric ${symbol} $${paramIndex}`,
+      needsValue: true
+    };
+  };
 };
 function buildElemMatchConditions(value, paramIndex) {
   if (typeof value !== "object" || Array.isArray(value)) {
@@ -147,46 +152,56 @@ var FILTER_OPERATORS = {
   $lt: createNumericOperator("<"),
   $lte: createNumericOperator("<="),
   // Array Operators
-  $in: (key, paramIndex) => ({
-    sql: `(
-      CASE
-        WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
-          EXISTS (
-            SELECT 1 FROM jsonb_array_elements_text(metadata->'${handleKey(key)}') as elem
-            WHERE elem = ANY($${paramIndex}::text[])
-          )
-        ELSE metadata#>>'{${handleKey(key)}}' = ANY($${paramIndex}::text[])
-      END
-    )`,
-    needsValue: true
-  }),
-  $nin: (key, paramIndex) => ({
-    sql: `(
-      CASE
-        WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
-          NOT EXISTS (
-            SELECT 1 FROM jsonb_array_elements_text(metadata->'${handleKey(key)}') as elem
-            WHERE elem = ANY($${paramIndex}::text[])
-          )
-        ELSE metadata#>>'{${handleKey(key)}}' != ALL($${paramIndex}::text[])
-      END
-    )`,
-    needsValue: true
-  }),
-  $all: (key, paramIndex) => ({
-    sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false 
-          ELSE (metadata#>'{${handleKey(key)}}')::jsonb ?& $${paramIndex}::text[] END`,
-    needsValue: true
-  }),
+  $in: (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `(
+        CASE
+          WHEN jsonb_typeof(metadata->'${jsonPathKey}') = 'array' THEN
+            EXISTS (
+              SELECT 1 FROM jsonb_array_elements_text(metadata->'${jsonPathKey}') as elem
+              WHERE elem = ANY($${paramIndex}::text[])
+            )
+          ELSE metadata#>>'{${jsonPathKey}}' = ANY($${paramIndex}::text[])
+        END
+      )`,
+      needsValue: true
+    };
+  },
+  $nin: (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `(
+        CASE
+          WHEN jsonb_typeof(metadata->'${jsonPathKey}') = 'array' THEN
+            NOT EXISTS (
+              SELECT 1 FROM jsonb_array_elements_text(metadata->'${jsonPathKey}') as elem
+              WHERE elem = ANY($${paramIndex}::text[])
+            )
+          ELSE metadata#>>'{${jsonPathKey}}' != ALL($${paramIndex}::text[])
+        END
+      )`,
+      needsValue: true
+    };
+  },
+  $all: (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false 
+            ELSE (metadata#>'{${jsonPathKey}}')::jsonb ?& $${paramIndex}::text[] END`,
+      needsValue: true
+    };
+  },
   $elemMatch: (key, paramIndex, value) => {
     const { sql, values } = buildElemMatchConditions(value, paramIndex);
+    const jsonPathKey = parseJsonPathKey(key);
     return {
       sql: `(
         CASE
-          WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
+          WHEN jsonb_typeof(metadata->'${jsonPathKey}') = 'array' THEN
             EXISTS (
               SELECT 1 
-              FROM jsonb_array_elements(metadata->'${handleKey(key)}') as elem
+              FROM jsonb_array_elements(metadata->'${jsonPathKey}') as elem
               WHERE ${sql}
             )
           ELSE FALSE
@@ -197,33 +212,40 @@ var FILTER_OPERATORS = {
     };
   },
   // Element Operators
-  $exists: (key) => ({
-    sql: `metadata ? '${key}'`,
-    needsValue: false
-  }),
+  $exists: (key) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `metadata ? '${jsonPathKey}'`,
+      needsValue: false
+    };
+  },
   // Logical Operators
   $and: (key) => ({ sql: `(${key})`, needsValue: false }),
   $or: (key) => ({ sql: `(${key})`, needsValue: false }),
   $not: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
   $nor: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
   // Regex Operators
-  $regex: (key, paramIndex) => ({
-    sql: `metadata#>>'{${handleKey(key)}}' ~ $${paramIndex}`,
-    needsValue: true
-  }),
+  $regex: (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `metadata#>>'{${jsonPathKey}}' ~ $${paramIndex}`,
+      needsValue: true
+    };
+  },
   $contains: (key, paramIndex, value) => {
+    const jsonPathKey = parseJsonPathKey(key);
     let sql;
     if (Array.isArray(value)) {
-      sql = `(metadata->'${handleKey(key)}') ?& $${paramIndex}`;
+      sql = `(metadata->'${jsonPathKey}') ?& $${paramIndex}`;
     } else if (typeof value === "string") {
-      sql = `metadata->>'${handleKey(key)}' ILIKE '%' || $${paramIndex} || '%'`;
+      sql = `metadata->>'${jsonPathKey}' ILIKE '%' || $${paramIndex} || '%' ESCAPE '\\'`;
     } else {
-      sql = `metadata->>'${handleKey(key)}' = $${paramIndex}`;
+      sql = `metadata->>'${jsonPathKey}' = $${paramIndex}`;
     }
     return {
       sql,
       needsValue: true,
-      transformValue: () => Array.isArray(value) ? value.map(String) : value
+      transformValue: () => Array.isArray(value) ? value.map(String) : typeof value === "string" ? escapeLikePattern(value) : value
     };
   },
   /**
@@ -238,29 +260,36 @@ var FILTER_OPERATORS = {
   //     return JSON.stringify(parts.reduceRight((value, key) => ({ [key]: value }), value));
   //   },
   // }),
-  $size: (key, paramIndex) => ({
-    sql: `(
+  $size: (key, paramIndex) => {
+    const jsonPathKey = parseJsonPathKey(key);
+    return {
+      sql: `(
       CASE
-        WHEN jsonb_typeof(metadata#>'{${handleKey(key)}}') = 'array' THEN 
-          jsonb_array_length(metadata#>'{${handleKey(key)}}') = $${paramIndex}
+        WHEN jsonb_typeof(metadata#>'{${jsonPathKey}}') = 'array' THEN 
+          jsonb_array_length(metadata#>'{${jsonPathKey}}') = $${paramIndex}
         ELSE FALSE
       END
     )`,
-    needsValue: true
-  })
+      needsValue: true
+    };
+  }
 };
-var handleKey = (key) => {
-  return key.replace(/\./g, ",");
+var parseJsonPathKey = (key) => {
+  const parsedKey = key !== "" ? utils.parseFieldKey(key) : "";
+  return parsedKey.replace(/\./g, ",");
 };
-function buildFilterQuery(filter, minScore) {
-  const values = [minScore];
+function escapeLikePattern(str) {
+  return str.replace(/([%_\\])/g, "\\$1");
+}
+function buildFilterQuery(filter, minScore, topK) {
+  const values = [minScore, topK];
   function buildCondition(key, value, parentPath) {
     if (["$and", "$or", "$not", "$nor"].includes(key)) {
       return handleLogicalOperator(key, value);
     }
     if (!value || typeof value !== "object") {
       values.push(value);
-      return `metadata#>>'{${handleKey(key)}}' = $${values.length}`;
+      return `metadata#>>'{${parseJsonPathKey(key)}}' = $${values.length}`;
     }
     const [[operator, operatorValue] = []] = Object.entries(value);
     if (operator === "$not") {
@@ -270,7 +299,7 @@ function buildFilterQuery(filter, minScore) {
           throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
         }
         const operatorFn2 = FILTER_OPERATORS[nestedOp];
-        const operatorResult2 = operatorFn2(key, values.length + 1);
+        const operatorResult2 = operatorFn2(key, values.length + 1, nestedValue);
         if (operatorResult2.needsValue) {
           values.push(nestedValue);
         }
@@ -389,7 +418,7 @@ var PgVector = class extends vector.MastraVector {
     void (async () => {
       const existingIndexes = await this.listIndexes();
       void existingIndexes.map(async (indexName) => {
-        const info = await this.getIndexInfo(indexName);
+        const info = await this.getIndexInfo({ indexName });
         const key = await this.getIndexCacheKey({
           indexName,
           metric: info.metric,
@@ -405,15 +434,19 @@ var PgVector = class extends vector.MastraVector {
     return this.mutexesByName.get(indexName);
   }
   getTableName(indexName) {
-    return this.schema ? `${this.schema}.${indexName}` : indexName;
+    const parsedIndexName = utils.parseSqlIdentifier(indexName, "index name");
+    const parsedSchemaName = this.schema ? utils.parseSqlIdentifier(this.schema, "schema name") : void 0;
+    return parsedSchemaName ? `${parsedSchemaName}.${parsedIndexName}` : parsedIndexName;
   }
   transformFilter(filter) {
     const translator = new PGFilterTranslator();
     return translator.translate(filter);
   }
-  async getIndexInfo(indexName) {
+  async getIndexInfo(...args) {
+    const params = this.normalizeArgs("getIndexInfo", args);
+    const { indexName } = params;
     if (!this.describeIndexCache.has(indexName)) {
-      this.describeIndexCache.set(indexName, await this.describeIndex(indexName));
+      this.describeIndexCache.set(indexName, await this.describeIndex({ indexName }));
     }
     return this.describeIndexCache.get(indexName);
   }
@@ -424,12 +457,18 @@ var PgVector = class extends vector.MastraVector {
       "probes"
     ]);
     const { indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0, ef, probes } = params;
+    if (!Number.isInteger(topK) || topK <= 0) {
+      throw new Error("topK must be a positive integer");
+    }
+    if (!Array.isArray(queryVector) || !queryVector.every((x) => typeof x === "number" && Number.isFinite(x))) {
+      throw new Error("queryVector must be an array of finite numbers");
+    }
     const client = await this.pool.connect();
     try {
       const vectorStr = `[${queryVector.join(",")}]`;
       const translatedFilter = this.transformFilter(filter);
-      const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
-      const indexInfo = await this.getIndexInfo(indexName);
+      const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore, topK);
+      const indexInfo = await this.getIndexInfo({ indexName });
       if (indexInfo.type === "hnsw") {
         const calculatedEf = ef ?? Math.max(topK, (indexInfo?.config?.m ?? 16) * topK);
         const searchEf = Math.min(1e3, Math.max(1, calculatedEf));
@@ -453,7 +492,7 @@ var PgVector = class extends vector.MastraVector {
         FROM vector_scores
         WHERE score > $1
         ORDER BY score DESC
-        LIMIT ${topK}`;
+        LIMIT $2`;
       const result = await client.query(query, filterValues);
       return result.rows.map(({ id, score, metadata, embedding }) => ({
         id,
@@ -716,7 +755,16 @@ var PgVector = class extends vector.MastraVector {
       client.release();
     }
   }
-  async describeIndex(indexName) {
+  /**
+   * Retrieves statistics about a vector index.
+   *
+   * @param params - The parameters for describing an index
+   * @param params.indexName - The name of the index to describe
+   * @returns A promise that resolves to the index statistics including dimension, count and metric
+   */
+  async describeIndex(...args) {
+    const params = this.normalizeArgs("describeIndex", args);
+    const { indexName } = params;
     const client = await this.pool.connect();
     try {
       const tableName = this.getTableName(indexName);
@@ -790,7 +838,9 @@ var PgVector = class extends vector.MastraVector {
       client.release();
     }
   }
-  async deleteIndex(indexName) {
+  async deleteIndex(...args) {
+    const params = this.normalizeArgs("deleteIndex", args);
+    const { indexName } = params;
     const client = await this.pool.connect();
     try {
       const tableName = this.getTableName(indexName);
@@ -803,7 +853,9 @@ var PgVector = class extends vector.MastraVector {
       client.release();
     }
   }
-  async truncateIndex(indexName) {
+  async truncateIndex(...args) {
+    const params = this.normalizeArgs("truncateIndex", args);
+    const { indexName } = params;
     const client = await this.pool.connect();
     try {
       const tableName = this.getTableName(indexName);
@@ -836,7 +888,7 @@ var PgVector = class extends vector.MastraVector {
       Please use updateVector() instead. 
       updateIndexById() will be removed on May 20th, 2025.`
     );
-    await this.updateVector(indexName, id, update);
+    await this.updateVector({ indexName, id, update });
   }
   /**
    * Updates a vector by its ID with the provided vector and/or metadata.
@@ -848,7 +900,9 @@ var PgVector = class extends vector.MastraVector {
    * @returns A promise that resolves when the update is complete.
    * @throws Will throw an error if no updates are provided or if the update operation fails.
    */
-  async updateVector(indexName, id, update) {
+  async updateVector(...args) {
+    const params = this.normalizeArgs("updateVector", args);
+    const { indexName, id, update } = params;
     if (!update.vector && !update.metadata) {
       throw new Error("No updates provided");
     }
@@ -897,7 +951,7 @@ var PgVector = class extends vector.MastraVector {
       Please use deleteVector() instead. 
       deleteIndexById() will be removed on May 20th, 2025.`
     );
-    await this.deleteVector(indexName, id);
+    await this.deleteVector({ indexName, id });
   }
   /**
    * Deletes a vector by its ID.
@@ -906,7 +960,9 @@ var PgVector = class extends vector.MastraVector {
    * @returns A promise that resolves when the deletion is complete.
    * @throws Will throw an error if the deletion operation fails.
    */
-  async deleteVector(indexName, id) {
+  async deleteVector(...args) {
+    const params = this.normalizeArgs("deleteVector", args);
+    const { indexName, id } = params;
     const client = await this.pool.connect();
     try {
       const tableName = this.getTableName(indexName);
@@ -965,7 +1021,9 @@ var PostgresStore = class extends storage.MastraStorage {
     );
   }
   getTableName(indexName) {
-    return this.schema ? `${this.schema}."${indexName}"` : `"${indexName}"`;
+    const parsedIndexName = utils.parseSqlIdentifier(indexName, "table name");
+    const parsedSchemaName = this.schema ? utils.parseSqlIdentifier(this.schema, "schema name") : void 0;
+    return parsedSchemaName ? `${parsedSchemaName}."${parsedIndexName}"` : `"${parsedIndexName}"`;
   }
   async getEvalsByAgentName(agentName, type) {
     try {
@@ -1040,12 +1098,14 @@ var PostgresStore = class extends storage.MastraStorage {
     }
     if (attributes) {
       Object.keys(attributes).forEach((key) => {
-        conditions.push(`attributes->>'${key}' = $${idx++}`);
+        const parsedKey = utils.parseSqlIdentifier(key, "attribute key");
+        conditions.push(`attributes->>'${parsedKey}' = $${idx++}`);
       });
     }
     if (filters) {
       Object.entries(filters).forEach(([key]) => {
-        conditions.push(`${key} = $${idx++}`);
+        const parsedKey = utils.parseSqlIdentifier(key, "filter key");
+        conditions.push(`${parsedKey} = $${idx++}`);
       });
     }
     if (fromDate) {
@@ -1147,10 +1207,11 @@ var PostgresStore = class extends storage.MastraStorage {
   }) {
     try {
       const columns = Object.entries(schema).map(([name, def]) => {
+        const parsedName = utils.parseSqlIdentifier(name, "column name");
         const constraints = [];
         if (def.primaryKey) constraints.push("PRIMARY KEY");
         if (!def.nullable) constraints.push("NOT NULL");
-        return `"${name}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
+        return `"${parsedName}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
       }).join(",\n");
       if (this.schema) {
         await this.setupSchema();
@@ -1187,7 +1248,7 @@ var PostgresStore = class extends storage.MastraStorage {
   }
   async insert({ tableName, record }) {
     try {
-      const columns = Object.keys(record);
+      const columns = Object.keys(record).map((col) => utils.parseSqlIdentifier(col, "column name"));
       const values = Object.values(record);
       const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
       await this.db.none(
@@ -1201,7 +1262,7 @@ var PostgresStore = class extends storage.MastraStorage {
   }
   async load({ tableName, keys }) {
     try {
-      const keyEntries = Object.entries(keys);
+      const keyEntries = Object.entries(keys).map(([key, value]) => [utils.parseSqlIdentifier(key, "column name"), value]);
       const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
       const values = keyEntries.map(([_, value]) => value);
       const result = await this.db.oneOrNone(