feat: add redirectHandler option (#176)

* add `DD_APM_FLUSH_DEADLINE_MILLISECONDS` support * Revert "add `DD_APM_FLUSH_DEADLINE_MILLISECONDS` support" This reverts commit 5373a3e. * add `redirectHandler` flag * redirect only when `redirectHandler` is `true` * add tests * add docs * Update README.md Co-authored-by: Bryce Eadie <bryce.eadie@datadoghq.com> --------- Co-authored-by: Bryce Eadie <bryce.eadie@datadoghq.com>
DataDog · May 17, 2023 · 4173cc4 · 4173cc4
1 parent ae3bb8b
commit 4173cc4
Show file tree

Hide file tree

Showing 8 changed files with 120 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -191,6 +191,7 @@ _Note_: The descriptions use the npm package parameters, but they also apply to
 | `encodeAuthorizerContext`     |`encode_authorizer_context` | When set to `true` for Lambda authorizers, the tracing context will be encoded into the response for propagation. Supported for NodeJS and Python. Defaults to `true`. |
 | `decodeAuthorizerContext`     |`decode_authorizer_context` | When set to `true` for Lambdas that are authorized via Lambda authorizers, it will parse and use the encoded tracing context (if found). Supported for NodeJS and Python. Defaults to `true`.                         |
 | `apmFlushDeadline` | Used to determine when to submit spans before a timeout occurs, in milliseconds. When the remaining time in an AWS Lambda invocation is less than the value set, the tracer attempts to submit the current active spans and all finished spans. Supported for NodeJS and Python. Defaults to `100` milliseconds. |
+| `redirectHandler` | `redirect_handler` | When set to `false`, skip redirecting handler to the Datadog Lambda Library's handler. Useful when only instrumenting with Datadog Lambda Extension. Defaults to `true`. |
 
 **Note**: Using the parameters above may override corresponding function level `DD_XXX` environment variables.
 ### Tracing

diff --git a/common/constants.ts b/common/constants.ts
@@ -34,6 +34,7 @@ export const DefaultDatadogProps = {
   architecture: "X86_64",
   captureLambdaPayload: false,
   sourceCodeIntegration: true,
+  redirectHandler: true,
 };
 
 export enum TagKeys {

diff --git a/common/datadogSharedLogic.ts b/common/datadogSharedLogic.ts
@@ -75,6 +75,7 @@ export function handleSettingPropDefaults(props: DatadogProps): DatadogStrictPro
   let enableDatadogLogs = props.enableDatadogLogs;
   let captureLambdaPayload = props.captureLambdaPayload;
   let sourceCodeIntegration = props.sourceCodeIntegration;
+  let redirectHandler = props.redirectHandler;
 
   if (addLayers === undefined) {
     log.debug(`No value provided for addLayers, defaulting to ${DefaultDatadogProps.addLayers}`);
@@ -110,6 +111,11 @@ export function handleSettingPropDefaults(props: DatadogProps): DatadogStrictPro
     sourceCodeIntegration = DefaultDatadogProps.sourceCodeIntegration;
   }
 
+  if (redirectHandler === undefined) {
+    log.debug(`No value provided for redirectHandler, default to ${DefaultDatadogProps.redirectHandler}`);
+    redirectHandler = DefaultDatadogProps.redirectHandler;
+  }
+
   return {
     addLayers: addLayers,
     enableDatadogTracing: enableDatadogTracing,
@@ -119,5 +125,6 @@ export function handleSettingPropDefaults(props: DatadogProps): DatadogStrictPro
     enableDatadogLogs: enableDatadogLogs,
     captureLambdaPayload: captureLambdaPayload,
     sourceCodeIntegration: sourceCodeIntegration,
+    redirectHandler: redirectHandler,
   };
 }
diff --git a/common/interfaces.ts b/common/interfaces.ts
@@ -37,6 +37,7 @@ export interface DatadogProps {
   readonly encodeAuthorizerContext?: boolean;
   readonly decodeAuthorizerContext?: boolean;
   readonly apmFlushDeadline?: string | number;
+  readonly redirectHandler?: boolean;
 }
 
 /*
@@ -61,6 +62,7 @@ export interface DatadogStrictProps {
   readonly apiKmsKey?: string;
   readonly logLevel?: string;
   readonly sourceCodeIntegration?: boolean;
+  readonly redirectHandler?: boolean;
 }
 
 export interface Runtime {

diff --git a/v1/src/datadog.ts b/v1/src/datadog.ts
@@ -71,7 +71,10 @@ export class Datadog extends cdk.Construct {
           this.props.extensionLayerVersion,
         );
       }
-      redirectHandlers(lambdaFunctions, baseProps.addLayers);
+
+      if (baseProps.redirectHandler) {
+        redirectHandlers(lambdaFunctions, baseProps.addLayers);
+      }
 
       if (this.props.forwarderArn !== undefined) {
         if (this.props.extensionLayerVersion !== undefined) {

diff --git a/v1/test/datadog.spec.ts b/v1/test/datadog.spec.ts
@@ -3,7 +3,7 @@ import * as lambda from "@aws-cdk/aws-lambda";
 import { LogGroup } from "@aws-cdk/aws-logs";
 import * as cdk from "@aws-cdk/core";
 import { Token } from "@aws-cdk/core";
-import { addCdkConstructVersionTag, checkForMultipleApiKeys, Datadog } from "../src/index";
+import { addCdkConstructVersionTag, checkForMultipleApiKeys, Datadog, DD_HANDLER_ENV_VAR } from "../src/index";
 const versionJson = require("../version.json");
 const EXTENSION_LAYER_VERSION = 5;
 const NODE_LAYER_VERSION = 1;
@@ -419,3 +419,54 @@ describe("setTags", () => {
     });
   });
 });
+
+describe("redirectHandler", () => {
+  it("doesn't redirect handler when explicitly set to `false`", () => {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, "stack", {
+      env: {
+        region: "sa-east-1",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_14_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+
+    const datadogCdk = new Datadog(stack, "Datadog", {
+      redirectHandler: false
+    });
+    datadogCdk.addLambdaFunctions([hello]);
+
+    expect(stack).toHaveResourceLike("AWS::Lambda::Function", {
+      Handler: "hello.handler"
+    });
+  });
+
+  it("redirects handler by default", () => {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, "stack", {
+      env: {
+        region: "sa-east-1",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_14_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+
+    const datadogCdk = new Datadog(stack, "Datadog", {});
+    datadogCdk.addLambdaFunctions([hello]);
+
+    expect(stack).toHaveResourceLike("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_HANDLER_ENV_VAR]: "hello.handler"
+        }
+      }
+    });
+  });
+});
+
diff --git a/v2/src/datadog.ts b/v2/src/datadog.ts
@@ -74,7 +74,10 @@ export class Datadog extends Construct {
           this.props.extensionLayerVersion,
         );
       }
-      redirectHandlers(lambdaFunctions, baseProps.addLayers);
+
+      if (baseProps.redirectHandler) {
+        redirectHandlers(lambdaFunctions, baseProps.addLayers);
+      }
 
       if (this.props.forwarderArn !== undefined) {
         if (this.props.extensionLayerVersion !== undefined) {

diff --git a/v2/test/datadog.spec.ts b/v2/test/datadog.spec.ts
@@ -2,7 +2,7 @@ import { App, Stack, Token } from "aws-cdk-lib";
 import { Template } from "aws-cdk-lib/assertions";
 import * as lambda from "aws-cdk-lib/aws-lambda";
 import { LogGroup } from "aws-cdk-lib/aws-logs";
-import { addCdkConstructVersionTag, checkForMultipleApiKeys, Datadog } from "../src/index";
+import { addCdkConstructVersionTag, checkForMultipleApiKeys, Datadog, DD_HANDLER_ENV_VAR } from "../src/index";
 const versionJson = require("../version.json");
 const EXTENSION_LAYER_VERSION = 5;
 const NODE_LAYER_VERSION = 1;
@@ -419,3 +419,51 @@ describe("setTags", () => {
     });
   });
 });
+
+describe("redirectHandler", () => {
+  it("doesn't redirect handler when explicitly set to `false`", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack", {
+      env: {
+        region: "us-west-2",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_16_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+    const datadogCDK = new Datadog(stack, "Datadog", {
+      redirectHandler: false,
+    });
+    datadogCDK.addLambdaFunctions([hello]);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Handler: "hello.handler",
+    });
+  });
+
+  it("redirects handler by default", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack", {
+      env: {
+        region: "us-west-2",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_16_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+    const datadogCDK = new Datadog(stack, "Datadog", {});
+    datadogCDK.addLambdaFunctions([hello]);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_HANDLER_ENV_VAR]: "hello.handler",
+        },
+      },
+    });
+  });
+});