chore: Make helpers take one, not multiple Lambda functions (3/x)

DataDog · Nov 25, 2024 · e0f12fb · e0f12fb
1 parent 07abcbc
commit e0f12fb
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 56 deletions.
diff --git a/src/datadog-lambda.ts b/src/datadog-lambda.ts
@@ -97,20 +97,20 @@ export class DatadogLambda extends Construct {
           this.props.useLayersFromAccount,
         );
       }
-    }
 
-    if (baseProps.extensionLayerVersion !== undefined) {
-      applyExtensionLayer(
-        this.scope,
-        region,
-        extractedLambdaFunctions,
-        baseProps.extensionLayerVersion,
-        this.props.useLayersFromAccount,
-      );
-    }
+      if (baseProps.extensionLayerVersion !== undefined) {
+        applyExtensionLayer(
+          this.scope,
+          region,
+          lambdaFunction,
+          baseProps.extensionLayerVersion,
+          this.props.useLayersFromAccount,
+        );
+      }
 
-    if (baseProps.redirectHandler) {
-      redirectHandlers(extractedLambdaFunctions, baseProps.addLayers);
+      if (baseProps.redirectHandler) {
+        redirectHandlers(lambdaFunction, baseProps.addLayers);
+      }
     }
 
     if (this.props.forwarderArn !== undefined) {

diff --git a/src/layer.ts b/src/layer.ts
@@ -98,28 +98,26 @@ export function applyLayers(
 export function applyExtensionLayer(
   scope: Construct,
   region: string,
-  lambdas: lambda.Function[],
+  lam: lambda.Function,
   extensionLayerVersion: number,
   useLayersFromAccount?: string,
 ): string[] {
   // TODO: check region availability
   const errors: string[] = [];
-  log.debug("Applying extension layer to Lambda functions...");
-  lambdas.forEach((lam) => {
-    const runtime: string = lam.runtime.name;
-    const lambdaRuntimeType: RuntimeType = runtimeLookup[runtime];
-    const isARM = lam.architecture?.dockerPlatform === Architecture.ARM_64.dockerPlatform;
-    const accountId = useLayersFromAccount;
-
-    if (lambdaRuntimeType === undefined || lambdaRuntimeType === RuntimeType.UNSUPPORTED) {
-      log.debug(`Unsupported runtime: ${runtime}`);
-      return;
-    }
+  log.debug("Applying extension layer to Lambda function...");
+  const runtime: string = lam.runtime.name;
+  const lambdaRuntimeType: RuntimeType = runtimeLookup[runtime];
+  const isARM = lam.architecture?.dockerPlatform === Architecture.ARM_64.dockerPlatform;
+  const accountId = useLayersFromAccount;
+
+  if (lambdaRuntimeType === undefined || lambdaRuntimeType === RuntimeType.UNSUPPORTED) {
+    log.debug(`Unsupported runtime: ${runtime}`);
+    return errors;
+  }
 
-    const extensionLayerArn = getExtensionLayerArn(region, extensionLayerVersion, isARM, accountId);
-    log.debug(`Using extension layer: ${extensionLayerArn}`);
-    addLayer(extensionLayerArn, true, scope, lam, runtime);
-  });
+  const extensionLayerArn = getExtensionLayerArn(region, extensionLayerVersion, isARM, accountId);
+  log.debug(`Using extension layer: ${extensionLayerArn}`);
+  addLayer(extensionLayerArn, true, scope, lam, runtime);
   return errors;
 }
 

diff --git a/src/redirect.ts b/src/redirect.ts
@@ -27,35 +27,33 @@ import {
  *
  * Unchanged aside from parameter type
  */
-export function redirectHandlers(lambdas: lambda.Function[], addLayers: boolean): void {
+export function redirectHandlers(lam: lambda.Function, addLayers: boolean): void {
   log.debug(`Wrapping Lambda function handlers with Datadog handler...`);
 
-  for (const l of lambdas) {
-    const runtime: string = l.runtime.name;
-    const runtimeType: RuntimeType = runtimeLookup[runtime];
+  const runtime: string = lam.runtime.name;
+  const runtimeType: RuntimeType = runtimeLookup[runtime];
 
-    if (runtimeType === RuntimeType.JAVA || runtimeType === RuntimeType.DOTNET) {
-      l.addEnvironment(AWS_LAMBDA_EXEC_WRAPPER_ENV_VAR, AWS_LAMBDA_EXEC_WRAPPER);
-      continue;
-    }
-
-    const cfnFuntion = l.node.defaultChild as lambda.CfnFunction;
-    if (cfnFuntion === undefined) {
-      log.debug("Unable to get Lambda Function handler");
-      continue;
-    }
+  if (runtimeType === RuntimeType.JAVA || runtimeType === RuntimeType.DOTNET) {
+    lam.addEnvironment(AWS_LAMBDA_EXEC_WRAPPER_ENV_VAR, AWS_LAMBDA_EXEC_WRAPPER);
+    return;
+  }
 
-    const originalHandler = cfnFuntion.handler as string;
-    l.addEnvironment(DD_HANDLER_ENV_VAR, originalHandler);
+  const cfnFuntion = lam.node.defaultChild as lambda.CfnFunction;
+  if (cfnFuntion === undefined) {
+    log.debug("Unable to get Lambda Function handler");
+    return;
+  }
 
-    const handler = getDDHandler(runtimeType, addLayers);
-    if (handler === null) {
-      log.debug("Unable to get Datadog handler");
-      continue;
-    }
+  const originalHandler = cfnFuntion.handler as string;
+  lam.addEnvironment(DD_HANDLER_ENV_VAR, originalHandler);
 
-    cfnFuntion.handler = handler;
+  const handler = getDDHandler(runtimeType, addLayers);
+  if (handler === null) {
+    log.debug("Unable to get Datadog handler");
+    return;
   }
+
+  cfnFuntion.handler = handler;
 }
 
 function getDDHandler(runtimeType: RuntimeType, addLayers: boolean): string | null {

diff --git a/test/redirect.spec.ts b/test/redirect.spec.ts
@@ -24,7 +24,7 @@ describe("redirectHandlers", () => {
       code: lambda.Code.fromInline("test"),
       handler: "hello.handler",
     });
-    redirectHandlers([hello], true);
+    redirectHandlers(hello, true);
     Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
       Handler: `${JS_HANDLER_WITH_LAYERS}`,
     });
@@ -42,7 +42,7 @@ describe("redirectHandlers", () => {
       code: lambda.Code.fromInline("test"),
       handler: "hello.handler",
     });
-    redirectHandlers([hello], false);
+    redirectHandlers(hello, false);
     Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
       Handler: `${JS_HANDLER}`,
     });
@@ -60,7 +60,7 @@ describe("redirectHandlers", () => {
       code: lambda.Code.fromInline("test"),
       handler: "hello.handler",
     });
-    redirectHandlers([hello], true);
+    redirectHandlers(hello, true);
     Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
       Handler: `${PYTHON_HANDLER}`,
     });
@@ -88,7 +88,7 @@ describe("redirectHandlers", () => {
       code: lambda.Code.fromAsset(__dirname + "/../integration_tests/lambda"),
       handler: "handleRequest",
     });
-    redirectHandlers([hello], true);
+    redirectHandlers(hello, true);
     Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
       Handler: "handleRequest",
     });
@@ -113,7 +113,7 @@ describe("redirectHandlers", () => {
       code: lambda.Code.fromAsset(__dirname + "/../integration_tests/lambda"),
       handler: "handleRequest",
     });
-    redirectHandlers([hello], true);
+    redirectHandlers(hello, true);
     Template.fromStack(stack).resourcePropertiesCountIs(
       "AWS::Lambda::Function",
       {
@@ -137,7 +137,7 @@ describe("redirectHandlers", () => {
     const hello = new lambda.DockerImageFunction(stack, "HelloHandler", {
       code: lambda.DockerImageCode.fromImageAsset("./test/assets"),
     });
-    redirectHandlers([hello], true);
+    redirectHandlers(hello, true);
 
     Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
       Handler: Match.absent(),