Add capture_payload_enabled.

Author: purple4reina

datadog_lambda/config.py

@@ -59,7 +59,8 @@ class Config:
     trace_enabled = _get_env("DD_TRACE_ENABLED", "true", as_bool)
     merge_xray_traces = _get_env("DD_MERGE_XRAY_TRACES", "false", as_bool)
     trace_extractor = _get_env("DD_TRACE_EXTRACTOR")
-    capture_payload_max_depth = _get_env("DD_CAPTURE_PAYLOAD_MAX_DEPTH", 10, int)
+    capture_payload_max_depth = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", 10, int)
+    capture_payload_enabled = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD", "false", as_bool)
     profiling_enabled = _get_env("DD_PROFILING_ENABLED", "false", as_bool)
     llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
     exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool)

datadog_lambda/wrapper.py

@@ -65,7 +65,6 @@
 DD_COLD_START_TRACING = "DD_COLD_START_TRACING"
 DD_MIN_COLD_START_DURATION = "DD_MIN_COLD_START_DURATION"
 DD_COLD_START_TRACE_SKIP_LIB = "DD_COLD_START_TRACE_SKIP_LIB"
-DD_CAPTURE_LAMBDA_PAYLOAD = "DD_CAPTURE_LAMBDA_PAYLOAD"
 DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
 DD_SERVICE = "DD_SERVICE"
 DD_ENV = "DD_ENV"
@@ -82,10 +81,6 @@ def get_env_as_int(env_key, default_value: int) -> int:
         return default_value
 
 
-dd_capture_lambda_payload_enabled = (
-    os.environ.get(DD_CAPTURE_LAMBDA_PAYLOAD, "false").lower() == "true"
-)
-
 init_timestamp_ns = time_ns()
 
 """
@@ -336,7 +331,7 @@ def _after(self, event, context):
                 trace_ctx = tracer.current_trace_context()
 
             if self.span:
-                if dd_capture_lambda_payload_enabled:
+                if config.capture_payload_enabled:
                     tag_object.tag_object(self.span, "function.request", event)
                     tag_object.tag_object(self.span, "function.response", self.response)
 

tests/test_config.py

@@ -130,13 +130,13 @@ def set_env(key, value):
     ("DD_TRACE_EXTRACTOR", "trace_extractor", None, None),
     ("DD_TRACE_EXTRACTOR", "trace_extractor", "", ""),
     ("DD_TRACE_EXTRACTOR", "trace_extractor", "my_extractor", "my_extractor"),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", None, 10),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "", 10),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "5", 5),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "0", 0),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "2.5", 10),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "-1", -1),
-    ("DD_CAPTURE_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "purple", 10),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", None, 10),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "", 10),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "5", 5),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "0", 0),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "2.5", 10),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "-1", -1),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", "purple", 10),
     ("DD_PROFILING_ENABLED", "profiling_enabled", None, False),
     ("DD_PROFILING_ENABLED", "profiling_enabled", "", False),
     ("DD_PROFILING_ENABLED", "profiling_enabled", "true", True),
@@ -167,6 +167,15 @@ def set_env(key, value):
     ("DD_ENV", "env", None, None),
     ("DD_ENV", "env", "", ""),
     ("DD_ENV", "env", "my_env", "my_env"),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", None, False),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "", False),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "true", True),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "TRUE", True),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "false", False),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "FALSE", False),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "1", True),  # CHANGED
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "0", False),
+    ("DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", "purple", False),
 )