@mastra/dynamodb 0.14.3 → 0.14.6-alpha.0

package/src/storage/index.ts

@@ -1,504 +0,0 @@
-import { DynamoDBClient, DescribeTableCommand } from '@aws-sdk/client-dynamodb';
-import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
-import type { MastraMessageContentV2 } from '@mastra/core/agent';
-import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
-import type { StorageThreadType, MastraMessageV2, MastraMessageV1 } from '@mastra/core/memory';
-
-import type { ScoreRowData, ScoringSource } from '@mastra/core/scores';
-import { MastraStorage } from '@mastra/core/storage';
-import type {
-  EvalRow,
-  StorageGetMessagesArg,
-  WorkflowRun,
-  WorkflowRuns,
-  TABLE_NAMES,
-  StorageGetTracesArg,
-  PaginationInfo,
-  StorageColumn,
-  StoragePagination,
-  StorageDomains,
-  PaginationArgs,
-  StorageResourceType,
-  ThreadSortOptions,
-} from '@mastra/core/storage';
-import type { Trace } from '@mastra/core/telemetry';
-import type { WorkflowRunState } from '@mastra/core/workflows';
-import type { Service } from 'electrodb';
-import { getElectroDbService } from '../entities';
-import { LegacyEvalsDynamoDB } from './domains/legacy-evals';
-import { MemoryStorageDynamoDB } from './domains/memory';
-import { StoreOperationsDynamoDB } from './domains/operations';
-import { ScoresStorageDynamoDB } from './domains/score';
-import { TracesStorageDynamoDB } from './domains/traces';
-import { WorkflowStorageDynamoDB } from './domains/workflows';
-
-export interface DynamoDBStoreConfig {
-  region?: string;
-  tableName: string;
-  endpoint?: string;
-  credentials?: {
-    accessKeyId: string;
-    secretAccessKey: string;
-  };
-}
-
-// Define a type for our service that allows string indexing
-type MastraService = Service<Record<string, any>> & {
-  [key: string]: any;
-};
-
-export class DynamoDBStore extends MastraStorage {
-  private tableName: string;
-  private client: DynamoDBDocumentClient;
-  private service: MastraService;
-  protected hasInitialized: Promise<boolean> | null = null;
-  stores: StorageDomains;
-
-  constructor({ name, config }: { name: string; config: DynamoDBStoreConfig }) {
-    super({ name });
-
-    // Validate required config
-    try {
-      if (!config.tableName || typeof config.tableName !== 'string' || config.tableName.trim() === '') {
-        throw new Error('DynamoDBStore: config.tableName must be provided and cannot be empty.');
-      }
-      // Validate tableName characters (basic check)
-      if (!/^[a-zA-Z0-9_.-]{3,255}$/.test(config.tableName)) {
-        throw new Error(
-          `DynamoDBStore: config.tableName "${config.tableName}" contains invalid characters or is not between 3 and 255 characters long.`,
-        );
-      }
-
-      const dynamoClient = new DynamoDBClient({
-        region: config.region || 'us-east-1',
-        endpoint: config.endpoint,
-        credentials: config.credentials,
-      });
-
-      this.tableName = config.tableName;
-      this.client = DynamoDBDocumentClient.from(dynamoClient);
-      this.service = getElectroDbService(this.client, this.tableName) as MastraService;
-
-      const operations = new StoreOperationsDynamoDB({
-        service: this.service,
-        tableName: this.tableName,
-        client: this.client,
-      });
-
-      const traces = new TracesStorageDynamoDB({ service: this.service, operations });
-
-      const workflows = new WorkflowStorageDynamoDB({ service: this.service });
-
-      const memory = new MemoryStorageDynamoDB({ service: this.service });
-
-      const scores = new ScoresStorageDynamoDB({ service: this.service });
-
-      this.stores = {
-        operations,
-        legacyEvals: new LegacyEvalsDynamoDB({ service: this.service, tableName: this.tableName }),
-        traces,
-        workflows,
-        memory,
-        scores,
-      };
-    } catch (error) {
-      throw new MastraError(
-        {
-          id: 'STORAGE_DYNAMODB_STORE_CONSTRUCTOR_FAILED',
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.USER,
-        },
-        error,
-      );
-    }
-
-    // We're using a single table design with ElectroDB,
-    // so we don't need to create multiple tables
-  }
-
-  get supports() {
-    return {
-      selectByIncludeResourceScope: true,
-      resourceWorkingMemory: true,
-      hasColumn: false,
-      createTable: false,
-      deleteMessages: false,
-    };
-  }
-
-  /**
-   * Validates that the required DynamoDB table exists and is accessible.
-   * This does not check the table structure - it assumes the table
-   * was created with the correct structure via CDK/CloudFormation.
-   */
-  private async validateTableExists(): Promise<boolean> {
-    try {
-      const command = new DescribeTableCommand({
-        TableName: this.tableName,
-      });
-
-      // If the table exists, this call will succeed
-      // If the table doesn't exist, it will throw a ResourceNotFoundException
-      await this.client.send(command);
-      return true;
-    } catch (error: any) {
-      // If the table doesn't exist, DynamoDB returns a ResourceNotFoundException
-      if (error.name === 'ResourceNotFoundException') {
-        return false;
-      }
-
-      // For other errors (like permissions issues), we should throw
-      throw new MastraError(
-        {
-          id: 'STORAGE_DYNAMODB_STORE_VALIDATE_TABLE_EXISTS_FAILED',
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.THIRD_PARTY,
-          details: { tableName: this.tableName },
-        },
-        error,
-      );
-    }
-  }
-
-  /**
-   * Initialize storage, validating the externally managed table is accessible.
-   * For the single-table design, we only validate once that we can access
-   * the table that was created via CDK/CloudFormation.
-   */
-  async init(): Promise<void> {
-    if (this.hasInitialized === null) {
-      // If no initialization promise exists, create and store it.
-      // This assignment ensures that even if multiple calls arrive here concurrently,
-      // they will all eventually await the same promise instance created by the first one
-      // to complete this assignment.
-      this.hasInitialized = this._performInitializationAndStore();
-    }
-
-    try {
-      // Await the stored promise.
-      // If initialization was successful, this resolves.
-      // If it failed, this will re-throw the error caught and re-thrown by _performInitializationAndStore.
-      await this.hasInitialized;
-    } catch (error) {
-      // The error has already been handled by _performInitializationAndStore
-      // (i.e., this.hasInitialized was reset). Re-throwing here ensures
-      // the caller of init() is aware of the failure.
-      throw new MastraError(
-        {
-          id: 'STORAGE_DYNAMODB_STORE_INIT_FAILED',
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.THIRD_PARTY,
-          details: { tableName: this.tableName },
-        },
-        error,
-      );
-    }
-  }
-
-  /**
-   * Performs the actual table validation and stores the promise.
-   * Handles resetting the stored promise on failure to allow retries.
-   */
-  private _performInitializationAndStore(): Promise<boolean> {
-    return this.validateTableExists()
-      .then(exists => {
-        if (!exists) {
-          throw new Error(
-            `Table ${this.tableName} does not exist or is not accessible. Ensure it's created via CDK/CloudFormation before using this store.`,
-          );
-        }
-        // Successfully initialized
-        return true;
-      })
-      .catch(err => {
-        // Initialization failed. Clear the stored promise to allow future calls to init() to retry.
-        this.hasInitialized = null;
-        // Re-throw the error so it can be caught by the awaiter in init()
-        throw err;
-      });
-  }
-
-  async createTable({ tableName, schema }: { tableName: TABLE_NAMES; schema: Record<string, any> }): Promise<void> {
-    return this.stores.operations.createTable({ tableName, schema });
-  }
-
-  async alterTable(_args: {
-    tableName: TABLE_NAMES;
-    schema: Record<string, StorageColumn>;
-    ifNotExists: string[];
-  }): Promise<void> {
-    return this.stores.operations.alterTable(_args);
-  }
-
-  async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
-    return this.stores.operations.clearTable({ tableName });
-  }
-
-  async dropTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
-    return this.stores.operations.dropTable({ tableName });
-  }
-
-  async insert({ tableName, record }: { tableName: TABLE_NAMES; record: Record<string, any> }): Promise<void> {
-    return this.stores.operations.insert({ tableName, record });
-  }
-
-  async batchInsert({ tableName, records }: { tableName: TABLE_NAMES; records: Record<string, any>[] }): Promise<void> {
-    return this.stores.operations.batchInsert({ tableName, records });
-  }
-
-  async load<R>({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, string> }): Promise<R | null> {
-    return this.stores.operations.load({ tableName, keys });
-  }
-
-  // Thread operations
-  async getThreadById({ threadId }: { threadId: string }): Promise<StorageThreadType | null> {
-    return this.stores.memory.getThreadById({ threadId });
-  }
-
-  async getThreadsByResourceId(args: { resourceId: string } & ThreadSortOptions): Promise<StorageThreadType[]> {
-    return this.stores.memory.getThreadsByResourceId(args);
-  }
-
-  async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
-    return this.stores.memory.saveThread({ thread });
-  }
-
-  async updateThread({
-    id,
-    title,
-    metadata,
-  }: {
-    id: string;
-    title: string;
-    metadata: Record<string, unknown>;
-  }): Promise<StorageThreadType> {
-    return this.stores.memory.updateThread({ id, title, metadata });
-  }
-
-  async deleteThread({ threadId }: { threadId: string }): Promise<void> {
-    return this.stores.memory.deleteThread({ threadId });
-  }
-
-  // Message operations
-  public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
-  public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
-  public async getMessages({
-    threadId,
-    resourceId,
-    selectBy,
-    format,
-  }: StorageGetMessagesArg & { format?: 'v1' | 'v2' }): Promise<MastraMessageV1[] | MastraMessageV2[]> {
-    return this.stores.memory.getMessages({ threadId, resourceId, selectBy, format });
-  }
-
-  async getMessagesById({ messageIds, format }: { messageIds: string[]; format: 'v1' }): Promise<MastraMessageV1[]>;
-  async getMessagesById({ messageIds, format }: { messageIds: string[]; format?: 'v2' }): Promise<MastraMessageV2[]>;
-  async getMessagesById({
-    messageIds,
-    format,
-  }: {
-    messageIds: string[];
-    format?: 'v1' | 'v2';
-  }): Promise<MastraMessageV1[] | MastraMessageV2[]> {
-    return this.stores.memory.getMessagesById({ messageIds, format });
-  }
-
-  async saveMessages(args: { messages: MastraMessageV1[]; format?: undefined | 'v1' }): Promise<MastraMessageV1[]>;
-  async saveMessages(args: { messages: MastraMessageV2[]; format: 'v2' }): Promise<MastraMessageV2[]>;
-  async saveMessages(
-    args: { messages: MastraMessageV1[]; format?: undefined | 'v1' } | { messages: MastraMessageV2[]; format: 'v2' },
-  ): Promise<MastraMessageV2[] | MastraMessageV1[]> {
-    return this.stores.memory.saveMessages(args);
-  }
-
-  async getThreadsByResourceIdPaginated(
-    args: {
-      resourceId: string;
-      page: number;
-      perPage: number;
-    } & ThreadSortOptions,
-  ): Promise<PaginationInfo & { threads: StorageThreadType[] }> {
-    return this.stores.memory.getThreadsByResourceIdPaginated(args);
-  }
-
-  async getMessagesPaginated(
-    args: StorageGetMessagesArg & { format?: 'v1' | 'v2' },
-  ): Promise<PaginationInfo & { messages: MastraMessageV1[] | MastraMessageV2[] }> {
-    return this.stores.memory.getMessagesPaginated(args);
-  }
-
-  async updateMessages(_args: {
-    messages: Partial<Omit<MastraMessageV2, 'createdAt'>> &
-      {
-        id: string;
-        content?: { metadata?: MastraMessageContentV2['metadata']; content?: MastraMessageContentV2['content'] };
-      }[];
-  }): Promise<MastraMessageV2[]> {
-    return this.stores.memory.updateMessages(_args);
-  }
-
-  // Trace operations
-  async getTraces(args: {
-    name?: string;
-    scope?: string;
-    page: number;
-    perPage: number;
-    attributes?: Record<string, string>;
-    filters?: Record<string, any>;
-  }): Promise<any[]> {
-    return this.stores.traces.getTraces(args);
-  }
-
-  async batchTraceInsert({ records }: { records: Record<string, any>[] }): Promise<void> {
-    return this.stores.traces.batchTraceInsert({ records });
-  }
-
-  async getTracesPaginated(_args: StorageGetTracesArg): Promise<PaginationInfo & { traces: Trace[] }> {
-    return this.stores.traces.getTracesPaginated(_args);
-  }
-
-  // Workflow operations
-  async persistWorkflowSnapshot({
-    workflowName,
-    runId,
-    snapshot,
-  }: {
-    workflowName: string;
-    runId: string;
-    snapshot: WorkflowRunState;
-  }): Promise<void> {
-    return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, snapshot });
-  }
-
-  async loadWorkflowSnapshot({
-    workflowName,
-    runId,
-  }: {
-    workflowName: string;
-    runId: string;
-  }): Promise<WorkflowRunState | null> {
-    return this.stores.workflows.loadWorkflowSnapshot({ workflowName, runId });
-  }
-
-  async getWorkflowRuns(args?: {
-    workflowName?: string;
-    fromDate?: Date;
-    toDate?: Date;
-    limit?: number;
-    offset?: number;
-    resourceId?: string;
-  }): Promise<WorkflowRuns> {
-    return this.stores.workflows.getWorkflowRuns(args);
-  }
-
-  async getWorkflowRunById(args: { runId: string; workflowName?: string }): Promise<WorkflowRun | null> {
-    return this.stores.workflows.getWorkflowRunById(args);
-  }
-
-  async getResourceById({ resourceId }: { resourceId: string }): Promise<StorageResourceType | null> {
-    return this.stores.memory.getResourceById({ resourceId });
-  }
-
-  async saveResource({ resource }: { resource: StorageResourceType }): Promise<StorageResourceType> {
-    return this.stores.memory.saveResource({ resource });
-  }
-
-  async updateResource({
-    resourceId,
-    workingMemory,
-    metadata,
-  }: {
-    resourceId: string;
-    workingMemory?: string;
-    metadata?: Record<string, any>;
-  }): Promise<StorageResourceType> {
-    return this.stores.memory.updateResource({ resourceId, workingMemory, metadata });
-  }
-
-  // Eval operations
-  async getEvalsByAgentName(agentName: string, type?: 'test' | 'live'): Promise<EvalRow[]> {
-    return this.stores.legacyEvals.getEvalsByAgentName(agentName, type);
-  }
-
-  async getEvals(
-    options: {
-      agentName?: string;
-      type?: 'test' | 'live';
-    } & PaginationArgs,
-  ): Promise<PaginationInfo & { evals: EvalRow[] }> {
-    return this.stores.legacyEvals.getEvals(options);
-  }
-
-  /**
-   * Closes the DynamoDB client connection and cleans up resources.
-   * Should be called when the store is no longer needed, e.g., at the end of tests or application shutdown.
-   */
-  public async close(): Promise<void> {
-    this.logger.debug('Closing DynamoDB client for store:', { name: this.name });
-    try {
-      this.client.destroy();
-      this.logger.debug('DynamoDB client closed successfully for store:', { name: this.name });
-    } catch (error) {
-      throw new MastraError(
-        {
-          id: 'STORAGE_DYNAMODB_STORE_CLOSE_FAILED',
-          domain: ErrorDomain.STORAGE,
-          category: ErrorCategory.THIRD_PARTY,
-        },
-        error,
-      );
-    }
-  }
-  /**
-   * SCORERS - Not implemented
-   */
-  async getScoreById({ id: _id }: { id: string }): Promise<ScoreRowData | null> {
-    return this.stores.scores.getScoreById({ id: _id });
-  }
-
-  async saveScore(_score: ScoreRowData): Promise<{ score: ScoreRowData }> {
-    return this.stores.scores.saveScore(_score);
-  }
-
-  async getScoresByRunId({
-    runId: _runId,
-    pagination: _pagination,
-  }: {
-    runId: string;
-    pagination: StoragePagination;
-  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
-    return this.stores.scores.getScoresByRunId({ runId: _runId, pagination: _pagination });
-  }
-
-  async getScoresByEntityId({
-    entityId: _entityId,
-    entityType: _entityType,
-    pagination: _pagination,
-  }: {
-    pagination: StoragePagination;
-    entityId: string;
-    entityType: string;
-  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
-    return this.stores.scores.getScoresByEntityId({
-      entityId: _entityId,
-      entityType: _entityType,
-      pagination: _pagination,
-    });
-  }
-
-  async getScoresByScorerId({
-    scorerId,
-    source,
-    entityId,
-    entityType,
-    pagination,
-  }: {
-    scorerId: string;
-    entityId?: string;
-    entityType?: string;
-    source?: ScoringSource;
-    pagination: StoragePagination;
-  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
-    return this.stores.scores.getScoresByScorerId({ scorerId, source, entityId, entityType, pagination });
-  }
-}