Add capture_payload_max_depth.

Author: purple4reina

datadog_lambda/config.py

@@ -6,14 +6,26 @@
 import logging
 import os
 
+logger = logging.getLogger(__name__)
+
 
 def _get_env(key, default=None, cast=None):
     @property
     def _getter(self):
         if not hasattr(self, prop_key):
             val = os.environ.get(key, default)
             if cast is not None:
-                val = cast(val)
+                try:
+                    val = cast(val)
+                except (ValueError, TypeError):
+                    logger.warning(
+                        "Failed to cast environment variable '%s' with value '%s' to type %s. Using default value '%s'.",
+                        key,
+                        val,
+                        cast.__name__,
+                        default,
+                    )
+                    val = default
             setattr(self, prop_key, val)
         return getattr(self, prop_key)
 
@@ -45,6 +57,9 @@ 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
+    )
 
     @property
     def fips_mode_enabled(self):
@@ -67,7 +82,6 @@ def _reset(self):
 config = Config()
 
 if config.is_gov_region or config.fips_mode_enabled:
-    logger = logging.getLogger(__name__)
     logger.debug(
         "Python Lambda Layer FIPS mode is %s.",
         "enabled" if config.fips_mode_enabled else "not enabled",

datadog_lambda/tag_object.py

@@ -4,18 +4,17 @@
 # Copyright 2021 Datadog, Inc.
 
 from decimal import Decimal
-import logging
 import ujson as json
 
+from datadog_lambda.config import config
+
 redactable_keys = ["authorization", "x-authorization", "password", "token"]
-max_depth = 10
-logger = logging.getLogger(__name__)
 
 
 def tag_object(span, key, obj, depth=0):
     if obj is None:
         return span.set_tag(key, obj)
-    if depth >= max_depth:
+    if depth >= config.capture_payload_max_depth:
         return span.set_tag(key, _redact_val(key, str(obj)[0:5000]))
     depth += 1
     if _should_try_string(obj):

datadog_lambda/wrapper.py

@@ -69,7 +69,6 @@
 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_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH = "DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH"
 DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
 DD_SERVICE = "DD_SERVICE"
 DD_ENV = "DD_ENV"
@@ -89,13 +88,6 @@ def get_env_as_int(env_key, default_value: int) -> int:
     os.environ.get(DD_CAPTURE_LAMBDA_PAYLOAD, "false").lower() == "true"
 )
 
-if dd_capture_lambda_payload_enabled:
-    import datadog_lambda.tag_object as tag_object
-
-    tag_object.max_depth = get_env_as_int(
-        DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH, tag_object.max_depth
-    )
-
 env_env_var = os.environ.get(DD_ENV, None)
 
 init_timestamp_ns = time_ns()

tests/test_config.py

@@ -130,6 +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),
 )
 
 

tests/test_tag_object.py

@@ -29,6 +29,7 @@ def test_tag_object(self):
             True,
         )
 
+    @patch("datadog_lambda.config.Config.capture_payload_max_depth", 2)
     def test_tag_object_max_depth(self):
         payload = {
             "hello": "world",
@@ -41,11 +42,8 @@ def test_tag_object_max_depth(self):
             "vals": [{"thingOne": 1}, {"thingTwo": 2}],
         }
         spanMock = MagicMock()
-        import datadog_lambda.tag_object as lib_ref
 
-        lib_ref.max_depth = 2  # setting up the test
         tag_object(spanMock, "function.request", payload)
-        lib_ref.max_depth = 10  # revert the setup
         spanMock.set_tag.assert_has_calls(
             [
                 call("function.request.vals.0", "{'thingOne': 1}"),
@@ -62,6 +60,7 @@ def test_tag_object_max_depth(self):
             True,
         )
 
+    @patch("datadog_lambda.config.Config.capture_payload_max_depth", 0)
     def test_tag_object_max_depth_0(self):
         payload = {
             "hello": "world",
@@ -74,11 +73,8 @@ def test_tag_object_max_depth_0(self):
             "vals": [{"thingOne": 1}, {"thingTwo": 2}],
         }
         spanMock = MagicMock()
-        import datadog_lambda.tag_object as lib_ref
 
-        lib_ref.max_depth = 0  # setting up the test
         tag_object(spanMock, "function.request", payload)
-        lib_ref.max_depth = 10  # revert the setup
         spanMock.set_tag.assert_has_calls(
             [
                 call(