[SVLS-9359] add setEnvironment method to seed construct-tracked env vars

Author: ava-silver

README.md

@@ -404,6 +404,8 @@ The DatadogLambda construct takes in a list of lambda functions and installs the
 
 **Configuring `DD_*` environment variables:** Set Datadog environment variables through `DatadogLambdaProps` (e.g. `enableDatadogTracing`, `logLevel`, `tags`). If you need to set a `DD_*` variable directly on a function, call `func.addEnvironment()` **after** `addLambdaFunctions()` -- values set before will be overridden by the construct.
 
+To set a value the construct should respect (rather than override), use `datadogLambda.setEnvironment(func, key, value)` **before** `addLambdaFunctions()`. This is useful for setting `DD_TAGS` per function: when source code integration is enabled, the construct appends git metadata (`git.commit.sha`, `git.repository_url`) to the value you set instead of overriding it.
+
 While Lambda function based log groups are handled by the `addLambdaFunctions` method automatically, the construct has an additional function `addForwarderToNonLambdaLogGroups` which subscribes the forwarder to any additional log groups of your choosing.
 
 ## Step Functions

src/datadog-lambda.ts

@@ -182,6 +182,20 @@ export class DatadogLambda extends Construct {
     }
   }
 
+  /**
+   * Sets an environment variable on the given Lambda function and records it in the
+   * construct's internal tracker, so the construct treats it as a value it manages.
+   *
+   * Call this before `addLambdaFunctions()` to seed values the construct will respect.
+   * The main use case is setting `DD_TAGS` per function: when source code integration is
+   * enabled, the construct appends git metadata (`git.commit.sha`, `git.repository_url`)
+   * to the tracked `DD_TAGS` rather than overriding it.
+   */
+  public setEnvironment(lambdaFunction: LambdaFunction, key: string, value: string): void {
+    const [extractedLambdaFunction] = extractSingletonFunctions([lambdaFunction]);
+    setTrackedEnv(extractedLambdaFunction, key, value);
+  }
+
   public overrideGitMetadata(gitCommitSha: string, gitRepoUrl?: string): void {
     if (gitCommitSha) {
       this.gitCommitShaOverride = gitCommitSha;

test/datadog-lambda.spec.ts

@@ -1096,3 +1096,74 @@ describe("overrideGitMetadata", () => {
     });
   });
 });
+
+describe("setEnvironment", () => {
+  it("appends git metadata to a DD_TAGS value seeded before addLambdaFunctions", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack");
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_18_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+    const datadogLambda = new DatadogLambda(stack, "Datadog", {
+      nodeLayerVersion: NODE_LAYER_VERSION,
+      extensionLayerVersion: EXTENSION_LAYER_VERSION,
+      apiKey: "ABC",
+      enableDatadogTracing: false,
+      flushMetricsToLogs: false,
+      logLevel: "debug",
+    });
+    datadogLambda.overrideGitMetadata("fake-sha", "fake-url");
+    datadogLambda.setEnvironment(hello, DD_TAGS, "team:serverless");
+    datadogLambda.addLambdaFunctions([hello], stack);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_TAGS]: Match.stringLikeRegexp("team:serverless"),
+        },
+      },
+    });
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_TAGS]: Match.stringLikeRegexp("git\\.commit\\.sha:fake-sha"),
+        },
+      },
+    });
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          [DD_TAGS]: Match.stringLikeRegexp("git\\.repository_url:fake-url"),
+        },
+      },
+    });
+  });
+
+  it("does not override a value seeded before addLambdaFunctions", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack");
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.NODEJS_18_X,
+      code: lambda.Code.fromInline("test"),
+      handler: "hello.handler",
+    });
+    const datadogLambda = new DatadogLambda(stack, "Datadog", {
+      nodeLayerVersion: NODE_LAYER_VERSION,
+      apiKey: "ABC",
+      enableDatadogTracing: true,
+      flushMetricsToLogs: false,
+    });
+    datadogLambda.setEnvironment(hello, "DD_TRACE_ENABLED", "false");
+    datadogLambda.addLambdaFunctions([hello], stack);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Environment: {
+        Variables: {
+          DD_TRACE_ENABLED: "false",
+        },
+      },
+    });
+  });
+});