@mastra/pg 0.12.3 → 0.12.4

package/src/storage/domains/traces/index.ts

@@ -0,0 +1,160 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { PaginationInfo, StorageGetTracesArg, StorageGetTracesPaginatedArg } from '@mastra/core/storage';
+import { TABLE_TRACES, TracesStorage, safelyParseJSON } from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
+import { parseFieldKey } from '@mastra/core/utils';
+import type { IDatabase } from 'pg-promise';
+import type { StoreOperationsPG } from '../operations';
+import { getSchemaName, getTableName } from '../utils';
+
+export class TracesPG extends TracesStorage {
+  public client: IDatabase<{}>;
+  private operations: StoreOperationsPG;
+  private schema?: string;
+
+  constructor({
+    client,
+    operations,
+    schema,
+  }: {
+    client: IDatabase<{}>;
+    operations: StoreOperationsPG;
+    schema?: string;
+  }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+    this.schema = schema;
+  }
+
+  async getTraces(args: StorageGetTracesArg): Promise<Trace[]> {
+    if (args.fromDate || args.toDate) {
+      (args as any).dateRange = {
+        start: args.fromDate,
+        end: args.toDate,
+      };
+    }
+    try {
+      const result = await this.getTracesPaginated(args);
+      return result.traces;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_TRACES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async getTracesPaginated(args: StorageGetTracesPaginatedArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+    const currentOffset = page * perPage;
+
+    const queryParams: any[] = [];
+    const conditions: string[] = [];
+    let paramIndex = 1;
+
+    if (name) {
+      conditions.push(`name LIKE $${paramIndex++}`);
+      queryParams.push(`${name}%`);
+    }
+    if (scope) {
+      conditions.push(`scope = $${paramIndex++}`);
+      queryParams.push(scope);
+    }
+    if (attributes) {
+      Object.entries(attributes).forEach(([key, value]) => {
+        const parsedKey = parseFieldKey(key);
+        conditions.push(`attributes->>'${parsedKey}' = $${paramIndex++}`);
+        queryParams.push(value);
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        const parsedKey = parseFieldKey(key);
+        conditions.push(`"${parsedKey}" = $${paramIndex++}`);
+        queryParams.push(value);
+      });
+    }
+    if (fromDate) {
+      conditions.push(`"createdAt" >= $${paramIndex++}`);
+      queryParams.push(fromDate);
+    }
+    if (toDate) {
+      conditions.push(`"createdAt" <= $${paramIndex++}`);
+      queryParams.push(toDate);
+    }
+
+    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+    try {
+      const countResult = await this.client.oneOrNone<{ count: string }>(
+        `SELECT COUNT(*) FROM ${getTableName({ indexName: TABLE_TRACES, schemaName: getSchemaName(this.schema) })} ${whereClause}`,
+        queryParams,
+      );
+      const total = Number(countResult?.count ?? 0);
+
+      if (total === 0) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      const dataResult = await this.client.manyOrNone<Record<string, any>>(
+        `SELECT * FROM ${getTableName({ indexName: TABLE_TRACES, schemaName: getSchemaName(this.schema) })} ${whereClause} ORDER BY "startTime" DESC LIMIT $${paramIndex++} OFFSET $${paramIndex++}`,
+        [...queryParams, perPage, currentOffset],
+      );
+
+      const traces = dataResult.map(row => ({
+        id: row.id,
+        parentSpanId: row.parentSpanId,
+        traceId: row.traceId,
+        name: row.name,
+        scope: row.scope,
+        kind: row.kind,
+        status: safelyParseJSON(row.status as string),
+        events: safelyParseJSON(row.events as string),
+        links: safelyParseJSON(row.links as string),
+        attributes: safelyParseJSON(row.attributes as string),
+        startTime: row.startTime,
+        endTime: row.endTime,
+        other: safelyParseJSON(row.other as string),
+        createdAt: row.createdAtZ || row.createdAt,
+      })) as Trace[];
+
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + traces.length < total,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async batchTraceInsert({ records }: { records: Record<string, any>[] }): Promise<void> {
+    this.logger.debug('Batch inserting traces', { count: records.length });
+    await this.operations.batchInsert({
+      tableName: TABLE_TRACES,
+      records,
+    });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,12 @@
+import { parseSqlIdentifier } from '@mastra/core/utils';
+
+export function getSchemaName(schema?: string) {
+  return schema ? `"${parseSqlIdentifier(schema, 'schema name')}"` : undefined;
+}
+
+export function getTableName({ indexName, schemaName }: { indexName: string; schemaName?: string }) {
+  const parsedIndexName = parseSqlIdentifier(indexName, 'index name');
+  const quotedIndexName = `"${parsedIndexName}"`;
+  const quotedSchemaName = schemaName;
+  return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
+}

package/src/storage/domains/workflows/index.ts

@@ -0,0 +1,253 @@
+import type { WorkflowRun, WorkflowRuns, WorkflowRunState } from '@mastra/core';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT } from '@mastra/core/storage';
+import type { IDatabase } from 'pg-promise';
+import type { StoreOperationsPG } from '../operations';
+import { getTableName } from '../utils';
+
+function parseWorkflowRun(row: Record<string, 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 as string,
+    runId: row.run_id as string,
+    snapshot: parsedSnapshot,
+    resourceId: row.resourceId as string,
+    createdAt: new Date(row.createdAtZ || (row.createdAt as string)),
+    updatedAt: new Date(row.updatedAtZ || (row.updatedAt as string)),
+  };
+}
+
+export class WorkflowsPG extends WorkflowsStorage {
+  public client: IDatabase<{}>;
+  private operations: StoreOperationsPG;
+  private schema: string;
+
+  constructor({
+    client,
+    operations,
+    schema,
+  }: {
+    client: IDatabase<{}>;
+    operations: StoreOperationsPG;
+    schema: string;
+  }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+    this.schema = schema;
+  }
+
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot,
+  }: {
+    workflowName: string;
+    runId: string;
+    snapshot: WorkflowRunState;
+  }): Promise<void> {
+    try {
+      const now = new Date().toISOString();
+      await this.client.none(
+        `INSERT INTO ${TABLE_WORKFLOW_SNAPSHOT} (workflow_name, run_id, snapshot, "createdAt", "updatedAt")
+                 VALUES ($1, $2, $3, $4, $5)
+                 ON CONFLICT (workflow_name, run_id) DO UPDATE
+                 SET snapshot = $3, "updatedAt" = $5`,
+        [workflowName, runId, JSON.stringify(snapshot), now, now],
+      );
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_PERSIST_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId,
+  }: {
+    workflowName: string;
+    runId: string;
+  }): Promise<WorkflowRunState | null> {
+    try {
+      const result = await this.operations.load<{ snapshot: WorkflowRunState }>({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: { workflow_name: workflowName, run_id: runId },
+      });
+
+      return result ? result.snapshot : null;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_LOAD_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        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++;
+      }
+
+      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 ${getTableName({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: this.schema })} 
+          ${whereClause} 
+        `;
+
+      const queryValues = values;
+
+      const result = await this.client.oneOrNone(query, queryValues);
+
+      if (!result) {
+        return null;
+      }
+
+      return parseWorkflowRun(result);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_WORKFLOW_RUN_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            runId,
+            workflowName: workflowName || '',
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset,
+    resourceId,
+  }: {
+    workflowName?: string;
+    fromDate?: Date;
+    toDate?: Date;
+    limit?: number;
+    offset?: number;
+    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 (resourceId) {
+        const hasResourceId = await this.operations.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.`);
+        }
+      }
+
+      if (fromDate) {
+        conditions.push(`"createdAt" >= $${paramIndex}`);
+        values.push(fromDate);
+        paramIndex++;
+      }
+
+      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.client.one(
+          `SELECT COUNT(*) as count FROM ${getTableName({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: this.schema })} ${whereClause}`,
+          values,
+        );
+        total = Number(countResult.count);
+      }
+
+      // Get results
+      const query = `
+          SELECT * FROM ${getTableName({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: this.schema })} 
+          ${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 result = await this.client.manyOrNone(query, queryValues);
+
+      const runs = (result || []).map(row => {
+        return parseWorkflowRun(row);
+      });
+
+      // Use runs.length as total when not paginating
+      return { runs, total: total || runs.length };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'MASTRA_STORAGE_PG_STORE_GET_WORKFLOW_RUNS_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            workflowName: workflowName || 'all',
+          },
+        },
+        error,
+      );
+    }
+  }
+}