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

Co-authored-by: Andrea Amorosi <dreamorosi@gmail.com>

Author: am29d

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) => {

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

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'];
 

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) => {

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) => {

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) => {

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) => {

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;
 

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);
     };

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,
+};