Refactored the code to make the API simpler for better error handling

sdangol · sdangol · commit 0160de47c684 · 2025-09-11T09:02:01.000+01:00
diff --git a/packages/batch/src/BasePartialBatchProcessor.ts b/packages/batch/src/BasePartialBatchProcessor.ts
@@ -1,4 +1,3 @@
-import type { StandardSchemaV1 } from '@standard-schema/spec';
 import type {
   DynamoDBRecord,
   KinesisStreamRecord,
@@ -12,7 +11,7 @@ import {
 } from './constants.js';
 import { FullBatchFailureError } from './errors.js';
 import type {
-  BasePartialBatchProcessorConfig,
+  BasePartialBatchProcessorParserConfig,
   EventSourceDataClassTypes,
   PartialItemFailureResponse,
   PartialItemFailures,
@@ -45,9 +44,9 @@ abstract class BasePartialBatchProcessor extends BasePartialProcessor {
   public eventType: keyof typeof EventType;
 
   /**
-   * The schema of the body of the event record for parsing
+   * The configuration options for the parser integration
    */
-  protected schema?: StandardSchemaV1;
+  protected parserConfig?: BasePartialBatchProcessorParserConfig;
 
   /**
    * Initializes base batch processing class
@@ -56,7 +55,7 @@ abstract class BasePartialBatchProcessor extends BasePartialProcessor {
    */
   public constructor(
     eventType: keyof typeof EventType,
-    config?: BasePartialBatchProcessorConfig
+    parserConfig?: BasePartialBatchProcessorParserConfig
   ) {
     super();
     this.eventType = eventType;
@@ -66,9 +65,7 @@ abstract class BasePartialBatchProcessor extends BasePartialProcessor {
       [EventType.KinesisDataStreams]: () => this.collectKinesisFailures(),
       [EventType.DynamoDBStreams]: () => this.collectDynamoDBFailures(),
     };
-    if (config) {
-      this.schema = config.schema;
-    }
+    this.parserConfig = parserConfig;
   }
 
   /**
diff --git a/packages/batch/src/BatchProcessor.ts b/packages/batch/src/BatchProcessor.ts
@@ -4,6 +4,7 @@ import { BasePartialBatchProcessor } from './BasePartialBatchProcessor.js';
 import { EventType, SchemaVendor } from './constants.js';
 import { BatchProcessingError } from './errors.js';
 import type {
+  BasePartialBatchProcessorParserConfig,
   BaseRecord,
   EventSourceDataClassTypes,
   FailureResponse,
@@ -108,13 +109,9 @@ class BatchProcessor extends BasePartialBatchProcessor {
     record: BaseRecord
   ): Promise<SuccessResponse | FailureResponse> {
     try {
-      const recordToProcess =
-        this.schema == null
-          ? record
-          : await this.#parseRecord(record, this.eventType, this.schema);
+      const recordToProcess = await this.#parseRecord(record, this.eventType);
       const data = this.toBatchType(recordToProcess, this.eventType);
       const result = await this.handler(data, this.options?.context);
-
       return this.successHandler(record, result);
     } catch (error) {
       return this.failureHandler(record, error as Error);
@@ -146,69 +143,53 @@ class BatchProcessor extends BasePartialBatchProcessor {
    */
   async #createExtendedSchema(options: {
     eventType: keyof typeof EventType;
-    schema: StandardSchemaV1;
-    useTransformers: boolean;
+    innerSchema: StandardSchemaV1;
+    transformer?: BasePartialBatchProcessorParserConfig['transformer'];
   }) {
-    const { eventType, schema, useTransformers } = options;
+    const { eventType, innerSchema, transformer } = options;
+    let schema = innerSchema;
+    switch (transformer) {
+      case 'json': {
+        const { JSONStringified } = await import(
+          '@aws-lambda-powertools/parser/helpers'
+        );
+        schema = JSONStringified(innerSchema as any);
+        break;
+      }
+      case 'base64': {
+        const { Base64Encoded } = await import(
+          '@aws-lambda-powertools/parser/helpers'
+        );
+        schema = Base64Encoded(innerSchema as any);
+        break;
+      }
+      case 'unmarshall': {
+        const { DynamoDBMarshalled } = await import(
+          '@aws-lambda-powertools/parser/helpers/dynamodb'
+        );
+        schema = DynamoDBMarshalled(innerSchema as any);
+        break;
+      }
+    }
     switch (eventType) {
       case EventType.SQS: {
-        if (useTransformers) {
-          const [{ JSONStringified }, { SqsRecordSchema }] = await Promise.all([
-            import('@aws-lambda-powertools/parser/helpers'),
-            import('@aws-lambda-powertools/parser/schemas/sqs'),
-          ]);
-          return SqsRecordSchema.extend({
-            body: JSONStringified(schema as any),
-          });
-        }
         const { SqsRecordSchema } = await import(
           '@aws-lambda-powertools/parser/schemas/sqs'
         );
-        return SqsRecordSchema.extend({ body: schema });
+        return SqsRecordSchema.extend({
+          body: schema,
+        });
       }
-
       case EventType.KinesisDataStreams: {
-        if (useTransformers) {
-          const [
-            { Base64Encoded },
-            { KinesisDataStreamRecord, KinesisDataStreamRecordPayload },
-          ] = await Promise.all([
-            import('@aws-lambda-powertools/parser/helpers'),
-            import('@aws-lambda-powertools/parser/schemas/kinesis'),
-          ]);
-          return KinesisDataStreamRecord.extend({
-            kinesis: KinesisDataStreamRecordPayload.extend({
-              data: Base64Encoded(schema as any),
-            }),
-          });
-        }
         const { KinesisDataStreamRecord, KinesisDataStreamRecordPayload } =
           await import('@aws-lambda-powertools/parser/schemas/kinesis');
         return KinesisDataStreamRecord.extend({
-          kinesis: KinesisDataStreamRecordPayload.extend({ data: schema }),
+          kinesis: KinesisDataStreamRecordPayload.extend({
+            data: schema,
+          }),
         });
       }
-
       case EventType.DynamoDBStreams: {
-        if (useTransformers) {
-          const [
-            { DynamoDBMarshalled },
-            { DynamoDBStreamRecord, DynamoDBStreamChangeRecordBase },
-          ] = await Promise.all([
-            import('@aws-lambda-powertools/parser/helpers/dynamodb'),
-            import('@aws-lambda-powertools/parser/schemas/dynamodb'),
-          ]);
-          return DynamoDBStreamRecord.extend({
-            dynamodb: DynamoDBStreamChangeRecordBase.extend({
-              OldImage: DynamoDBMarshalled<StreamRecord['OldImage']>(
-                schema as any
-              ).optional(),
-              NewImage: DynamoDBMarshalled<StreamRecord['NewImage']>(
-                schema as any
-              ).optional(),
-            }),
-          });
-        }
         const { DynamoDBStreamRecord, DynamoDBStreamChangeRecordBase } =
           await import('@aws-lambda-powertools/parser/schemas/dynamodb');
         return DynamoDBStreamRecord.extend({
@@ -218,7 +199,6 @@ class BatchProcessor extends BasePartialBatchProcessor {
           }),
         });
       }
-
       default: {
         console.warn(
           `The event type provided is not supported. Supported events: ${Object.values(EventType).join(',')}`
@@ -238,57 +218,77 @@ class BatchProcessor extends BasePartialBatchProcessor {
    *
    * @param record - The record to be parsed
    * @param eventType - The type of event to process
-   * @param schema - The StandardSchema to be used for parsing
    */
   async #parseRecord(
     record: EventSourceDataClassTypes,
-    eventType: keyof typeof EventType,
-    schema: StandardSchemaV1
+    eventType: keyof typeof EventType
   ): Promise<EventSourceDataClassTypes> {
-    const { parse } = await import('@aws-lambda-powertools/parser');
-    // Try parsing with the original schema first
-    const extendedSchemaParsing = parse(record, undefined, schema, true);
-    if (extendedSchemaParsing.success) {
-      return extendedSchemaParsing.data as EventSourceDataClassTypes;
+    if (this.parserConfig == null) {
+      return record;
     }
-    // Only proceed with schema extension if it's a Zod schema
-    if (schema['~standard'].vendor !== SchemaVendor.Zod) {
-      console.warn(
-        'The schema provided is not supported. Only Zod schemas are supported for extension.'
+    const { parse } = await import('@aws-lambda-powertools/parser');
+    const { schema, innerSchema, transformer } = this.parserConfig;
+    // If the external schema is specified, use it to parse the record
+    if (schema != null) {
+      const extendedSchemaParsing = parse(record, undefined, schema, true);
+      if (extendedSchemaParsing.success) {
+        return extendedSchemaParsing.data as EventSourceDataClassTypes;
+      }
+      const issues = extendedSchemaParsing.error
+        .cause as ReadonlyArray<StandardSchemaV1.Issue>;
+      throw new Error(
+        `Failed to parse record: ${issues.map((issue) => `${issue?.path?.join('.')}: ${issue.message}`).join('; ')} `
       );
-      throw new Error('Unsupported schema type');
-    }
-    // Handle schema extension based on event type
-    // Try without transformers first, then with transformers
-    const schemaWithoutTransformers = await this.#createExtendedSchema({
-      eventType,
-      schema,
-      useTransformers: false,
-    });
-    const schemaWithoutTransformersParsing = parse(
-      record,
-      undefined,
-      schemaWithoutTransformers,
-      true
-    );
-    if (schemaWithoutTransformersParsing.success) {
-      return schemaWithoutTransformersParsing.data as EventSourceDataClassTypes;
     }
-    const schemaWithTransformers = await this.#createExtendedSchema({
-      eventType,
-      schema,
-      useTransformers: true,
-    });
-    const schemaWithTransformersParsing = parse(
-      record,
-      undefined,
-      schemaWithTransformers,
-      true
-    );
-    if (schemaWithTransformersParsing.success) {
-      return schemaWithTransformersParsing.data as EventSourceDataClassTypes;
+    if (innerSchema != null) {
+      // Only proceed with schema extension if it's a Zod schema
+      if (innerSchema['~standard'].vendor !== SchemaVendor.Zod) {
+        console.warn(
+          'The schema provided is not supported. Only Zod schemas are supported for extension.'
+        );
+        throw new Error('Unsupported schema type');
+      }
+      if (transformer != null) {
+        const schemaWithTransformers = await this.#createExtendedSchema({
+          eventType,
+          innerSchema,
+          transformer,
+        });
+        const schemaWithTransformersParsing = parse(
+          record,
+          undefined,
+          schemaWithTransformers,
+          true
+        );
+        if (schemaWithTransformersParsing.success) {
+          return schemaWithTransformersParsing.data as EventSourceDataClassTypes;
+        }
+        const issues = schemaWithTransformersParsing.error
+          .cause as ReadonlyArray<StandardSchemaV1.Issue>;
+        throw new Error(
+          `Failed to parse record: ${issues.map((issue) => `${issue?.path?.join('.')}: ${issue.message}`).join('; ')} `
+        );
+      }
+      const schemaWithoutTransformers = await this.#createExtendedSchema({
+        eventType,
+        innerSchema,
+      });
+      const schemaWithoutTransformersParsing = parse(
+        record,
+        undefined,
+        schemaWithoutTransformers,
+        true
+      );
+      if (schemaWithoutTransformersParsing.success) {
+        return schemaWithoutTransformersParsing.data as EventSourceDataClassTypes;
+      }
+      const issues = schemaWithoutTransformersParsing.error
+        .cause as ReadonlyArray<StandardSchemaV1.Issue>;
+      throw new Error(
+        `Failed to parse record: ${issues.map((issue) => `${issue?.path?.join('.')}: ${issue.message}`).join('; ')} `
+      );
     }
-    throw new Error('Failed to parse record');
+    throw new Error('Either schema or innerSchema is required for parsing');
   }
 }
 
diff --git a/packages/batch/src/types.ts b/packages/batch/src/types.ts
@@ -5,7 +5,7 @@ import type {
   KinesisStreamRecord,
   SQSRecord,
 } from 'aws-lambda';
-
+import type { GenericLogger } from '../../commons/lib/esm/types/GenericLogger.js';
 import type { BasePartialBatchProcessor } from './BasePartialBatchProcessor.js';
 import type { SqsFifoPartialProcessor } from './SqsFifoPartialProcessor.js';
 import type { SqsFifoPartialProcessorAsync } from './SqsFifoPartialProcessorAsync.js';
@@ -92,18 +92,37 @@ type PartialItemFailures = { itemIdentifier: string };
 type PartialItemFailureResponse = { batchItemFailures: PartialItemFailures[] };
 
 /**
- * Type representing the configuration options passed to the BasePartialBatchProcessor class.
+ * Type representing the parser configuration options passed to the BasePartialBatchProcessor class.
  *
  * @property schema - The schema to be used for parsing
+ * @property innerSchema - The schema for the inner payload
+ * @property transformer - The transformer to be used for parsing the payload
+ * @property logger - The logger to be used for logging debug and warning messages.
  */
-type BasePartialBatchProcessorConfig = {
+type BasePartialBatchProcessorParserConfig = {
+  /**
+   * The schema for the full event including the extended inner payload schema.
+   *
+   * StandardSchema is supported.
+   */
+  schema?: StandardSchemaV1;
+  /**
+   * The schema for the inner payload of the event.
+   * Only Zod schemas are supported.
+   */
+  innerSchema?: StandardSchemaV1;
+  /**
+   * The transformer to be used for parsing the payload.
+   * Supported transformers are:
+   * 1. 'json': Uses JSONStringified helper
+   * 2. 'base64': Uses Base64Encoded helper
+   * 3. 'unmarshall': Uses DynamoDBMarshalled helper
+   */
+  transformer?: 'json' | 'base64' | 'unmarshall';
   /**
-   * The schema be either of the following:
-   * 1. An internal schema of the payload of the supported event types.
-   * 2. An internal schema along with helper transformer functions.
-   * 3. An extended schema of the supported event type.
+   * The logger to be used for logging debug and warning messages.
    */
-  schema: StandardSchemaV1;
+  logger?: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
 };
 
 export type {
@@ -114,5 +133,5 @@ export type {
   FailureResponse,
   PartialItemFailures,
   PartialItemFailureResponse,
-  BasePartialBatchProcessorConfig,
+  BasePartialBatchProcessorParserConfig,
 };
diff --git a/packages/batch/tests/unit/BatchProcessor.test.ts b/packages/batch/tests/unit/BatchProcessor.test.ts
