@mastra/pg 0.3.1-alpha.2 → 0.3.1-alpha.4

package/src/storage/index.ts

@@ -8,7 +8,14 @@ import {
   TABLE_WORKFLOW_SNAPSHOT,
   TABLE_EVALS,
 } from '@mastra/core/storage';
-import type { EvalRow, StorageColumn, StorageGetMessagesArg, TABLE_NAMES } from '@mastra/core/storage';
+import type {
+  EvalRow,
+  StorageColumn,
+  StorageGetMessagesArg,
+  TABLE_NAMES,
+  WorkflowRun,
+  WorkflowRuns,
+} from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import pgPromise from 'pg-promise';
 import type { ISSLConfig } from 'pg-promise/typescript/pg-subset';
@@ -561,7 +568,7 @@ export class PostgresStore extends MastraStorage {
     }
   }
 
-  async getMessages<T = unknown>({ threadId, selectBy }: StorageGetMessagesArg): Promise<T> {
+  async getMessages<T = unknown>({ threadId, selectBy }: StorageGetMessagesArg): Promise<T[]> {
     try {
       const messages: any[] = [];
       const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
@@ -645,7 +652,7 @@ export class PostgresStore extends MastraStorage {
         }
       });
 
-      return messages as T;
+      return messages as T[];
     } catch (error) {
       console.error('Error getting messages:', error);
       throw error;
@@ -748,96 +755,166 @@ export class PostgresStore extends MastraStorage {
     }
   }
 
+  private async hasColumn(table: string, column: string): Promise<boolean> {
+    // Use this.schema to scope the check
+    const schema = this.schema || 'public';
+    const result = await this.db.oneOrNone(
+      `SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND (column_name = $3 OR column_name = $4)`,
+      [schema, table, column, column.toLowerCase()],
+    );
+    return !!result;
+  }
+
+  private parseWorkflowRun(row: any): WorkflowRun {
+    let parsedSnapshot: WorkflowRunState | string = row.snapshot as string;
+    if (typeof parsedSnapshot === 'string') {
+      try {
+        parsedSnapshot = JSON.parse(row.snapshot as string) as WorkflowRunState;
+      } catch (e) {
+        // If parsing fails, return the raw snapshot string
+        console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+      }
+    }
+
+    return {
+      workflowName: row.workflow_name,
+      runId: row.run_id,
+      snapshot: parsedSnapshot,
+      createdAt: row.createdAt,
+      updatedAt: row.updatedAt,
+      resourceId: row.resourceId,
+    };
+  }
+
   async getWorkflowRuns({
     workflowName,
     fromDate,
     toDate,
     limit,
     offset,
+    resourceId,
   }: {
     workflowName?: string;
     fromDate?: Date;
     toDate?: Date;
     limit?: number;
     offset?: number;
-  } = {}): Promise<{
-    runs: Array<{
-      workflowName: string;
-      runId: string;
-      snapshot: WorkflowRunState | string;
-      createdAt: Date;
-      updatedAt: Date;
-    }>;
-    total: number;
-  }> {
-    const conditions: string[] = [];
-    const values: any[] = [];
-    let paramIndex = 1;
-
-    if (workflowName) {
-      conditions.push(`workflow_name = $${paramIndex}`);
-      values.push(workflowName);
-      paramIndex++;
-    }
-
-    if (fromDate) {
-      conditions.push(`"createdAt" >= $${paramIndex}`);
-      values.push(fromDate);
-      paramIndex++;
-    }
+    resourceId?: string;
+  } = {}): Promise<WorkflowRuns> {
+    try {
+      const conditions: string[] = [];
+      const values: any[] = [];
+      let paramIndex = 1;
+
+      if (workflowName) {
+        conditions.push(`workflow_name = $${paramIndex}`);
+        values.push(workflowName);
+        paramIndex++;
+      }
 
-    if (toDate) {
-      conditions.push(`"createdAt" <= $${paramIndex}`);
-      values.push(toDate);
-      paramIndex++;
-    }
+      if (resourceId) {
+        const hasResourceId = await this.hasColumn(TABLE_WORKFLOW_SNAPSHOT, 'resourceId');
+        if (hasResourceId) {
+          conditions.push(`"resourceId" = $${paramIndex}`);
+          values.push(resourceId);
+          paramIndex++;
+        } else {
+          console.warn(`[${TABLE_WORKFLOW_SNAPSHOT}] resourceId column not found. Skipping resourceId filter.`);
+        }
+      }
 
-    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+      if (fromDate) {
+        conditions.push(`"createdAt" >= $${paramIndex}`);
+        values.push(fromDate);
+        paramIndex++;
+      }
 
-    let total = 0;
-    // Only get total count when using pagination
-    if (limit !== undefined && offset !== undefined) {
-      const countResult = await this.db.one(
-        `SELECT COUNT(*) as count FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} ${whereClause}`,
-        values,
-      );
-      total = Number(countResult.count);
-    }
+      if (toDate) {
+        conditions.push(`"createdAt" <= $${paramIndex}`);
+        values.push(toDate);
+        paramIndex++;
+      }
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+      let total = 0;
+      // Only get total count when using pagination
+      if (limit !== undefined && offset !== undefined) {
+        const countResult = await this.db.one(
+          `SELECT COUNT(*) as count FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} ${whereClause}`,
+          values,
+        );
+        total = Number(countResult.count);
+      }
 
-    // Get results
-    const query = `
+      // Get results
+      const query = `
       SELECT * FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} 
       ${whereClause} 
       ORDER BY "createdAt" DESC
       ${limit !== undefined && offset !== undefined ? ` LIMIT $${paramIndex} OFFSET $${paramIndex + 1}` : ''}
     `;
 
-    const queryValues = limit !== undefined && offset !== undefined ? [...values, limit, offset] : values;
+      const queryValues = limit !== undefined && offset !== undefined ? [...values, limit, offset] : values;
 
-    const result = await this.db.manyOrNone(query, queryValues);
+      const result = await this.db.manyOrNone(query, queryValues);
 
-    const runs = (result || []).map(row => {
-      let parsedSnapshot: WorkflowRunState | string = row.snapshot as string;
-      if (typeof parsedSnapshot === 'string') {
-        try {
-          parsedSnapshot = JSON.parse(row.snapshot as string) as WorkflowRunState;
-        } catch (e) {
-          // If parsing fails, return the raw snapshot string
-          console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
-        }
+      const runs = (result || []).map(row => {
+        return this.parseWorkflowRun(row);
+      });
+
+      // Use runs.length as total when not paginating
+      return { runs, total: total || runs.length };
+    } catch (error) {
+      console.error('Error getting workflow runs:', error);
+      throw error;
+    }
+  }
+
+  async getWorkflowRunById({
+    runId,
+    workflowName,
+  }: {
+    runId: string;
+    workflowName?: string;
+  }): Promise<WorkflowRun | null> {
+    try {
+      const conditions: string[] = [];
+      const values: any[] = [];
+      let paramIndex = 1;
+
+      if (runId) {
+        conditions.push(`run_id = $${paramIndex}`);
+        values.push(runId);
+        paramIndex++;
       }
 
-      return {
-        workflowName: row.workflow_name,
-        runId: row.run_id,
-        snapshot: parsedSnapshot,
-        createdAt: row.createdAt,
-        updatedAt: row.updatedAt,
-      };
-    });
+      if (workflowName) {
+        conditions.push(`workflow_name = $${paramIndex}`);
+        values.push(workflowName);
+        paramIndex++;
+      }
+
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+      // Get results
+      const query = `
+      SELECT * FROM ${this.getTableName(TABLE_WORKFLOW_SNAPSHOT)} 
+      ${whereClause} 
+    `;
+
+      const queryValues = values;
 
-    // Use runs.length as total when not paginating
-    return { runs, total: total || runs.length };
+      const result = await this.db.oneOrNone(query, queryValues);
+
+      if (!result) {
+        return null;
+      }
+
+      return this.parseWorkflowRun(result);
+    } catch (error) {
+      console.error('Error getting workflow run by ID:', error);
+      throw error;
+    }
   }
 
   async close(): Promise<void> {

package/src/vector/index.test.ts

@@ -2436,4 +2436,45 @@ describe('PgVector', () => {
       });
     });
   });
+
+  describe('PoolConfig Custom Options', () => {
+    it('should apply custom values to properties with default values', async () => {
+      const db = new PgVector({
+        connectionString,
+        pgPoolOptions: {
+          max: 5,
+          idleTimeoutMillis: 10000,
+          connectionTimeoutMillis: 1000,
+        },
+      });
+
+      expect(db['pool'].options.max).toBe(5);
+      expect(db['pool'].options.idleTimeoutMillis).toBe(10000);
+      expect(db['pool'].options.connectionTimeoutMillis).toBe(1000);
+    });
+
+    it('should pass properties with no default values', async () => {
+      const db = new PgVector({
+        connectionString,
+        pgPoolOptions: {
+          ssl: false,
+        },
+      });
+
+      expect(db['pool'].options.ssl).toBe(false);
+    });
+    it('should keep default values when custom values are added', async () => {
+      const db = new PgVector({
+        connectionString,
+        pgPoolOptions: {
+          ssl: false,
+        },
+      });
+
+      expect(db['pool'].options.max).toBe(20);
+      expect(db['pool'].options.idleTimeoutMillis).toBe(30000);
+      expect(db['pool'].options.connectionTimeoutMillis).toBe(2000);
+      expect(db['pool'].options.ssl).toBe(false);
+    });
+  });
 });

package/src/vector/index.ts

@@ -71,23 +71,49 @@ export class PgVector extends MastraVector {
   private schemaSetupComplete: boolean | undefined = undefined;
 
   constructor(connectionString: string);
-  constructor(config: { connectionString: string; schemaName?: string });
-  constructor(config: string | { connectionString: string; schemaName?: string }) {
-    const connectionString = typeof config === 'string' ? config : config.connectionString;
-    if (!connectionString || typeof connectionString !== 'string' || connectionString.trim() === '') {
+  constructor(config: {
+    connectionString: string;
+    schemaName?: string;
+    pgPoolOptions?: Omit<pg.PoolConfig, 'connectionString'>;
+  });
+  constructor(
+    config:
+      | string
+      | {
+          connectionString: string;
+          schemaName?: string;
+          pgPoolOptions?: Omit<pg.PoolConfig, 'connectionString'>;
+        },
+  ) {
+    let connectionString: string;
+    let pgPoolOptions: Omit<pg.PoolConfig, 'connectionString'> | undefined;
+    let schemaName: string | undefined;
+
+    if (typeof config === 'string') {
+      connectionString = config;
+      schemaName = undefined;
+      pgPoolOptions = undefined;
+    } else {
+      connectionString = config.connectionString;
+      schemaName = config.schemaName;
+      pgPoolOptions = config.pgPoolOptions;
+    }
+
+    if (!connectionString || connectionString.trim() === '') {
       throw new Error(
         'PgVector: connectionString must be provided and cannot be empty. Passing an empty string may cause fallback to local Postgres defaults.',
       );
     }
     super();
 
-    this.schema = typeof config === 'string' ? undefined : config.schemaName;
+    this.schema = schemaName;
 
     const basePool = new pg.Pool({
       connectionString,
       max: 20, // Maximum number of clients in the pool
       idleTimeoutMillis: 30000, // Close idle connections after 30 seconds
       connectionTimeoutMillis: 2000, // Fail fast if can't connect
+      ...pgPoolOptions,
     });
 
     const telemetry = this.__getTelemetry();