Merge branch 'main' of https://github.com/middyjs/middy

packages/core/__tests__/index.js

@@ -1,6 +1,12 @@
 import { setTimeout } from 'node:timers/promises'
 import test from 'ava'
 import sinon from 'sinon'
+import {
+  createReadableStream,
+  createPassThroughStream,
+  createWritableStream,
+  pipejoin
+} from '@datastream/core'
 import middy from '../index.js'
 
 const event = {}
@@ -771,3 +777,113 @@ test('Should not invoke timeoutEarlyResponse on error', async (t) => {
 
   t.false(timeoutCalled)
 })
+
+// streamifyResponse
+globalThis.awslambda = {
+  streamifyResponse: (cb) => cb,
+  HttpResponseStream: {
+    from: (responseStream, metadata) => {
+      return responseStream
+    }
+  }
+}
+
+test('Should return with streamifyResponse:true using undefined', async (t) => {
+  const input = ''
+  const handler = middy(
+    (event, context, { signal }) => {
+      return {
+        statusCode: 200,
+        headers: {
+          'Content-Type': 'plain/text'
+        }
+      }
+    },
+    {
+      streamifyResponse: true
+    }
+  )
+
+  let chunkResponse = ''
+  const responseStream = createWritableStream((chunk) => {
+    chunkResponse += chunk
+  })
+  const response = await handler(event, responseStream, context)
+  t.is(response, undefined)
+  t.is(chunkResponse, input)
+})
+
+test('Should return with streamifyResponse:true using string', async (t) => {
+  const input = 'x'.repeat(1024 * 1024)
+  const handler = middy({
+    streamifyResponse: true
+  }).handler((event, context, { signal }) => {
+    return {
+      statusCode: 200,
+      headers: {
+        'Content-Type': 'plain/text'
+      },
+      body: input
+    }
+  })
+
+  let chunkResponse = ''
+  const responseStream = createWritableStream((chunk) => {
+    chunkResponse += chunk
+  })
+  const response = await handler(event, responseStream, context)
+  t.is(response, undefined)
+  t.is(chunkResponse, input)
+})
+
+test('Should return with streamifyResponse:true using ReadableStream', async (t) => {
+  const input = 'x'.repeat(1024 * 1024)
+  const handler = middy(
+    async (event, context, { signal }) => {
+      return {
+        statusCode: 200,
+        headers: {
+          'Content-Type': 'plain/text'
+        },
+        body: createReadableStream(input)
+      }
+    },
+    {
+      streamifyResponse: true
+    }
+  )
+
+  let chunkResponse = ''
+  const responseStream = createWritableStream((chunk) => {
+    chunkResponse += chunk
+  })
+  const response = await handler(event, responseStream, context)
+  t.is(response, undefined)
+  t.is(chunkResponse, input)
+})
+
+test('Should return with streamifyResponse:true using ReadableStream.pipe(...)', async (t) => {
+  const input = 'x'.repeat(1024 * 1024)
+  const handler = middy(
+    async (event, context, { signal }) => {
+      return {
+        statusCode: 200,
+        headers: {
+          'Content-Type': 'plain/text'
+        },
+        body: pipejoin([createReadableStream(input), createPassThroughStream()])
+      }
+    },
+    {
+      streamifyResponse: true
+    }
+  )
+
+  let chunkResponse = ''
+  const responseStream = createWritableStream((chunk) => {
+    chunkResponse += chunk
+  })
+  const response = await handler(event, responseStream, context)
+  t.is(response, undefined)
+  t.is(chunkResponse, input)
+})

packages/core/index.d.ts

@@ -6,7 +6,9 @@ import {
 
 declare type PluginHook = () => void
 declare type PluginHookWithMiddlewareName = (middlewareName: string) => void
-declare type PluginHookPromise = (request: Request) => Promise<unknown> | unknown
+declare type PluginHookPromise = (
+  request: Request
+) => Promise<unknown> | unknown
 
 interface PluginObject {
   internal?: any
@@ -19,9 +21,15 @@ interface PluginObject {
   timeoutEarlyResponse?: PluginHook
   afterHandler?: PluginHook
   requestEnd?: PluginHookPromise
+  streamifyResponse?: Boolean
 }
 
-export interface Request<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> {
+export interface Request<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> {
   event: TEvent
   context: TContext
   response: TResult | null
@@ -31,50 +39,121 @@ export interface Request<TEvent = any, TResult = any, TErr = Error, TContext ext
   }
 }
 
-declare type MiddlewareFn<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> = (request: Request<TEvent, TResult, TErr, TContext>) => any
+declare type MiddlewareFn<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> = (request: Request<TEvent, TResult, TErr, TContext>) => any
 
-export interface MiddlewareObj<TEvent = unknown, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> {
+export interface MiddlewareObj<
+  TEvent = unknown,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> {
   before?: MiddlewareFn<TEvent, TResult, TErr, TContext>
   after?: MiddlewareFn<TEvent, TResult, TErr, TContext>
   onError?: MiddlewareFn<TEvent, TResult, TErr, TContext>
 }
 
 // The AWS provided Handler type uses void | Promise<TResult> so we have no choice but to follow and suppress the linter warning
 // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
-type MiddyInputHandler<TEvent, TResult, TContext extends LambdaContext = LambdaContext> = (event: TEvent, context: TContext, callback: LambdaCallback<TResult>) => void | Promise<TResult>
-type MiddyInputPromiseHandler<TEvent, TResult, TContext extends LambdaContext = LambdaContext> = (event: TEvent, context: TContext,) => Promise<TResult>
+type MiddyInputHandler<
+  TEvent,
+  TResult,
+  TContext extends LambdaContext = LambdaContext
+> = (
+  event: TEvent,
+  context: TContext,
+  callback: LambdaCallback<TResult>
+) => // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+void | Promise<TResult>
+type MiddyInputPromiseHandler<
+  TEvent,
+  TResult,
+  TContext extends LambdaContext = LambdaContext
+> = (event: TEvent, context: TContext) => Promise<TResult>
 
-export interface MiddyfiedHandler<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> extends MiddyInputHandler<TEvent, TResult, TContext>,
+export interface MiddyfiedHandler<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> extends MiddyInputHandler<TEvent, TResult, TContext>,
   MiddyInputPromiseHandler<TEvent, TResult, TContext> {
   use: UseFn<TEvent, TResult, TErr, TContext>
   before: AttachMiddlewareFn<TEvent, TResult, TErr, TContext>
   after: AttachMiddlewareFn<TEvent, TResult, TErr, TContext>
   onError: AttachMiddlewareFn<TEvent, TResult, TErr, TContext>
-  handler: <TAdditional>(handler: MiddlewareHandler<LambdaHandler<TEvent & TAdditional, TResult>, TContext>) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
+  handler: <TAdditional>(
+    handler: MiddlewareHandler<
+    LambdaHandler<TEvent & TAdditional, TResult>,
+    TContext
+    >
+  ) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
 }
 
-declare type AttachMiddlewareFn<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> =
-  (middleware: MiddlewareFn<TEvent, TResult, TErr, TContext>) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
+declare type AttachMiddlewareFn<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> = (
+  middleware: MiddlewareFn<TEvent, TResult, TErr, TContext>
+) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
 
-declare type AttachMiddlewareObj<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> =
-  (middleware: MiddlewareObj<TEvent, TResult, TErr, TContext>) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
+declare type AttachMiddlewareObj<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> = (
+  middleware: MiddlewareObj<TEvent, TResult, TErr, TContext>
+) => MiddyfiedHandler<TEvent, TResult, TErr, TContext>
 
-declare type UseFn<TEvent = any, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> =
-  <TMiddleware extends MiddlewareObj<any, any, Error, any>>(middlewares: TMiddleware | TMiddleware[]) => TMiddleware extends MiddlewareObj<infer TMiddlewareEvent, any, Error, infer TMiddlewareContext>
-    ? MiddyfiedHandler<TMiddlewareEvent & TEvent, TResult, TErr, TMiddlewareContext & TContext> // always true
-    : never
+declare type UseFn<
+  TEvent = any,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> = <TMiddleware extends MiddlewareObj<any, any, Error, any>>(
+  middlewares: TMiddleware | TMiddleware[]
+) => TMiddleware extends MiddlewareObj<
+infer TMiddlewareEvent,
+any,
+Error,
+infer TMiddlewareContext
+>
+  ? MiddyfiedHandler<
+  TMiddlewareEvent & TEvent,
+  TResult,
+  TErr,
+  TMiddlewareContext & TContext
+  > // always true
+  : never
 
-declare type MiddlewareHandler<THandler extends LambdaHandler<any, any>, TContext extends LambdaContext = LambdaContext> =
-  THandler extends LambdaHandler<infer TEvent, infer TResult> // always true
-    ? MiddyInputHandler<TEvent, TResult, TContext>
-    : never
+declare type MiddlewareHandler<
+  THandler extends LambdaHandler<any, any>,
+  TContext extends LambdaContext = LambdaContext
+> = THandler extends LambdaHandler<infer TEvent, infer TResult> // always true
+  ? MiddyInputHandler<TEvent, TResult, TContext>
+  : never
 
 /**
  * Middy factory function. Use it to wrap your existing handler to enable middlewares on it.
  * @param handler your original AWS Lambda function
  * @param plugin wraps around each middleware and handler to add custom lifecycle behaviours (e.g. to profile performance)
  */
-declare function middy<TEvent = unknown, TResult = any, TErr = Error, TContext extends LambdaContext = LambdaContext> (handler?: MiddlewareHandler<LambdaHandler<TEvent, TResult>, TContext>, plugin?: PluginObject): MiddyfiedHandler<TEvent, TResult, TErr, TContext>
+declare function middy<
+  TEvent = unknown,
+  TResult = any,
+  TErr = Error,
+  TContext extends LambdaContext = LambdaContext
+> (
+  handler?: MiddlewareHandler<LambdaHandler<TEvent, TResult>, TContext>,
+  plugin?: PluginObject
+): MiddyfiedHandler<TEvent, TResult, TErr, TContext>
 
 declare namespace middy {
   export {

packages/core/index.js

@@ -1,11 +1,15 @@
+/* global awslambda */
+import { Readable } from 'node:stream'
+import { pipeline } from 'node:stream/promises'
 import { setTimeout } from 'node:timers/promises'
 
 const defaultLambdaHandler = () => {}
 const defaultPlugin = {
   timeoutEarlyInMillis: 5,
   timeoutEarlyResponse: () => {
     throw new Error('Timeout')
-  }
+  },
+  streamifyResponse: false
 }
 
 const middy = (lambdaHandler = defaultLambdaHandler, plugin = {}) => {
@@ -22,7 +26,7 @@ const middy = (lambdaHandler = defaultLambdaHandler, plugin = {}) => {
   const afterMiddlewares = []
   const onErrorMiddlewares = []
 
-  const middy = (event = {}, context = {}) => {
+  const middyHandler = (event = {}, context = {}) => {
     plugin.requestStart?.()
     const request = {
       event,
@@ -41,6 +45,35 @@ const middy = (lambdaHandler = defaultLambdaHandler, plugin = {}) => {
       plugin
     )
   }
+  const middy = plugin.streamifyResponse
+    ? awslambda.streamifyResponse(async (event, responseStream, context) => {
+      const response = await middyHandler(event, context)
+      response.body ??= ''
+      let { body } = response
+
+      // Source @datastream/core (MIT)
+      if (typeof body === 'string') {
+        function * iterator (input) {
+          const size = 16 * 1024 // Node.js default
+          let position = 0
+          const length = input.length
+          while (position < length) {
+            yield input.substring(position, position + size)
+            position += size
+          }
+        }
+        body = Readable.from(iterator(response.body))
+      }
+
+      // delete response.body // Not needed
+      responseStream = awslambda.HttpResponseStream.from(
+        responseStream,
+        response
+      )
+
+      await pipeline(body, responseStream)
+    })
+    : middyHandler
 
   middy.use = (middlewares) => {
     if (!Array.isArray(middlewares)) {