From dd693238f6433cf1d12a4778dbdd6f65a1872458 Mon Sep 17 00:00:00 2001
From: mathisobadia <mathis.obadia@gmail.com>
Date: Sat, 20 Nov 2021 22:21:57 +0100
Subject: [PATCH] feat(nextjs-cdk-construct): add runtime options to
 regeneration (#2050)

---
 .../nextjs-cdk-construct/src/index.ts         | 38 ++++++++++---------
 .../nextjs-cdk-construct/src/props.ts         |  7 +++-
 .../src/utils/toLambdaOption.ts               |  5 ++-
 3 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/packages/serverless-components/nextjs-cdk-construct/src/index.ts b/packages/serverless-components/nextjs-cdk-construct/src/index.ts
index 0f8408bed2..7578750f5a 100644
--- a/packages/serverless-components/nextjs-cdk-construct/src/index.ts
+++ b/packages/serverless-components/nextjs-cdk-construct/src/index.ts
@@ -89,6 +89,20 @@ export class NextJSLambdaEdge extends cdk.Construct {
       ...(props.s3Props || {})
     });
 
+    this.edgeLambdaRole = new Role(this, "NextEdgeLambdaRole", {
+      assumedBy: new CompositePrincipal(
+        new ServicePrincipal("lambda.amazonaws.com"),
+        new ServicePrincipal("edgelambda.amazonaws.com")
+      ),
+      managedPolicies: [
+        ManagedPolicy.fromManagedPolicyArn(
+          this,
+          "NextApiLambdaPolicy",
+          "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+        )
+      ]
+    });
+
     const hasISRPages = Object.keys(this.prerenderManifest.routes).some(
       (key) =>
         typeof this.prerenderManifest.routes[key].initialRevalidateSeconds ===
@@ -116,11 +130,15 @@ export class NextJSLambdaEdge extends cdk.Construct {
         "RegenerationFunction",
         {
           handler: "index.handler",
-          runtime: lambda.Runtime.NODEJS_14_X,
-          timeout: Duration.seconds(30),
           code: lambda.Code.fromAsset(
             path.join(this.props.serverlessBuildOutDir, "regeneration-lambda")
-          )
+          ),
+          runtime:
+            toLambdaOption("regenerationLambda", props.runtime) ??
+            lambda.Runtime.NODEJS_14_X,
+          memorySize: toLambdaOption("regenerationLambda", props.memory) ?? undefined,
+          timeout:
+            toLambdaOption("regenerationLambda", props.timeout) ?? Duration.seconds(30)
         }
       );
 
@@ -129,20 +147,6 @@ export class NextJSLambdaEdge extends cdk.Construct {
       );
     }
 
-    this.edgeLambdaRole = new Role(this, "NextEdgeLambdaRole", {
-      assumedBy: new CompositePrincipal(
-        new ServicePrincipal("lambda.amazonaws.com"),
-        new ServicePrincipal("edgelambda.amazonaws.com")
-      ),
-      managedPolicies: [
-        ManagedPolicy.fromManagedPolicyArn(
-          this,
-          "NextApiLambdaPolicy",
-          "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
-        )
-      ]
-    });
-
     this.defaultNextLambda = new lambda.Function(this, "NextLambda", {
       functionName: toLambdaOption("defaultLambda", props.name),
       description: `Default Lambda@Edge for Next CloudFront distribution`,
diff --git a/packages/serverless-components/nextjs-cdk-construct/src/props.ts b/packages/serverless-components/nextjs-cdk-construct/src/props.ts
index 4a6ffc1da5..9ee38d0744 100644
--- a/packages/serverless-components/nextjs-cdk-construct/src/props.ts
+++ b/packages/serverless-components/nextjs-cdk-construct/src/props.ts
@@ -7,7 +7,12 @@ import { Duration, StackProps } from "@aws-cdk/core";
 
 export type LambdaOption<T> =
   | T
-  | { defaultLambda?: T; apiLambda?: T; imageLambda?: T };
+  | {
+      defaultLambda?: T;
+      apiLambda?: T;
+      imageLambda?: T;
+      regenerationLambda?: T;
+    };
 
 export interface Props extends StackProps {
   /**
diff --git a/packages/serverless-components/nextjs-cdk-construct/src/utils/toLambdaOption.ts b/packages/serverless-components/nextjs-cdk-construct/src/utils/toLambdaOption.ts
index b641ff9351..cf014619c3 100644
--- a/packages/serverless-components/nextjs-cdk-construct/src/utils/toLambdaOption.ts
+++ b/packages/serverless-components/nextjs-cdk-construct/src/utils/toLambdaOption.ts
@@ -3,7 +3,7 @@
 import { LambdaOption } from "../props";
 
 export const toLambdaOption = <T extends unknown>(
-  key: "defaultLambda" | "apiLambda" | "imageLambda",
+  key: "defaultLambda" | "apiLambda" | "imageLambda" | "regenerationLambda",
   option?: LambdaOption<T>
 ): T | undefined => {
   if (
@@ -11,7 +11,8 @@ export const toLambdaOption = <T extends unknown>(
     !(
       "defaultLambda" in option ||
       "apiLambda" in option ||
-      "imageLambda" in option
+      "imageLambda" in option ||
+      "regenerationLambda" in option
     )
   ) {
     return option as T | undefined;