Support CDK Aspects

README.md

@@ -97,6 +97,17 @@ datadog.addLambdaFunctions([<LAMBDA_FUNCTIONS>])
 datadog.addForwarderToNonLambdaLogGroups([<LOG_GROUPS>])
 ```
 
+Or you can use CDK Aspects:
+```
+import { Aspects } from "aws-cdk-lib";
+import { Datadog, DatadogAspect } from "datadog-cdk-constructs-v2";
+const datadog = new Datadog(this as any, "Datadog", {
+  // ... Same as above
+});
+
+Aspects.of(this).add(new DatadogAspect(datadog))
+```
+
 ## Source Code Integration
 [Source code integration](https://docs.datadoghq.com/integrations/guide/source-code-integration/) is enabled by default through automatic lambda tagging, and will work if:
 

v2/src/datadog.ts

@@ -6,11 +6,11 @@
  * Copyright 2021 Datadog, Inc.
  */
 
-import { Tags } from "aws-cdk-lib";
+import { Tags, IAspect } from "aws-cdk-lib";
 import * as lambda from "aws-cdk-lib/aws-lambda";
 import * as logs from "aws-cdk-lib/aws-logs";
 import { ISecret, Secret } from "aws-cdk-lib/aws-secretsmanager";
-import { Construct } from "constructs";
+import { Construct, IConstruct } from "constructs";
 import log from "loglevel";
 import {
   applyLayers,
@@ -30,6 +30,16 @@ import { LambdaFunction } from "./interfaces";
 
 const versionJson = require("../version.json");
 
+export class DatadogAspect implements IAspect {
+  constructor(private datadog: Datadog) {}
+
+  visit(node: IConstruct): void {
+    if (node.constructor.name === lambda.Function.name) {
+      this.datadog.addLambdaFunctions([node as LambdaFunction]);
+    }
+  }
+}
+
 export class Datadog extends Construct {
   scope: Construct;
   props: DatadogProps;

v2/test/datadog.spec.ts

@@ -1,8 +1,14 @@
-import { App, Stack, Token } from "aws-cdk-lib";
+import { App, Stack, Token, Aspects } 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, DD_HANDLER_ENV_VAR } from "../src/index";
+import {
+  addCdkConstructVersionTag,
+  checkForMultipleApiKeys,
+  Datadog,
+  DatadogAspect,
+  DD_HANDLER_ENV_VAR,
+} from "../src/index";
 const { ISecret } = require("aws-cdk-lib/aws-secretsmanager");
 const versionJson = require("../version.json");
 const EXTENSION_LAYER_VERSION = 5;
@@ -666,3 +672,37 @@ describe("redirectHandler", () => {
     });
   });
 });
+
+describe("DatadogAspect", () => {
+  it("redirects handler", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack", {
+      env: {
+        region: "us-west-2",
+      },
+    });
+    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", {
+      nodeLayerVersion: NODE_LAYER_VERSION,
+      extensionLayerVersion: EXTENSION_LAYER_VERSION,
+      apiKey: "1234",
+    });
+
+    Aspects.of(stack).add(new DatadogAspect(datadogCDK));
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_HANDLER_ENV_VAR]: "hello.handler",
+        },
+      },
+      Layers: [
+        "arn:aws:lambda:us-west-2:464622532012:layer:Datadog-Node16-x:91",
+        "arn:aws:lambda:us-west-2:464622532012:layer:Datadog-Extension:5",
+      ],
+    });
+  });
+});