feat(parser): allow parser set event type of handler with middy (#2786)

Co-authored-by: Andrea Amorosi <[email protected]>

Author: am29d

Diff for: packages/parser/src/envelopes/cloudwatch.ts

@@ -17,7 +17,7 @@ export class CloudWatchEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     const parsedEnvelope = CloudWatchLogsSchema.parse(data);
 
     return parsedEnvelope.awslogs.data.logEvents.map((record) => {

Diff for: packages/parser/src/envelopes/dynamodb.ts

@@ -3,11 +3,7 @@ import { DynamoDBStreamSchema } from '../schemas/index.js';
 import type { ParsedResult, ParsedResultError } from '../types/index.js';
 import { Envelope } from './envelope.js';
 import { ParseError } from '../errors.js';
-
-type DynamoDBStreamEnvelopeResponse<T extends ZodSchema> = {
-  NewImage: z.infer<T>;
-  OldImage: z.infer<T>;
-};
+import type { DynamoDBStreamEnvelopeResponse } from '../types/envelope.js';
 
 /**
  * DynamoDB Stream Envelope to extract data within NewImage/OldImage

Diff for: packages/parser/src/envelopes/kafka.ts

@@ -20,7 +20,7 @@ export class KafkaEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     // manually fetch event source to deside between Msk or SelfManaged
     const eventSource = (data as KafkaMskEvent)['eventSource'];
 

Diff for: packages/parser/src/envelopes/kinesis-firehose.ts

@@ -20,7 +20,7 @@ export class KinesisFirehoseEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     const parsedEnvelope = KinesisFirehoseSchema.parse(data);
 
     return parsedEnvelope.records.map((record) => {

Diff for: packages/parser/src/envelopes/kinesis.ts

@@ -18,7 +18,7 @@ export class KinesisEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     const parsedEnvelope = KinesisDataStreamSchema.parse(data);
 
     return parsedEnvelope.Records.map((record) => {

Diff for: packages/parser/src/envelopes/sns.ts

@@ -18,7 +18,7 @@ export class SnsEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     const parsedEnvelope = SnsSchema.parse(data);
 
     return parsedEnvelope.Records.map((record) => {

Diff for: packages/parser/src/envelopes/sqs.ts

@@ -17,7 +17,7 @@ export class SqsEnvelope extends Envelope {
   public static parse<T extends ZodSchema>(
     data: unknown,
     schema: T
-  ): z.infer<T> {
+  ): z.infer<T>[] {
     const parsedEnvelope = SqsSchema.parse(data);
 
     return parsedEnvelope.Records.map((record) => {

Diff for: packages/parser/src/middleware/parser.ts

@@ -1,8 +1,9 @@
 import { type MiddyLikeRequest } from '@aws-lambda-powertools/commons/types';
 import { type MiddlewareObj } from '@middy/core';
-import { type ZodSchema } from 'zod';
-import { type ParserOptions } from '../types/parser.js';
+import { ZodType } from 'zod';
+import type { ParserOptions, ParserOutput } from '../types/parser.js';
 import { parse } from '../parser.js';
+import type { Envelope } from '../types/envelope.js';
 
 /**
  * A middiy middleware to parse your event.
@@ -32,9 +33,13 @@ import { parse } from '../parser.js';
  *
  * @param options
  */
-const parser = <S extends ZodSchema>(
-  options: ParserOptions<S>
-): MiddlewareObj => {
+const parser = <
+  TSchema extends ZodType,
+  TEnvelope extends Envelope = undefined,
+  TSafeParse extends boolean = false,
+>(
+  options: ParserOptions<TSchema, TEnvelope, TSafeParse>
+): MiddlewareObj<ParserOutput<TSchema, TEnvelope, TSafeParse>> => {
   const before = (request: MiddyLikeRequest): void => {
     const { schema, envelope, safeParse } = options;
 

Diff for: packages/parser/src/parserDecorator.ts

@@ -1,8 +1,9 @@
 import type { HandlerMethodDecorator } from '@aws-lambda-powertools/commons/types';
 import type { Context, Handler } from 'aws-lambda';
-import { ZodSchema, z } from 'zod';
+import { type ZodSchema } from 'zod';
 import { parse } from './parser.js';
-import type { ParserOptions, ParsedResult } from './types/index.js';
+import type { ParserOptions, Envelope } from './types/index.js';
+import type { ParserOutput } from './types/parser.js';
 
 /**
  * A decorator to parse your event.
@@ -67,8 +68,12 @@ import type { ParserOptions, ParsedResult } from './types/index.js';
  *
  * @param options Configure the parser with the `schema`, `envelope` and whether to `safeParse` or not
  */
-export const parser = <S extends ZodSchema>(
-  options: ParserOptions<S>
+export const parser = <
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope = undefined,
+  TSafeParse extends boolean = false,
+>(
+  options: ParserOptions<TSchema, TEnvelope, TSafeParse>
 ): HandlerMethodDecorator => {
   return (_target, _propertyKey, descriptor) => {
     const original = descriptor.value!;
@@ -77,14 +82,11 @@ export const parser = <S extends ZodSchema>(
 
     descriptor.value = async function (
       this: Handler,
-      event: unknown,
+      event: ParserOutput<TSchema, TEnvelope, TSafeParse>,
       context: Context,
       callback
     ) {
-      const parsedEvent: ParsedResult<
-        typeof event,
-        z.infer<typeof schema>
-      > = parse(event, envelope, schema, safeParse);
+      const parsedEvent = parse(event, envelope, schema, safeParse);
 
       return original.call(this, parsedEvent, context, callback);
     };

Diff for: packages/parser/src/types/envelope.ts

@@ -14,8 +14,14 @@ import type {
   VpcLatticeEnvelope,
   VpcLatticeV2Envelope,
 } from '../envelopes/index.js';
+import { z, type ZodSchema } from 'zod';
 
-export type Envelope =
+type DynamoDBStreamEnvelopeResponse<Schema extends ZodSchema> = {
+  NewImage: z.infer<Schema>;
+  OldImage: z.infer<Schema>;
+};
+
+type Envelope =
   | typeof ApiGatewayEnvelope
   | typeof ApiGatewayV2Envelope
   | typeof CloudWatchEnvelope
@@ -29,4 +35,23 @@ export type Envelope =
   | typeof SnsSqsEnvelope
   | typeof SqsEnvelope
   | typeof VpcLatticeEnvelope
-  | typeof VpcLatticeV2Envelope;
+  | typeof VpcLatticeV2Envelope
+  | undefined;
+
+/**
+ * Envelopes that return an array, needed to narrow down the return type of the parser
+ */
+type EnvelopeArrayReturnType =
+  | typeof CloudWatchEnvelope
+  | typeof DynamoDBStreamEnvelope
+  | typeof KafkaEnvelope
+  | typeof KinesisEnvelope
+  | typeof KinesisFirehoseEnvelope
+  | typeof SnsEnvelope
+  | typeof SqsEnvelope;
+
+export type {
+  Envelope,
+  DynamoDBStreamEnvelopeResponse,
+  EnvelopeArrayReturnType,
+};

Diff for: packages/parser/src/types/parser.ts

@@ -1,13 +1,17 @@
-import type { ZodSchema, ZodError } from 'zod';
-import type { Envelope } from './envelope.js';
+import { type ZodSchema, type ZodError, z } from 'zod';
+import type { Envelope, EnvelopeArrayReturnType } from './envelope.js';
 
 /**
  * Options for the parser used in middy middleware and decorator
  */
-type ParserOptions<S extends ZodSchema> = {
-  schema: S;
-  envelope?: Envelope;
-  safeParse?: boolean;
+type ParserOptions<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+  TSafeParse extends boolean,
+> = {
+  schema: TSchema;
+  envelope?: TEnvelope;
+  safeParse?: TSafeParse;
 };
 
 /**
@@ -34,9 +38,46 @@ type ParsedResult<Input = unknown, Output = unknown> =
   | ParsedResultSuccess<Output>
   | ParsedResultError<Input>;
 
+/**
+ * The inferred result of the schema, can be either an array or a single object depending on the envelope
+ */
+type ZodInferredResult<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+> = undefined extends TEnvelope
+  ? z.infer<TSchema>
+  : TEnvelope extends EnvelopeArrayReturnType
+    ? z.infer<TSchema>[]
+    : z.infer<TSchema>;
+
+type ZodInferredSafeParseResult<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+> = undefined extends TEnvelope
+  ? ParsedResult<unknown, z.infer<TSchema>>
+  : TEnvelope extends EnvelopeArrayReturnType
+    ? ParsedResult<unknown, z.infer<TSchema>>
+    : ParsedResult<unknown, z.infer<TSchema>[]>;
+
+/**
+ * The output of the parser function, can be either schema inferred type or a ParsedResult
+ */
+type ParserOutput<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+  TSafeParse = false,
+> = undefined extends TSafeParse
+  ? ZodInferredResult<TSchema, TEnvelope>
+  : TSafeParse extends true
+    ? ZodInferredSafeParseResult<TSchema, TEnvelope>
+    : TSafeParse extends false
+      ? ZodInferredResult<TSchema, TEnvelope>
+      : never;
+
 export type {
   ParserOptions,
   ParsedResult,
   ParsedResultError,
   ParsedResultSuccess,
+  ParserOutput,
 };

Diff for: packages/parser/tests/unit/parser.decorator.test.ts

@@ -27,7 +27,7 @@ describe('Parser Decorator', () => {
     public async handler(
       event: TestEvent,
       _context: Context
-    ): Promise<unknown> {
+    ): Promise<TestEvent> {
       return event;
     }
 
@@ -60,7 +60,7 @@ describe('Parser Decorator', () => {
       safeParse: true,
     })
     public async handlerWithSchemaAndSafeParse(
-      event: ParsedResult<TestEvent, TestEvent>,
+      event: ParsedResult<unknown, TestEvent>,
       _context: Context
     ): Promise<ParsedResult> {
       return event;
@@ -99,9 +99,7 @@ describe('Parser Decorator', () => {
     testEvent.detail = customPayload;
 
     const resp = await lambda.handlerWithSchemaAndEnvelope(
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      testEvent,
+      testEvent as unknown as TestEvent,
       {} as Context
     );
 
@@ -130,9 +128,7 @@ describe('Parser Decorator', () => {
     testEvent.detail = customPayload;
 
     const resp = await lambda.handlerWithParserCallsAnotherMethod(
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      testEvent,
+      testEvent as unknown as TestEvent,
       {} as Context
     );
 
@@ -143,9 +139,7 @@ describe('Parser Decorator', () => {
     const testEvent = generateMock(TestSchema);
 
     const resp = await lambda.handlerWithSchemaAndSafeParse(
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      testEvent,
+      testEvent as unknown as ParsedResult<unknown, TestEvent>,
       {} as Context
     );
 
@@ -157,9 +151,10 @@ describe('Parser Decorator', () => {
 
   it('should parse event with schema and safeParse and return error', async () => {
     expect(
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      await lambda.handlerWithSchemaAndSafeParse({ foo: 'bar' }, {} as Context)
+      await lambda.handlerWithSchemaAndSafeParse(
+        { foo: 'bar' } as unknown as ParsedResult<unknown, TestEvent>,
+        {} as Context
+      )
     ).toEqual({
       error: expect.any(ParseError),
       success: false,
@@ -173,9 +168,7 @@ describe('Parser Decorator', () => {
     event.detail = testEvent;
 
     const resp = await lambda.harndlerWithEnvelopeAndSafeParse(
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      event,
+      event as unknown as ParsedResult<TestEvent, TestEvent>,
       {} as Context
     );
 
@@ -188,9 +181,7 @@ describe('Parser Decorator', () => {
   it('should parse event with envelope and safeParse and return error', async () => {
     expect(
       await lambda.harndlerWithEnvelopeAndSafeParse(
-        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-        // @ts-ignore
-        { foo: 'bar' },
+        { foo: 'bar' } as unknown as ParsedResult<TestEvent, TestEvent>,
         {} as Context
       )
     ).toEqual({