@mastra/pg 0.3.4 → 0.3.5-alpha.0

package/dist/index.js

@@ -1,3 +1,4 @@
+import { parseSqlIdentifier, parseFieldKey } from '@mastra/core/utils';
 import { MastraVector } from '@mastra/core/vector';
 import { Mutex } from 'async-mutex';
 import pg from 'pg';
@@ -73,22 +74,26 @@ var PGFilterTranslator = class extends 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)) {
@@ -139,46 +144,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
@@ -189,33 +204,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
     };
   },
   /**
@@ -230,29 +252,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 !== "" ? 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") {
@@ -262,7 +291,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);
         }
@@ -381,7 +410,7 @@ var PgVector = class extends 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,
@@ -397,15 +426,19 @@ var PgVector = class extends MastraVector {
     return this.mutexesByName.get(indexName);
   }
   getTableName(indexName) {
-    return this.schema ? `${this.schema}.${indexName}` : indexName;
+    const parsedIndexName = parseSqlIdentifier(indexName, "index name");
+    const parsedSchemaName = this.schema ? 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);
   }
@@ -416,12 +449,18 @@ var PgVector = class extends 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));
@@ -445,7 +484,7 @@ var PgVector = class extends 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,
@@ -708,7 +747,16 @@ var PgVector = class extends 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);
@@ -782,7 +830,9 @@ var PgVector = class extends 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);
@@ -795,7 +845,9 @@ var PgVector = class extends 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);
@@ -828,7 +880,7 @@ var PgVector = class extends 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.
@@ -840,7 +892,9 @@ var PgVector = class extends 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");
     }
@@ -889,7 +943,7 @@ var PgVector = class extends 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.
@@ -898,7 +952,9 @@ var PgVector = class extends 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);
@@ -957,7 +1013,9 @@ var PostgresStore = class extends MastraStorage {
     );
   }
   getTableName(indexName) {
-    return this.schema ? `${this.schema}."${indexName}"` : `"${indexName}"`;
+    const parsedIndexName = parseSqlIdentifier(indexName, "table name");
+    const parsedSchemaName = this.schema ? parseSqlIdentifier(this.schema, "schema name") : void 0;
+    return parsedSchemaName ? `${parsedSchemaName}."${parsedIndexName}"` : `"${parsedIndexName}"`;
   }
   async getEvalsByAgentName(agentName, type) {
     try {
@@ -1032,12 +1090,14 @@ var PostgresStore = class extends MastraStorage {
     }
     if (attributes) {
       Object.keys(attributes).forEach((key) => {
-        conditions.push(`attributes->>'${key}' = $${idx++}`);
+        const parsedKey = parseSqlIdentifier(key, "attribute key");
+        conditions.push(`attributes->>'${parsedKey}' = $${idx++}`);
       });
     }
     if (filters) {
       Object.entries(filters).forEach(([key]) => {
-        conditions.push(`${key} = $${idx++}`);
+        const parsedKey = parseSqlIdentifier(key, "filter key");
+        conditions.push(`${parsedKey} = $${idx++}`);
       });
     }
     if (fromDate) {
@@ -1139,10 +1199,11 @@ var PostgresStore = class extends MastraStorage {
   }) {
     try {
       const columns = Object.entries(schema).map(([name, def]) => {
+        const parsedName = 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();
@@ -1179,7 +1240,7 @@ var PostgresStore = class extends MastraStorage {
   }
   async insert({ tableName, record }) {
     try {
-      const columns = Object.keys(record);
+      const columns = Object.keys(record).map((col) => parseSqlIdentifier(col, "column name"));
       const values = Object.values(record);
       const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
       await this.db.none(
@@ -1193,7 +1254,7 @@ var PostgresStore = class extends MastraStorage {
   }
   async load({ tableName, keys }) {
     try {
-      const keyEntries = Object.entries(keys);
+      const keyEntries = Object.entries(keys).map(([key, value]) => [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(

package/docker-compose.perf.yaml

@@ -3,19 +3,19 @@ services:
     image: pgvector/pgvector:pg16
     container_name: 'pg-perf-test-db'
     ports:
-      - '5435:5432' 
+      - '5435:5432'
     environment:
       POSTGRES_USER: ${POSTGRES_USER:-postgres}
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-postgres}
       POSTGRES_DB: ${POSTGRES_DB:-mastra}
     shm_size: 1gb
     command:
-      - "postgres"
-      - "-c"
-      - "shared_buffers=512MB"
-      - "-c"
-      - "maintenance_work_mem=1024MB"
-      - "-c"
-      - "work_mem=512MB"
+      - 'postgres'
+      - '-c'
+      - 'shared_buffers=512MB'
+      - '-c'
+      - 'maintenance_work_mem=1024MB'
+      - '-c'
+      - 'work_mem=512MB'
     tmpfs:
-      - /var/lib/postgresql/data 
+      - /var/lib/postgresql/data

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/pg",
-  "version": "0.3.4",
+  "version": "0.3.5-alpha.0",
   "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.9.4"
+    "@mastra/core": "^0.9.5-alpha.0"
   },
   "devDependencies": {
     "@microsoft/api-extractor": "^7.52.5",

package/src/storage/index.ts

@@ -16,6 +16,7 @@ import type {
   WorkflowRun,
   WorkflowRuns,
 } from '@mastra/core/storage';
+import { parseSqlIdentifier } from '@mastra/core/utils';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import pgPromise from 'pg-promise';
 import type { ISSLConfig } from 'pg-promise/typescript/pg-subset';
@@ -93,7 +94,9 @@ export class PostgresStore extends MastraStorage {
   }
 
   private getTableName(indexName: string) {
-    return this.schema ? `${this.schema}."${indexName}"` : `"${indexName}"`;
+    const parsedIndexName = parseSqlIdentifier(indexName, 'table name');
+    const parsedSchemaName = this.schema ? parseSqlIdentifier(this.schema, 'schema name') : undefined;
+    return parsedSchemaName ? `${parsedSchemaName}."${parsedIndexName}"` : `"${parsedIndexName}"`;
   }
 
   async getEvalsByAgentName(agentName: string, type?: 'test' | 'live'): Promise<EvalRow[]> {
@@ -192,13 +195,15 @@ export class PostgresStore extends MastraStorage {
     }
     if (attributes) {
       Object.keys(attributes).forEach(key => {
-        conditions.push(`attributes->>'${key}' = \$${idx++}`);
+        const parsedKey = parseSqlIdentifier(key, 'attribute key');
+        conditions.push(`attributes->>'${parsedKey}' = \$${idx++}`);
       });
     }
 
     if (filters) {
       Object.entries(filters).forEach(([key]) => {
-        conditions.push(`${key} = \$${idx++}`);
+        const parsedKey = parseSqlIdentifier(key, 'filter key');
+        conditions.push(`${parsedKey} = \$${idx++}`);
       });
     }
 
@@ -341,10 +346,11 @@ export class PostgresStore extends MastraStorage {
     try {
       const columns = Object.entries(schema)
         .map(([name, def]) => {
+          const parsedName = 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');
 
@@ -392,7 +398,7 @@ export class PostgresStore extends MastraStorage {
 
   async insert({ tableName, record }: { tableName: TABLE_NAMES; record: Record<string, any> }): Promise<void> {
     try {
-      const columns = Object.keys(record);
+      const columns = Object.keys(record).map(col => parseSqlIdentifier(col, 'column name'));
       const values = Object.values(record);
       const placeholders = values.map((_, i) => `$${i + 1}`).join(', ');
 
@@ -408,7 +414,7 @@ export class PostgresStore extends MastraStorage {
 
   async load<R>({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, string> }): Promise<R | null> {
     try {
-      const keyEntries = Object.entries(keys);
+      const keyEntries = Object.entries(keys).map(([key, value]) => [parseSqlIdentifier(key, 'column name'), value]);
       const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(' AND ');
       const values = keyEntries.map(([_, value]) => value);