@mastra/pg 0.15.2 → 0.15.3-alpha.1

package/dist/index.js

@@ -5,7 +5,7 @@ import { Mutex } from 'async-mutex';
 import pg from 'pg';
 import xxhash from 'xxhash-wasm';
 import { BaseFilterTranslator } from '@mastra/core/vector/filter';
-import { MastraStorage, StoreOperations, TABLE_WORKFLOW_SNAPSHOT, ScoresStorage, TABLE_SCORERS, TracesStorage, safelyParseJSON, TABLE_TRACES, WorkflowsStorage, LegacyEvalsStorage, MemoryStorage, resolveMessageLimit, TABLE_RESOURCES, TABLE_EVALS, TABLE_THREADS, TABLE_MESSAGES } from '@mastra/core/storage';
+import { MastraStorage, StoreOperations, TABLE_WORKFLOW_SNAPSHOT, TABLE_THREADS, TABLE_MESSAGES, TABLE_TRACES, TABLE_EVALS, ScoresStorage, TABLE_SCORERS, TracesStorage, safelyParseJSON, WorkflowsStorage, LegacyEvalsStorage, MemoryStorage, resolveMessageLimit, TABLE_RESOURCES } from '@mastra/core/storage';
 import pgPromise from 'pg-promise';
 import { MessageList } from '@mastra/core/agent';
 
@@ -2366,6 +2366,303 @@ var StoreOperationsPG = class extends StoreOperations {
       );
     }
   }
+  /**
+   * Create a new index on a table
+   */
+  async createIndex(options) {
+    try {
+      const {
+        name,
+        table,
+        columns,
+        unique = false,
+        concurrent = true,
+        where,
+        method = "btree",
+        opclass,
+        storage,
+        tablespace
+      } = options;
+      const schemaName = this.schemaName || "public";
+      const fullTableName = getTableName({
+        indexName: table,
+        schemaName: getSchemaName(this.schemaName)
+      });
+      const indexExists = await this.client.oneOrNone(
+        `SELECT 1 FROM pg_indexes 
+         WHERE indexname = $1 
+         AND schemaname = $2`,
+        [name, schemaName]
+      );
+      if (indexExists) {
+        return;
+      }
+      const uniqueStr = unique ? "UNIQUE " : "";
+      const concurrentStr = concurrent ? "CONCURRENTLY " : "";
+      const methodStr = method !== "btree" ? `USING ${method} ` : "";
+      const columnsStr = columns.map((col) => {
+        if (col.includes(" DESC") || col.includes(" ASC")) {
+          const [colName, ...modifiers] = col.split(" ");
+          if (!colName) {
+            throw new Error(`Invalid column specification: ${col}`);
+          }
+          const quotedCol2 = `"${parseSqlIdentifier(colName, "column name")}" ${modifiers.join(" ")}`;
+          return opclass ? `${quotedCol2} ${opclass}` : quotedCol2;
+        }
+        const quotedCol = `"${parseSqlIdentifier(col, "column name")}"`;
+        return opclass ? `${quotedCol} ${opclass}` : quotedCol;
+      }).join(", ");
+      const whereStr = where ? ` WHERE ${where}` : "";
+      const tablespaceStr = tablespace ? ` TABLESPACE ${tablespace}` : "";
+      let withStr = "";
+      if (storage && Object.keys(storage).length > 0) {
+        const storageParams = Object.entries(storage).map(([key, value]) => `${key} = ${value}`).join(", ");
+        withStr = ` WITH (${storageParams})`;
+      }
+      const sql = `CREATE ${uniqueStr}INDEX ${concurrentStr}${name} ON ${fullTableName} ${methodStr}(${columnsStr})${withStr}${tablespaceStr}${whereStr}`;
+      await this.client.none(sql);
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("CONCURRENTLY")) {
+        const retryOptions = { ...options, concurrent: false };
+        return this.createIndex(retryOptions);
+      }
+      throw new MastraError(
+        {
+          id: "MASTRA_STORAGE_PG_INDEX_CREATE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            indexName: options.name,
+            tableName: options.table
+          }
+        },
+        error
+      );
+    }
+  }
+  /**
+   * Drop an existing index
+   */
+  async dropIndex(indexName) {
+    try {
+      const schemaName = this.schemaName || "public";
+      const indexExists = await this.client.oneOrNone(
+        `SELECT 1 FROM pg_indexes 
+         WHERE indexname = $1 
+         AND schemaname = $2`,
+        [indexName, schemaName]
+      );
+      if (!indexExists) {
+        return;
+      }
+      const sql = `DROP INDEX IF EXISTS ${getSchemaName(this.schemaName)}.${indexName}`;
+      await this.client.none(sql);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MASTRA_STORAGE_PG_INDEX_DROP_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            indexName
+          }
+        },
+        error
+      );
+    }
+  }
+  /**
+   * List indexes for a specific table or all tables
+   */
+  async listIndexes(tableName) {
+    try {
+      const schemaName = this.schemaName || "public";
+      let query;
+      let params;
+      if (tableName) {
+        query = `
+          SELECT 
+            i.indexname as name,
+            i.tablename as table,
+            i.indexdef as definition,
+            ix.indisunique as is_unique,
+            pg_size_pretty(pg_relation_size(c.oid)) as size,
+            array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
+          FROM pg_indexes i
+          JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
+          JOIN pg_index ix ON ix.indexrelid = c.oid
+          JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
+          WHERE i.schemaname = $1 
+          AND i.tablename = $2
+          GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid
+        `;
+        params = [schemaName, tableName];
+      } else {
+        query = `
+          SELECT 
+            i.indexname as name,
+            i.tablename as table,
+            i.indexdef as definition,
+            ix.indisunique as is_unique,
+            pg_size_pretty(pg_relation_size(c.oid)) as size,
+            array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
+          FROM pg_indexes i
+          JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
+          JOIN pg_index ix ON ix.indexrelid = c.oid
+          JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
+          WHERE i.schemaname = $1
+          GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid
+        `;
+        params = [schemaName];
+      }
+      const results = await this.client.manyOrNone(query, params);
+      return results.map((row) => {
+        let columns = [];
+        if (typeof row.columns === "string" && row.columns.startsWith("{") && row.columns.endsWith("}")) {
+          const arrayContent = row.columns.slice(1, -1);
+          columns = arrayContent ? arrayContent.split(",") : [];
+        } else if (Array.isArray(row.columns)) {
+          columns = row.columns;
+        }
+        return {
+          name: row.name,
+          table: row.table,
+          columns,
+          unique: row.is_unique || false,
+          size: row.size || "0",
+          definition: row.definition || ""
+        };
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MASTRA_STORAGE_PG_INDEX_LIST_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: tableName ? {
+            tableName
+          } : {}
+        },
+        error
+      );
+    }
+  }
+  /**
+   * Creates automatic indexes for optimal query performance
+   * These composite indexes cover both filtering and sorting in single index
+   */
+  async createAutomaticIndexes() {
+    try {
+      const schemaPrefix = this.schemaName ? `${this.schemaName}_` : "";
+      const indexes = [
+        // Composite index for threads (filter + sort)
+        {
+          name: `${schemaPrefix}mastra_threads_resourceid_createdat_idx`,
+          table: TABLE_THREADS,
+          columns: ["resourceId", "createdAt DESC"]
+        },
+        // Composite index for messages (filter + sort)
+        {
+          name: `${schemaPrefix}mastra_messages_thread_id_createdat_idx`,
+          table: TABLE_MESSAGES,
+          columns: ["thread_id", "createdAt DESC"]
+        },
+        // Composite index for traces (filter + sort)
+        {
+          name: `${schemaPrefix}mastra_traces_name_starttime_idx`,
+          table: TABLE_TRACES,
+          columns: ["name", "startTime DESC"]
+        },
+        // Composite index for evals (filter + sort)
+        {
+          name: `${schemaPrefix}mastra_evals_agent_name_created_at_idx`,
+          table: TABLE_EVALS,
+          columns: ["agent_name", "created_at DESC"]
+        }
+      ];
+      for (const indexOptions of indexes) {
+        try {
+          await this.createIndex(indexOptions);
+        } catch (error) {
+          this.logger?.warn?.(`Failed to create index ${indexOptions.name}:`, error);
+        }
+      }
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MASTRA_STORAGE_PG_STORE_CREATE_PERFORMANCE_INDEXES_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY
+        },
+        error
+      );
+    }
+  }
+  /**
+   * Get detailed statistics for a specific index
+   */
+  async describeIndex(indexName) {
+    try {
+      const schemaName = this.schemaName || "public";
+      const query = `
+        SELECT 
+          i.indexname as name,
+          i.tablename as table,
+          i.indexdef as definition,
+          ix.indisunique as is_unique,
+          pg_size_pretty(pg_relation_size(c.oid)) as size,
+          array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns,
+          am.amname as method,
+          s.idx_scan as scans,
+          s.idx_tup_read as tuples_read,
+          s.idx_tup_fetch as tuples_fetched
+        FROM pg_indexes i
+        JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
+        JOIN pg_index ix ON ix.indexrelid = c.oid
+        JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
+        JOIN pg_am am ON c.relam = am.oid
+        LEFT JOIN pg_stat_user_indexes s ON s.indexrelname = i.indexname AND s.schemaname = i.schemaname
+        WHERE i.schemaname = $1 
+        AND i.indexname = $2
+        GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid, am.amname, s.idx_scan, s.idx_tup_read, s.idx_tup_fetch
+      `;
+      const result = await this.client.oneOrNone(query, [schemaName, indexName]);
+      if (!result) {
+        throw new Error(`Index "${indexName}" not found in schema "${schemaName}"`);
+      }
+      let columns = [];
+      if (typeof result.columns === "string" && result.columns.startsWith("{") && result.columns.endsWith("}")) {
+        const arrayContent = result.columns.slice(1, -1);
+        columns = arrayContent ? arrayContent.split(",") : [];
+      } else if (Array.isArray(result.columns)) {
+        columns = result.columns;
+      }
+      return {
+        name: result.name,
+        table: result.table,
+        columns,
+        unique: result.is_unique || false,
+        size: result.size || "0",
+        definition: result.definition || "",
+        method: result.method || "btree",
+        scans: parseInt(result.scans) || 0,
+        tuples_read: parseInt(result.tuples_read) || 0,
+        tuples_fetched: parseInt(result.tuples_fetched) || 0
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MASTRA_STORAGE_PG_INDEX_DESCRIBE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            indexName
+          }
+        },
+        error
+      );
+    }
+  }
 };
 function transformScoreRow(row) {
   return {
@@ -2800,16 +3097,17 @@ var WorkflowsPG = class extends WorkflowsStorage {
   async persistWorkflowSnapshot({
     workflowName,
     runId,
+    resourceId,
     snapshot
   }) {
     try {
       const now = (/* @__PURE__ */ new Date()).toISOString();
       await this.client.none(
-        `INSERT INTO ${getTableName({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: this.schema })} (workflow_name, run_id, snapshot, "createdAt", "updatedAt")
-                 VALUES ($1, $2, $3, $4, $5)
+        `INSERT INTO ${getTableName({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: this.schema })} (workflow_name, run_id, "resourceId", snapshot, "createdAt", "updatedAt")
+                 VALUES ($1, $2, $3, $4, $5, $6)
                  ON CONFLICT (workflow_name, run_id) DO UPDATE
-                 SET snapshot = $3, "updatedAt" = $5`,
-        [workflowName, runId, JSON.stringify(snapshot), now, now]
+                 SET "resourceId" = $3, snapshot = $4, "updatedAt" = $6`,
+        [workflowName, runId, resourceId, JSON.stringify(snapshot), now, now]
       );
     } catch (error) {
       throw new MastraError(
@@ -2971,37 +3269,68 @@ var PostgresStore = class extends MastraStorage {
   isConnected = false;
   stores;
   constructor(config) {
+    const isConnectionStringConfig = (cfg) => {
+      return "connectionString" in cfg;
+    };
+    const isHostConfig = (cfg) => {
+      return "host" in cfg && "database" in cfg && "user" in cfg && "password" in cfg;
+    };
+    const isCloudSqlConfig = (cfg) => {
+      return "stream" in cfg || "password" in cfg && typeof cfg.password === "function";
+    };
     try {
-      if ("connectionString" in config) {
+      if (isConnectionStringConfig(config)) {
         if (!config.connectionString || typeof config.connectionString !== "string" || config.connectionString.trim() === "") {
           throw new Error(
             "PostgresStore: connectionString must be provided and cannot be empty. Passing an empty string may cause fallback to local Postgres defaults."
           );
         }
-      } else {
+      } else if (isCloudSqlConfig(config)) ; else if (isHostConfig(config)) {
         const required = ["host", "database", "user", "password"];
         for (const key of required) {
-          if (!(key in config) || typeof config[key] !== "string" || config[key].trim() === "") {
+          if (!config[key] || typeof config[key] !== "string" || config[key].trim() === "") {
             throw new Error(
               `PostgresStore: ${key} must be provided and cannot be empty. Passing an empty string may cause fallback to local Postgres defaults.`
             );
           }
         }
+      } else {
+        throw new Error(
+          "PostgresStore: invalid config. Provide either {connectionString}, {host,port,database,user,password}, or a pg ClientConfig (e.g., Cloud SQL connector with `stream`)."
+        );
       }
       super({ name: "PostgresStore" });
       this.schema = config.schemaName || "public";
-      this.#config = {
-        max: config.max,
-        idleTimeoutMillis: config.idleTimeoutMillis,
-        ...`connectionString` in config ? { connectionString: config.connectionString } : {
+      if (isConnectionStringConfig(config)) {
+        this.#config = {
+          connectionString: config.connectionString,
+          max: config.max,
+          idleTimeoutMillis: config.idleTimeoutMillis
+        };
+      } else if (isCloudSqlConfig(config)) {
+        this.#config = {
+          ...config,
+          max: config.max,
+          idleTimeoutMillis: config.idleTimeoutMillis
+        };
+      } else if (isHostConfig(config)) {
+        this.#config = {
           host: config.host,
           port: config.port,
           database: config.database,
           user: config.user,
           password: config.password,
-          ssl: config.ssl
-        }
-      };
+          ssl: config.ssl,
+          max: config.max,
+          idleTimeoutMillis: config.idleTimeoutMillis
+        };
+      } else {
+        this.#config = {
+          ...config,
+          max: config.max,
+          idleTimeoutMillis: config.idleTimeoutMillis
+        };
+      }
       this.stores = {};
     } catch (e) {
       throw new MastraError(
@@ -3037,6 +3366,11 @@ var PostgresStore = class extends MastraStorage {
         memory
       };
       await super.init();
+      try {
+        await operations.createAutomaticIndexes();
+      } catch (indexError) {
+        console.warn("Failed to create indexes:", indexError);
+      }
     } catch (error) {
       this.isConnected = false;
       throw new MastraError(
@@ -3067,7 +3401,9 @@ var PostgresStore = class extends MastraStorage {
       resourceWorkingMemory: true,
       hasColumn: true,
       createTable: true,
-      deleteMessages: true
+      deleteMessages: true,
+      aiTracing: false,
+      indexManagement: true
     };
   }
   /** @deprecated use getEvals instead */
@@ -3203,9 +3539,10 @@ var PostgresStore = class extends MastraStorage {
   async persistWorkflowSnapshot({
     workflowName,
     runId,
+    resourceId,
     snapshot
   }) {
-    return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, snapshot });
+    return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot });
   }
   async loadWorkflowSnapshot({
     workflowName,