Lambda Stream Support

src/awslambda.ts

@@ -0,0 +1,30 @@
+import type { Writable } from 'stream';
+import type { Handler, Context } from 'aws-lambda';
+
+const anyGlobal: any = global;
+
+export namespace awslambda {
+  export type HttpMetadata = {
+    statusCode: number;
+    headers: Record<string, string>;
+    cookies?: string[];
+  };
+
+  export namespace HttpResponseStream {
+    export function from(writable: Writable, metadata: HttpMetadata): Writable {
+      return anyGlobal.awslambda.HttpResponseStream.from(writable, metadata);
+    }
+  }
+
+  export type StreamHandler<Event> = (
+    event: Event,
+    responseStream: Writable,
+    context: Context,
+  ) => void;
+
+  export function streamifyResponse<Event>(
+    handler: StreamHandler<Event>,
+  ): Handler<any, any> {
+    return anyGlobal.awslambda.streamifyResponse(handler);
+  }
+}

src/lambdaHandler.ts

@@ -5,50 +5,57 @@ import type {
 } from '@apollo/server';
 import type { WithRequired } from '@apollo/utils.withrequired';
 import type { Context, Handler } from 'aws-lambda';
-import type { LambdaResponse, MiddlewareFn } from './middleware';
-import type {
-  RequestHandler,
-  RequestHandlerEvent,
-  RequestHandlerResult,
+import {
+  runMiddleware,
+  type LambdaResponse,
+  type MiddlewareFn,
+} from './middleware';
+import {
+  isStreamRequestHandler,
+  type RequestHandler,
+  type RequestHandlerEvent,
+  type RequestHandlerResult,
+  type StreamRequestHandler,
 } from './request-handlers/_create';
+import { awslambda } from './awslambda';
+import type { Writable } from 'stream';
 
 export interface LambdaContextFunctionArgument<
-  RH extends RequestHandler<any, any>,
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
 > {
-  event: RH extends RequestHandler<infer EventType, any> ? EventType : never;
+  event: RequestHandlerEvent<RH>;
   context: Context;
 }
 
 export interface LambdaHandlerOptions<
-  RH extends RequestHandler<any, any>,
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
   TContext extends BaseContext,
 > {
   middleware?: Array<MiddlewareFn<RH>>;
   context?: ContextFunction<[LambdaContextFunctionArgument<RH>], TContext>;
 }
 
-export type LambdaHandler<RH extends RequestHandler<any, any>> = Handler<
-  RequestHandlerEvent<RH>,
-  RequestHandlerResult<RH>
->;
+export type LambdaHandler<
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
+> = Handler<RequestHandlerEvent<RH>, RequestHandlerResult<RH>>;
 
 export function startServerAndCreateLambdaHandler<
-  RH extends RequestHandler<any, any>,
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
 >(
   server: ApolloServer<BaseContext>,
   handler: RH,
   options?: LambdaHandlerOptions<RH, BaseContext>,
 ): LambdaHandler<RH>;
 export function startServerAndCreateLambdaHandler<
-  RH extends RequestHandler<any, any>,
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
   TContext extends BaseContext,
 >(
   server: ApolloServer<TContext>,
   handler: RH,
   options: WithRequired<LambdaHandlerOptions<RH, TContext>, 'context'>,
 ): LambdaHandler<RH>;
 export function startServerAndCreateLambdaHandler<
-  RH extends RequestHandler<any, any>,
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
   TContext extends BaseContext,
 >(
   server: ApolloServer<TContext>,
@@ -70,24 +77,95 @@ export function startServerAndCreateLambdaHandler<
     TContext
   > = options?.context ?? defaultContext;
 
+  if (isStreamRequestHandler(handler)) {
+    return awslambda.streamifyResponse<RequestHandlerEvent<RH>>(
+      async (event, responseStream, context) => {
+        let resultMiddlewareFns: Array<
+          LambdaResponse<RequestHandlerResult<RH>>
+        > = [];
+        let httpResponseStream: Writable | undefined;
+        try {
+          const middlewareResult = await runMiddleware(
+            event,
+            options?.middleware ?? [],
+            handler,
+          );
+          if (middlewareResult.status === 'result') {
+            httpResponseStream = awslambda.HttpResponseStream.from(
+              responseStream,
+              middlewareResult.result,
+            );
+            httpResponseStream.end();
+            return;
+          }
+          resultMiddlewareFns = middlewareResult.middleware;
+
+          const httpGraphQLRequest = handler.fromEvent(event);
+
+          const response = await server.executeHTTPGraphQLRequest({
+            httpGraphQLRequest,
+            context: () => {
+              return contextFunction({
+                event,
+                context,
+              });
+            },
+          });
+
+          const metadata = await handler.buildHTTPMetadata(response);
+
+          httpResponseStream = awslambda.HttpResponseStream.from(
+            responseStream,
+            metadata,
+          );
+
+          if (response.body.kind === 'complete') {
+            httpResponseStream.write(response.body.string);
+            httpResponseStream.end();
+            return;
+          }
+
+          for await (const chunk of response.body.asyncIterator) {
+            httpResponseStream.write(chunk);
+          }
+          httpResponseStream.end();
+        } catch (e) {
+          const { metadata, body } = await handler.toErrorResult(e);
+
+          if (httpResponseStream) {
+            httpResponseStream.write(body);
+            httpResponseStream.end();
+            return;
+          }
+
+          for (const resultMiddlewareFn of resultMiddlewareFns) {
+            await resultMiddlewareFn(metadata as any);
+          }
+
+          httpResponseStream = awslambda.HttpResponseStream.from(
+            responseStream,
+            metadata,
+          );
+          httpResponseStream.write(body);
+          httpResponseStream.end();
+        }
+      },
+    );
+  }
+
   return async function (event, context) {
-    const resultMiddlewareFns: Array<LambdaResponse<RequestHandlerResult<RH>>> =
+    let resultMiddlewareFns: Array<LambdaResponse<RequestHandlerResult<RH>>> =
       [];
     try {
-      for (const middlewareFn of options?.middleware ?? []) {
-        const middlewareReturnValue = await middlewareFn(event);
-        // If the middleware returns an object, we assume it's a LambdaResponse
-        if (
-          typeof middlewareReturnValue === 'object' &&
-          middlewareReturnValue !== null
-        ) {
-          return middlewareReturnValue;
-        }
-        // If the middleware returns a function, we assume it's a result callback
-        if (middlewareReturnValue) {
-          resultMiddlewareFns.push(middlewareReturnValue);
-        }
+      const middlewareResult = await runMiddleware(
+        event,
+        options?.middleware ?? [],
+        handler,
+      );
+      if (middlewareResult.status === 'result') {
+        return middlewareResult.result;
       }
+      resultMiddlewareFns = middlewareResult.middleware;
 
       const httpGraphQLRequest = handler.fromEvent(event);
 

src/middleware.ts

@@ -1,12 +1,67 @@
-import type { RequestHandler } from './request-handlers/_create';
+import type {
+  RequestHandler,
+  RequestHandlerEvent,
+  RequestHandlerResult,
+  StreamRequestHandler,
+} from './request-handlers/_create';
 
 export type LambdaResponse<ResultType> = (result: ResultType) => Promise<void>;
 
 export type LambdaRequest<EventType, ResultType> = (
   event: EventType,
 ) => Promise<LambdaResponse<ResultType> | ResultType | void>;
 
-export type MiddlewareFn<RH extends RequestHandler<any, any>> =
-  RH extends RequestHandler<infer EventType, infer ResultType>
-    ? LambdaRequest<EventType, ResultType>
-    : never;
+export type MiddlewareFn<
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
+> = LambdaRequest<RequestHandlerEvent<RH>, RequestHandlerResult<RH>>;
+
+export async function runMiddleware<
+  RH extends RequestHandler<any, any> | StreamRequestHandler<any>,
+>(
+  event: RequestHandlerEvent<RH>,
+  middleware: Array<MiddlewareFn<RH>>,
+  handler: RH,
+): Promise<
+  | {
+      status: 'result';
+      result: RequestHandlerResult<RH>;
+    }
+  | {
+      status: 'continue';
+      middleware: Array<LambdaResponse<RequestHandlerResult<RH>>>;
+    }
+> {
+  const resultMiddlewareFns: Array<LambdaResponse<RequestHandlerResult<RH>>> =
+    [];
+  try {
+    for (const middlewareFn of middleware) {
+      const middlewareReturnValue = await middlewareFn(event);
+      // If the middleware returns an object, we assume it's a LambdaResponse
+      if (
+        typeof middlewareReturnValue === 'object' &&
+        middlewareReturnValue !== null
+      ) {
+        return middlewareReturnValue;
+      }
+      // If the middleware returns a function, we assume it's a result callback
+      if (middlewareReturnValue) {
+        resultMiddlewareFns.push(middlewareReturnValue);
+      }
+    }
+    return {
+      status: 'continue',
+      middleware: resultMiddlewareFns,
+    };
+  } catch (e) {
+    const result = handler.toErrorResult(e);
+
+    for (const resultMiddlewareFn of resultMiddlewareFns) {
+      await resultMiddlewareFn(result);
+    }
+
+    return {
+      status: 'result',
+      result,
+    };
+  }
+}

src/request-handlers/APIGatewayProxyEventV2StreamRequestHandler.ts

@@ -0,0 +1,38 @@
+import type { APIGatewayProxyEventV2 } from 'aws-lambda';
+import { createStreamRequestHandler } from './_create';
+import { HeaderMap } from '@apollo/server';
+
+export const createAPIGatewayProxyEventV2StreamRequestHandler = <
+  Event extends APIGatewayProxyEventV2 = APIGatewayProxyEventV2,
+>() => {
+  return createStreamRequestHandler<Event>({
+    parseHttpMethod(event) {
+      return event.requestContext.http.method;
+    },
+    parseHeaders(event) {
+      const headerMap = new HeaderMap();
+      for (const [key, value] of Object.entries(event.headers ?? {})) {
+        headerMap.set(key, value ?? '');
+      }
+      return headerMap;
+    },
+    parseBody(event, headers) {
+      if (event.body) {
+        const contentType = headers.get('content-type');
+        const parsedBody = event.isBase64Encoded
+          ? Buffer.from(event.body, 'base64').toString('utf8')
+          : event.body;
+        if (contentType?.startsWith('application/json')) {
+          return JSON.parse(parsedBody);
+        }
+        if (contentType?.startsWith('text/plain')) {
+          return parsedBody;
+        }
+      }
+      return '';
+    },
+    parseQueryParams(event) {
+      return event.rawQueryString;
+    },
+  });
+};