expose DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH [SVLS-3853] (#387)

* expose DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH

* Update README.md

Co-authored-by: Brett Blue <[email protected]>

* improve readme.md

---------

Co-authored-by: Brett Blue <[email protected]>

Author: joeyzhao2018

README.md

@@ -16,6 +16,21 @@ Follow the [installation instructions](https://docs.datadoghq.com/serverless/ins
 
 Follow the [configuration instructions](https://docs.datadoghq.com/serverless/configuration) to tag your telemetry, capture request/response payloads, filter or scrub sensitive information from logs or traces, and more.
 
+For additional tracing configuration options, check out the [official documentation for Datadog trace client](https://ddtrace.readthedocs.io/en/stable/configuration.html).
+
+Besides the environment variables supported by dd-trace-py, the datadog-lambda-python library added following environment variables.
+
+| Environment Variables | Description | Default Value |
+| -------------------- | ------------ | ------------- |
+| DD_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. | `true` |
+| DD_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. | `true` |
+| DD_COLD_START_TRACING | Set to `false` to disable Cold Start Tracing. Used in NodeJS and Python. | `true` |
+| DD_MIN_COLD_START_DURATION |  Sets the minimum duration (in milliseconds) for a module load event to be traced via Cold Start Tracing. Number. | `3` |
+| DD_COLD_START_TRACE_SKIP_LIB | optionally skip creating Cold Start Spans for a comma-separated list of libraries. Useful to limit depth or skip known libraries. | `ddtrace.internal.compat,ddtrace.filters` |
+| DD_CAPTURE_LAMBDA_PAYLOAD | [Captures incoming and outgoing AWS Lambda payloads][1] in the Datadog APM spans for Lambda invocations. | `false` |
+| DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH | Determines the level of detail captured from AWS Lambda payloads, which are then assigned as tags for the `aws.lambda` span. It specifies the nesting depth of the JSON payload structure to process. Once the specified maximum depth is reached, the tag's value is set to the stringified value of any nested elements beyond this level.  <br> For example, given the input payload: <pre>{<br>  "lv1" : {<br>    "lv2": {<br>      "lv3": "val"<br>    }<br>  }<br>}</pre> If the depth is set to `2`, the resulting tag's key is set to `function.request.lv1.lv2` and the value is `{\"lv3\": \"val\"}`. <br> If the depth is set to `0`, the resulting tag's key is set to `function.request` and value is `{\"lv1\":{\"lv2\":{\"lv3\": \"val\"}}}` | `10` |
+
+
 ## Opening Issues
 
 If you encounter a bug with this package, we want to hear about it. Before opening a new issue, search the existing issues to avoid duplicates.
@@ -51,3 +66,5 @@ For product feedback and questions, join the `#serverless` channel in the [Datad
 Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
 
 This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2019 Datadog, Inc.
+
+[1]: https://www.datadoghq.com/blog/troubleshoot-lambda-function-request-response-payloads/

datadog_lambda/tag_object.py

@@ -13,12 +13,11 @@
 
 
 def tag_object(span, key, obj, depth=0):
-    if depth >= max_depth:
-        return
-    else:
-        depth += 1
     if obj is None:
         return span.set_tag(key, obj)
+    if depth >= max_depth:
+        return tag_object(span, key, _redact_val(key, str(obj)[0:5000]))
+    depth += 1
     if _should_try_string(obj):
         parsed = None
         try:

datadog_lambda/wrapper.py

@@ -49,18 +49,13 @@
     extract_trigger_tags,
     extract_http_status_code_tag,
 )
-from datadog_lambda.tag_object import tag_object
 
 profiling_env_var = os.environ.get("DD_PROFILING_ENABLED", "false").lower() == "true"
 if profiling_env_var:
     from ddtrace.profiling import profiler
 
 logger = logging.getLogger(__name__)
 
-dd_capture_lambda_payload_enabled = (
-    os.environ.get("DD_CAPTURE_LAMBDA_PAYLOAD", "false").lower() == "true"
-)
-
 DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG"
 DD_LOGS_INJECTION = "DD_LOGS_INJECTION"
 DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES"
@@ -72,10 +67,34 @@
 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_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"
 
+
+def get_env_as_int(env_key, default_value: int) -> int:
+    try:
+        return int(os.environ.get(env_key, default_value))
+    except Exception as e:
+        logger.warn(
+            f"Failed to parse {env_key} as int. Using default value: {default_value}. Error: {e}"
+        )
+        return default_value
+
+
+dd_capture_lambda_payload_enabled = (
+    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()
@@ -161,14 +180,9 @@ def __init__(self, func):
             self.cold_start_tracing = depends_on_dd_tracing_enabled(
                 os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
             )
-            self.min_cold_start_trace_duration = 3
-            if DD_MIN_COLD_START_DURATION in os.environ:
-                try:
-                    self.min_cold_start_trace_duration = int(
-                        os.environ[DD_MIN_COLD_START_DURATION]
-                    )
-                except Exception:
-                    logger.debug(f"Malformatted env {DD_MIN_COLD_START_DURATION}")
+            self.min_cold_start_trace_duration = get_env_as_int(
+                DD_MIN_COLD_START_DURATION, 3
+            )
             self.cold_start_trace_skip_lib = [
                 "ddtrace.internal.compat",
                 "ddtrace.filters",
@@ -307,16 +321,14 @@ def _after(self, event, context):
                 create_dd_dummy_metadata_subsegment(
                     self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
                 )
-            should_trace_cold_start = (
-                dd_tracing_enabled and self.cold_start_tracing and is_new_sandbox()
-            )
+            should_trace_cold_start = self.cold_start_tracing and is_new_sandbox()
             if should_trace_cold_start:
                 trace_ctx = tracer.current_trace_context()
 
             if self.span:
                 if dd_capture_lambda_payload_enabled:
-                    tag_object(self.span, "function.request", event)
-                    tag_object(self.span, "function.response", self.response)
+                    tag_object.tag_object(self.span, "function.request", event)
+                    tag_object.tag_object(self.span, "function.response", self.response)
 
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)

datadog_lambda/xray.py

@@ -75,7 +75,6 @@ def generate_random_id():
 
 
 def build_segment(context, key, metadata):
-
     segment = json.dumps(
         {
             "id": generate_random_id(),

scripts/run_integration_tests.sh

@@ -200,8 +200,8 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 sed -E "s/(api_key=|'api_key': '|DD-API-KEY:)[a-z0-9\.\-]+/\1XXXX/g" |
                 # Normalize package version so that these snapshots aren't broken on version bumps
                 sed -E "s/(dd_lambda_layer:datadog-python[0-9]+_)[0-9]+\.[0-9]+\.[0-9]+/\1X\.X\.X/g" |
-                sed -E "s/(datadog_lambda:v)([0-9]+\.[0-9]+\.[0-9])/\1XX/g" |
-                sed -E "s/(datadogpy\/)([0-9]+\.[0-9]+\.[0-9])/\1XX/g" |
+                sed -E "s/(datadog_lambda:v)([0-9]+\.[0-9]+\.[0-9]+)/\1XX/g" |
+                sed -E "s/(datadogpy\/)([0-9]+\.[0-9]+\.[0-9]+)/\1XX/g" |
                 sed -E "s/(python )([0-9]\.[0-9]+\.[0-9]+)/\1XX/g" |
                 # Strip out run ID (from function name, resource, etc.)
                 sed -E "s/${!run_id}/XXXX/g" |
@@ -231,7 +231,7 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 sed -E "s/(\"connection_id\"\:\ \")[a-zA-Z0-9\-]+/\1XXXX/g" |
                 sed -E "s/(\"shardId\-)([0-9]+)\:([a-zA-Z0-9]+)[a-zA-Z0-9]/\1XXXX:XXXX/g" |
                 sed -E "s/(\"shardId\-)[0-9a-zA-Z]+/\1XXXX/g" |
-                sed -E "s/(\"datadog_lambda\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g" |
+                sed -E "s/(\"datadog_lambda\"\: \")([0-9]+\.[0-9]+\.[0-9]+)/\1X.X.X/g" |
                 sed -E "s/(\"partition_key\"\:\ \")[a-zA-Z0-9\-]+/\1XXXX/g" |
                 sed -E "s/(\"object_etag\"\:\ \")[a-zA-Z0-9\-]+/\1XXXX/g" |
                 sed -E "s/(\"dd_trace\"\: \")([0-9]+\.[0-9]+\.[0-9]+)/\1X.X.X/g" |

tests/test_tag_object.py

@@ -28,7 +28,39 @@ def test_tag_object(self):
             ],
             True,
         )
-        self.assertEqual(1, 1)
+
+    def test_tag_object_max_depth(self):
+        payload = {
+            "hello": "world",
+            "level1": {
+                "level2_dict": {"level3": 3},
+                "level2_list": [None, True, "nice", {"l3": "v3"}],
+                "level2_bool": True,
+                "level2_int": 2,
+            },
+            "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}"),
+                call("function.request.vals.1", "{'thingTwo': 2}"),
+                call("function.request.hello", "world"),
+                call("function.request.level1.level2_dict", "{'level3': 3}"),
+                call(
+                    "function.request.level1.level2_list",
+                    "[None, True, 'nice', {'l3': 'v3'}]",
+                ),
+                call("function.request.level1.level2_bool", "True"),
+                call("function.request.level1.level2_int", "2"),
+            ],
+            True,
+        )
 
     def test_redacted_tag_object(self):
         payload = {