feat(parser): integrate parser with Batch Processing

package.json

@@ -10,9 +10,9 @@
     "packages/tracer",
     "packages/parameters",
     "packages/idempotency",
+    "packages/parser",
     "packages/batch",
     "packages/testing",
-    "packages/parser",
     "examples/snippets",
     "layers",
     "examples/app",

packages/batch/package.json

@@ -73,6 +73,7 @@
     "nodejs"
   ],
   "devDependencies": {
-    "@aws-lambda-powertools/testing-utils": "file:../testing"
+    "@aws-lambda-powertools/testing-utils": "file:../testing",
+    "@aws-lambda-powertools/parser": "2.25.2"
   }
 }

packages/batch/src/BasePartialBatchProcessor.ts

@@ -1,3 +1,4 @@
+import type { StandardSchemaV1 } from '@standard-schema/spec';
 import type {
   DynamoDBRecord,
   KinesisStreamRecord,
@@ -11,6 +12,7 @@ import {
 } from './constants.js';
 import { FullBatchFailureError } from './errors.js';
 import type {
+  BasePartialBatchProcessorConfig,
   EventSourceDataClassTypes,
   PartialItemFailureResponse,
   PartialItemFailures,
@@ -42,12 +44,20 @@ abstract class BasePartialBatchProcessor extends BasePartialProcessor {
    */
   public eventType: keyof typeof EventType;
 
+  /**
+   * The schema of the body of the event record for parsing
+   */
+  protected schema?: StandardSchemaV1;
+
   /**
    * Initializes base batch processing class
    *
    * @param eventType The type of event to process (SQS, Kinesis, DynamoDB)
    */
-  public constructor(eventType: keyof typeof EventType) {
+  public constructor(
+    eventType: keyof typeof EventType,
+    config?: BasePartialBatchProcessorConfig
+  ) {
     super();
     this.eventType = eventType;
     this.batchResponse = DEFAULT_RESPONSE;
@@ -56,6 +66,9 @@ abstract class BasePartialBatchProcessor extends BasePartialProcessor {
       [EventType.KinesisDataStreams]: () => this.collectKinesisFailures(),
       [EventType.DynamoDBStreams]: () => this.collectDynamoDBFailures(),
     };
+    if (config) {
+      this.schema = config.schema;
+    }
   }
 
   /**

packages/batch/src/BatchProcessor.ts

@@ -1,6 +1,19 @@
+import type { StandardSchemaV1 } from '@standard-schema/spec';
+import type {
+  DynamoDBRecord,
+  KinesisStreamRecord,
+  SQSRecord,
+  StreamRecord,
+} from 'aws-lambda';
 import { BasePartialBatchProcessor } from './BasePartialBatchProcessor.js';
+import { EventType, SchemaType } from './constants.js';
 import { BatchProcessingError } from './errors.js';
-import type { BaseRecord, FailureResponse, SuccessResponse } from './types.js';
+import type {
+  BaseRecord,
+  EventSourceDataClassTypes,
+  FailureResponse,
+  SuccessResponse,
+} from './types.js';
 
 /**
  * Process records in a batch asynchronously and handle partial failure cases.
@@ -100,7 +113,11 @@ class BatchProcessor extends BasePartialBatchProcessor {
     record: BaseRecord
   ): Promise<SuccessResponse | FailureResponse> {
     try {
-      const data = this.toBatchType(record, this.eventType);
+      const recordToProcess =
+        this.schema == null
+          ? record
+          : await this.parseRecord(record, this.eventType, this.schema);
+      const data = this.toBatchType(recordToProcess, this.eventType);
       const result = await this.handler(data, this.options?.context);
 
       return this.successHandler(record, result);
@@ -121,6 +138,156 @@ class BatchProcessor extends BasePartialBatchProcessor {
       'Not implemented. Use asyncProcess() instead.'
     );
   }
+
+  /**
+   * Create an extended schema according to the event type passed.
+   *
+   * If useTransformers is true, parsing with transformers
+   * else parse without transformers
+   *
+   * @param eventType The type of event to process (SQS, Kinesis, DynamoDB)
+   * @param schema The StandardSchema to be used for parsing
+   * @param useTransformers Whether to use transformers for parsing
+   */
+  private async createExtendedSchema(
+    eventType: keyof typeof EventType,
+    schema: StandardSchemaV1,
+    useTransformers: boolean
+  ) {
+    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 });
+      }
+
+      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 }),
+        });
+      }
+
+      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({
+          dynamodb: DynamoDBStreamChangeRecordBase.extend({
+            OldImage: (schema as any).optional(),
+            NewImage: (schema as any).optional(),
+          }),
+        });
+      }
+
+      default: {
+        console.warn(
+          `The event type provided is not supported. Supported events: ${Object.values(EventType).join(',')}`
+        );
+        throw new Error('Unsupported event type');
+      }
+    }
+  }
+
+  /**
+   * Parse the record according to the schema passed.
+   *
+   * If the passed schema is already an extended schema,
+   * it directly uses the schema to parse the record
+   *
+   * Only Zod Schemas are supported for automatic schema extension
+   *
+   * @param record The record to be parsed
+   * @param eventType The type of event to process
+   * @param schema The StandardSchema to be used for parsing
+   */
+  private async parseRecord(
+    record: EventSourceDataClassTypes,
+    eventType: keyof typeof EventType,
+    schema: StandardSchemaV1
+  ): 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;
+    }
+    // Only proceed with schema extension if it's a Zod schema
+    if (schema['~standard'].vendor !== SchemaType.Zod) {
+      console.warn(
+        'The schema provided is not supported. Only Zod schemas are supported for extension.'
+      );
+      throw new Error('Unsupported schema type');
+    }
+    // Handle schema extension based on event type
+    try {
+      // Try without transformers first, then with transformers
+      const extendedSchemaWithoutTransformers = await this.createExtendedSchema(
+        eventType,
+        schema,
+        false
+      );
+      return parse(
+        record,
+        undefined,
+        extendedSchemaWithoutTransformers
+      ) as EventSourceDataClassTypes;
+    } catch {
+      const extendedSchemaWithTransformers = await this.createExtendedSchema(
+        eventType,
+        schema,
+        true
+      );
+      return parse(
+        record,
+        undefined,
+        extendedSchemaWithTransformers
+      ) as EventSourceDataClassTypes;
+    }
+  }
 }
 
 export { BatchProcessor };

packages/batch/src/constants.ts

@@ -17,6 +17,13 @@ const EventType = {
   DynamoDBStreams: 'DynamoDBStreams',
 } as const;
 
+/**
+ * Enum of supported schema types for the utility
+ */
+const SchemaType = {
+  Zod: 'zod',
+} as const;
+
 /**
  * Default response for the partial batch processor
  */
@@ -35,4 +42,4 @@ const DATA_CLASS_MAPPING = {
     record as DynamoDBRecord,
 };
 
-export { EventType, DEFAULT_RESPONSE, DATA_CLASS_MAPPING };
+export { EventType, SchemaType, DEFAULT_RESPONSE, DATA_CLASS_MAPPING };

packages/batch/src/types.ts

@@ -1,9 +1,11 @@
+import type { StandardSchemaV1 } from '@standard-schema/spec';
 import type {
   Context,
   DynamoDBRecord,
   KinesisStreamRecord,
   SQSRecord,
 } from 'aws-lambda';
+
 import type { BasePartialBatchProcessor } from './BasePartialBatchProcessor.js';
 import type { SqsFifoPartialProcessor } from './SqsFifoPartialProcessor.js';
 import type { SqsFifoPartialProcessorAsync } from './SqsFifoPartialProcessorAsync.js';
@@ -89,6 +91,10 @@ type PartialItemFailures = { itemIdentifier: string };
  */
 type PartialItemFailureResponse = { batchItemFailures: PartialItemFailures[] };
 
+type BasePartialBatchProcessorConfig = {
+  schema: StandardSchemaV1;
+};
+
 export type {
   BatchProcessingOptions,
   BaseRecord,
@@ -97,4 +103,5 @@ export type {
   FailureResponse,
   PartialItemFailures,
   PartialItemFailureResponse,
+  BasePartialBatchProcessorConfig,
 };