@mastra/dynamodb 0.13.0 → 0.13.1

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

@@ -0,0 +1,285 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { ScoreRowData } from '@mastra/core/scores';
+import { ScoresStorage } from '@mastra/core/storage';
+import type { PaginationInfo, StoragePagination } from '@mastra/core/storage';
+import type { Service } from 'electrodb';
+
+export class ScoresStorageDynamoDB extends ScoresStorage {
+  private service: Service<Record<string, any>>;
+  constructor({ service }: { service: Service<Record<string, any>> }) {
+    super();
+    this.service = service;
+  }
+
+  // Helper function to parse score data (handle JSON fields)
+  private parseScoreData(data: any): ScoreRowData {
+    return {
+      ...data,
+      // Convert date strings back to Date objects for consistency
+      createdAt: data.createdAt ? new Date(data.createdAt) : new Date(),
+      updatedAt: data.updatedAt ? new Date(data.updatedAt) : new Date(),
+      // JSON fields are already transformed by the entity's getters
+    } as ScoreRowData;
+  }
+
+  async getScoreById({ id }: { id: string }): Promise<ScoreRowData | null> {
+    this.logger.debug('Getting score by ID', { id });
+    try {
+      const result = await this.service.entities.score.get({ entity: 'score', id }).go();
+
+      if (!result.data) {
+        return null;
+      }
+
+      return this.parseScoreData(result.data);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_SCORE_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id },
+        },
+        error,
+      );
+    }
+  }
+
+  async saveScore(score: Omit<ScoreRowData, 'id' | 'createdAt' | 'updatedAt'>): Promise<{ score: ScoreRowData }> {
+    this.logger.debug('Saving score', { scorerId: score.scorerId, runId: score.runId });
+
+    const now = new Date();
+    const scoreId = `score-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
+
+    const scoreData = {
+      entity: 'score',
+      id: scoreId,
+      scorerId: score.scorerId,
+      traceId: score.traceId || '',
+      runId: score.runId,
+      scorer: typeof score.scorer === 'string' ? score.scorer : JSON.stringify(score.scorer),
+      extractStepResult:
+        typeof score.extractStepResult === 'string' ? score.extractStepResult : JSON.stringify(score.extractStepResult),
+      analyzeStepResult:
+        typeof score.analyzeStepResult === 'string' ? score.analyzeStepResult : JSON.stringify(score.analyzeStepResult),
+      score: score.score,
+      reason: score.reason,
+      extractPrompt: score.extractPrompt,
+      analyzePrompt: score.analyzePrompt,
+      reasonPrompt: score.reasonPrompt,
+      input: typeof score.input === 'string' ? score.input : JSON.stringify(score.input),
+      output: typeof score.output === 'string' ? score.output : JSON.stringify(score.output),
+      additionalContext:
+        typeof score.additionalContext === 'string' ? score.additionalContext : JSON.stringify(score.additionalContext),
+      runtimeContext:
+        typeof score.runtimeContext === 'string' ? score.runtimeContext : JSON.stringify(score.runtimeContext),
+      entityType: score.entityType,
+      entityData: typeof score.entity === 'string' ? score.entity : JSON.stringify(score.entity),
+      entityId: score.entityId,
+      source: score.source,
+      resourceId: score.resourceId || '',
+      threadId: score.threadId || '',
+      createdAt: now.toISOString(),
+      updatedAt: now.toISOString(),
+    };
+
+    try {
+      await this.service.entities.score.upsert(scoreData).go();
+
+      const savedScore: ScoreRowData = {
+        ...score,
+        id: scoreId,
+        createdAt: now,
+        updatedAt: now,
+      };
+
+      return { score: savedScore };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_SAVE_SCORE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { scorerId: score.scorerId, runId: score.runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByScorerId({
+    scorerId,
+    pagination,
+    entityId,
+    entityType,
+  }: {
+    scorerId: string;
+    pagination: StoragePagination;
+    entityId?: string;
+    entityType?: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    this.logger.debug('Getting scores by scorer ID', { scorerId, pagination, entityId, entityType });
+
+    try {
+      // Query scores by scorer ID using the GSI
+      const query = this.service.entities.score.query.byScorer({ entity: 'score', scorerId });
+
+      // Get all scores for this scorer ID (DynamoDB doesn't support OFFSET/LIMIT)
+      const results = await query.go();
+      let allScores = results.data.map((data: any) => this.parseScoreData(data));
+
+      // Apply additional filters if provided
+      if (entityId) {
+        allScores = allScores.filter((score: ScoreRowData) => score.entityId === entityId);
+      }
+      if (entityType) {
+        allScores = allScores.filter((score: ScoreRowData) => score.entityType === entityType);
+      }
+
+      // Sort by createdAt DESC (newest first)
+      allScores.sort((a: ScoreRowData, b: ScoreRowData) => b.createdAt.getTime() - a.createdAt.getTime());
+
+      // Apply pagination in memory
+      const startIndex = pagination.page * pagination.perPage;
+      const endIndex = startIndex + pagination.perPage;
+      const paginatedScores = allScores.slice(startIndex, endIndex);
+
+      // Calculate pagination info
+      const total = allScores.length;
+      const hasMore = endIndex < total;
+
+      return {
+        scores: paginatedScores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_SCORES_BY_SCORER_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            scorerId: scorerId || '',
+            entityId: entityId || '',
+            entityType: entityType || '',
+            page: pagination.page,
+            perPage: pagination.perPage,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByRunId({
+    runId,
+    pagination,
+  }: {
+    runId: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    this.logger.debug('Getting scores by run ID', { runId, pagination });
+
+    try {
+      // Query scores by run ID using the GSI
+      const query = this.service.entities.score.query.byRun({ entity: 'score', runId });
+
+      // Get all scores for this run ID
+      const results = await query.go();
+      const allScores = results.data.map((data: any) => this.parseScoreData(data));
+
+      // Sort by createdAt DESC (newest first)
+      allScores.sort((a: ScoreRowData, b: ScoreRowData) => b.createdAt.getTime() - a.createdAt.getTime());
+
+      // Apply pagination in memory
+      const startIndex = pagination.page * pagination.perPage;
+      const endIndex = startIndex + pagination.perPage;
+      const paginatedScores = allScores.slice(startIndex, endIndex);
+
+      // Calculate pagination info
+      const total = allScores.length;
+      const hasMore = endIndex < total;
+
+      return {
+        scores: paginatedScores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_SCORES_BY_RUN_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId, page: pagination.page, perPage: pagination.perPage },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByEntityId({
+    entityId,
+    entityType,
+    pagination,
+  }: {
+    entityId: string;
+    entityType: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    this.logger.debug('Getting scores by entity ID', { entityId, entityType, pagination });
+
+    try {
+      // Use the byEntityData index which uses entityId as the primary key
+      const query = this.service.entities.score.query.byEntityData({ entity: 'score', entityId });
+
+      // Get all scores for this entity ID
+      const results = await query.go();
+      let allScores = results.data.map((data: any) => this.parseScoreData(data));
+
+      // Filter by entityType since the index only uses entityId
+      allScores = allScores.filter((score: ScoreRowData) => score.entityType === entityType);
+
+      // Sort by createdAt DESC (newest first)
+      allScores.sort((a: ScoreRowData, b: ScoreRowData) => b.createdAt.getTime() - a.createdAt.getTime());
+
+      // Apply pagination in memory
+      const startIndex = pagination.page * pagination.perPage;
+      const endIndex = startIndex + pagination.perPage;
+      const paginatedScores = allScores.slice(startIndex, endIndex);
+
+      // Calculate pagination info
+      const total = allScores.length;
+      const hasMore = endIndex < total;
+
+      return {
+        scores: paginatedScores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_SCORES_BY_ENTITY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { entityId, entityType, page: pagination.page, perPage: pagination.perPage },
+        },
+        error,
+      );
+    }
+  }
+}

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

@@ -0,0 +1,286 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { TABLE_TRACES, TracesStorage } from '@mastra/core/storage';
+import type { PaginationInfo, StorageGetTracesPaginatedArg } from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
+import type { Service } from 'electrodb';
+import type { StoreOperationsDynamoDB } from '../operations';
+
+export class TracesStorageDynamoDB extends TracesStorage {
+  private service: Service<Record<string, any>>;
+  private operations: StoreOperationsDynamoDB;
+  constructor({ service, operations }: { service: Service<Record<string, any>>; operations: StoreOperationsDynamoDB }) {
+    super();
+
+    this.service = service;
+    this.operations = operations;
+  }
+
+  // Trace operations
+  async getTraces(args: {
+    name?: string;
+    scope?: string;
+    page: number;
+    perPage: number;
+    attributes?: Record<string, string>;
+    filters?: Record<string, any>;
+  }): Promise<any[]> {
+    const { name, scope, page, perPage } = args;
+    this.logger.debug('Getting traces', { name, scope, page, perPage });
+
+    try {
+      let query;
+
+      // Determine which index to use based on the provided filters
+      // Provide *all* composite key components for the relevant index
+      if (name) {
+        query = this.service.entities.trace.query.byName({ entity: 'trace', name });
+      } else if (scope) {
+        query = this.service.entities.trace.query.byScope({ entity: 'trace', scope });
+      } else {
+        this.logger.warn('Performing a scan operation on traces - consider using a more specific query');
+        query = this.service.entities.trace.scan;
+      }
+
+      let items: any[] = [];
+      let cursor = null;
+      let pagesFetched = 0;
+      const startPage = page > 0 ? page : 1;
+
+      do {
+        const results: { data: any[]; cursor: string | null } = await query.go({ cursor, limit: perPage });
+        pagesFetched++;
+        if (pagesFetched === startPage) {
+          items = results.data;
+          break;
+        }
+        cursor = results.cursor;
+        if (!cursor && results.data.length > 0 && pagesFetched < startPage) {
+          break;
+        }
+      } while (cursor && pagesFetched < startPage);
+
+      return items;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_TRACES_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 });
+
+    if (!records.length) {
+      return;
+    }
+
+    try {
+      // Add 'entity' type to each record before passing to generic batchInsert
+      const recordsToSave = records.map(rec => ({ entity: 'trace', ...rec }));
+      await this.operations.batchInsert({
+        tableName: TABLE_TRACES,
+        records: recordsToSave, // Pass records with 'entity' included
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_BATCH_TRACE_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { count: records.length },
+        },
+        error,
+      );
+    }
+  }
+
+  async getTracesPaginated(args: StorageGetTracesPaginatedArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    this.logger.debug('Getting traces with pagination', { name, scope, page, perPage, attributes, filters, dateRange });
+
+    try {
+      let query;
+
+      // Determine which index to use based on the provided filters
+      if (name) {
+        query = this.service.entities.trace.query.byName({ entity: 'trace', name });
+      } else if (scope) {
+        query = this.service.entities.trace.query.byScope({ entity: 'trace', scope });
+      } else {
+        this.logger.warn('Performing a scan operation on traces - consider using a more specific query');
+        query = this.service.entities.trace.scan;
+      }
+
+      // For DynamoDB, we need to fetch all data and apply pagination in memory
+      // since DynamoDB doesn't support traditional offset-based pagination
+      const results = await query.go({
+        order: 'desc',
+        pages: 'all', // Get all pages to apply filtering and pagination
+      });
+
+      if (!results.data.length) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      // Apply filters in memory
+      let filteredData = results.data;
+
+      // Filter by attributes if provided
+      if (attributes) {
+        filteredData = filteredData.filter((item: Record<string, any>) => {
+          try {
+            // Handle the case where attributes might be stored as "[object Object]" or JSON string
+            let itemAttributes: Record<string, any> = {};
+
+            if (item.attributes) {
+              if (typeof item.attributes === 'string') {
+                if (item.attributes === '[object Object]') {
+                  // This means the object was stringified incorrectly
+                  itemAttributes = {};
+                } else {
+                  try {
+                    itemAttributes = JSON.parse(item.attributes);
+                  } catch {
+                    itemAttributes = {};
+                  }
+                }
+              } else if (typeof item.attributes === 'object') {
+                itemAttributes = item.attributes;
+              }
+            }
+
+            return Object.entries(attributes).every(([key, value]) => itemAttributes[key] === value);
+          } catch (e) {
+            this.logger.warn('Failed to parse attributes during filtering', { item, error: e });
+            return false;
+          }
+        });
+      }
+
+      // Filter by date range if provided
+      if (dateRange?.start) {
+        filteredData = filteredData.filter((item: Record<string, any>) => {
+          const itemDate = new Date(item.createdAt);
+          return itemDate >= dateRange.start!;
+        });
+      }
+
+      if (dateRange?.end) {
+        filteredData = filteredData.filter((item: Record<string, any>) => {
+          const itemDate = new Date(item.createdAt);
+          return itemDate <= dateRange.end!;
+        });
+      }
+
+      // Apply pagination
+      const total = filteredData.length;
+      const start = page * perPage;
+      const end = start + perPage;
+      const paginatedData = filteredData.slice(start, end);
+
+      const traces = paginatedData.map((item: any) => {
+        // Handle the case where attributes might be stored as "[object Object]" or JSON string
+        let attributes: Record<string, any> | undefined;
+        if (item.attributes) {
+          if (typeof item.attributes === 'string') {
+            if (item.attributes === '[object Object]') {
+              attributes = undefined;
+            } else {
+              try {
+                attributes = JSON.parse(item.attributes);
+              } catch {
+                attributes = undefined;
+              }
+            }
+          } else if (typeof item.attributes === 'object') {
+            attributes = item.attributes;
+          }
+        }
+
+        let status: Record<string, any> | undefined;
+        if (item.status) {
+          if (typeof item.status === 'string') {
+            try {
+              status = JSON.parse(item.status);
+            } catch {
+              status = undefined;
+            }
+          } else if (typeof item.status === 'object') {
+            status = item.status;
+          }
+        }
+
+        let events: any[] | undefined;
+        if (item.events) {
+          if (typeof item.events === 'string') {
+            try {
+              events = JSON.parse(item.events);
+            } catch {
+              events = undefined;
+            }
+          } else if (Array.isArray(item.events)) {
+            events = item.events;
+          }
+        }
+
+        let links: any[] | undefined;
+        if (item.links) {
+          if (typeof item.links === 'string') {
+            try {
+              links = JSON.parse(item.links);
+            } catch {
+              links = undefined;
+            }
+          } else if (Array.isArray(item.links)) {
+            links = item.links;
+          }
+        }
+
+        return {
+          id: item.id,
+          parentSpanId: item.parentSpanId,
+          name: item.name,
+          traceId: item.traceId,
+          scope: item.scope,
+          kind: item.kind,
+          attributes,
+          status,
+          events,
+          links,
+          other: item.other,
+          startTime: item.startTime,
+          endTime: item.endTime,
+          createdAt: item.createdAt,
+        };
+      });
+
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: end < total,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_DYNAMODB_STORE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+}