Feat: Add ability to tag Lambda request and response payloads for 2.7… (#180)

* Update expected breaking change date (#114)

* Feat: Add ability to tag Lambda request and response payloads for 2.7 and 3.x

* feat: fix names of functions. Add tests for redacted keys as well as deeply tagging a json object

* feat: test coverage for should_try_string across py versions

* feat: add black formatter, used in ci, to vscode settings

* lint: run black

* feat: add integration test support

Co-authored-by: Jorie Helwig <[email protected]>

Author: astuyve

.vscode/settings.json

@@ -1,3 +1,4 @@
 {
-    "python.pythonPath": "/usr/local/bin/python3"
-}
+  "python.pythonPath": "/usr/local/bin/python3",
+  "python.formatting.provider": "black"
+}

CONTRIBUTING.md

@@ -1,21 +1,23 @@
 # Contributing
 
-We love pull requests. For new features, consider opening an issue to discuss the idea first. When you're ready to open a pull requset, here's a quick guide.
+We love pull requests. For new features, consider opening an issue to discuss the idea first. When you're ready to open a pull request, here's a quick guide.
 
 1. Fork, clone and branch off `main`:
-    ```bash
-    git clone [email protected]:<your-username>/datadog-lambda-python.git
-    git checkout -b <my-branch>
-    ```
-1. Make your changes. Ensure your code is compatible with both Python 2.7 and 3.X. 
+   ```bash
+   git clone [email protected]:<your-username>/datadog-lambda-python.git
+   git checkout -b <my-branch>
+   ```
+1. Make your changes. Ensure your code is compatible with both Python 2.7 and 3.X.
 1. Test your Lambda function against the locally modified version of Datadog Lambda library.
-   * The easiest approach is to create a soft link of the `datadog_lambda` folder in your project's root. Note, this only overrides the `datadog_lambda` module, and you still need to install the `datadog_lambda` package or the Lambda layer to have the required dependencies.
+
+   - The easiest approach is to create a soft link of the `datadog_lambda` folder in your project's root. Note, this only overrides the `datadog_lambda` module, and you still need to install the `datadog_lambda` package or the Lambda layer to have the required dependencies.
 
      ```bash
      ln -s /PATH/TO/datadog-lambda-python/datadog_lambda /PATH/TO/MY/PROJECT
      ```
-   * Another option is to install the `datadog_lambda` module from the local folder. E.g., add `/PATH/TO/datadog-lambda-python/` to your `requirements.txt`. This approach only work in a Linux environment, because the dependency `ddtrace` utilizes the native C extension.
-   * You can also build and publish a Lambda layer to your own AWS account and use it for testing.
+
+   - Another option is to install the `datadog_lambda` module from the local folder. E.g., add `/PATH/TO/datadog-lambda-python/` to your `requirements.txt`. This approach only work in a Linux environment, because the dependency `ddtrace` utilizes the native C extension.
+   - You can also build and publish a Lambda layer to your own AWS account and use it for testing.
 
      ```bash
      # Build layers using docker
@@ -27,9 +29,9 @@ We love pull requests. For new features, consider opening an issue to discuss th
      ```
 
 1. Ensure the unit tests pass (install Docker if you haven't):
-    ```bash
-    ./scripts/run_tests.sh
-    ```
+   ```bash
+   ./scripts/run_tests.sh
+   ```
 1. Run the integration tests against your own AWS account and Datadog org (or ask a Datadog member to run):
    ```bash
    BUILD_LAYERS=true DD_API_KEY=<your Datadog api key> ./scripts/run_integration_tests.sh

README.md

@@ -6,9 +6,9 @@
 [![Slack](https://chat.datadoghq.com/badge.svg?bg=632CA6)](https://chat.datadoghq.com/)
 [![License](https://img.shields.io/badge/license-Apache--2.0-blue)](https://github.com/DataDog/datadog-lambda-python/blob/main/LICENSE)
 
-Datadog Lambda Library for Python (2.7, 3.6, 3.7, 3.8, and 3.9) enables enhanced Lambda metrics, distributed tracing, and custom metric submission from AWS Lambda functions.  
+Datadog Lambda Library for Python (2.7, 3.6, 3.7, 3.8, and 3.9) enables enhanced Lambda metrics, distributed tracing, and custom metric submission from AWS Lambda functions.
 
-**IMPORTANT NOTE:** AWS Lambda is expected to receive a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **December 1, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to the latest.
+**IMPORTANT NOTE:** AWS Lambda is expected to receive a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **March 31, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to the latest.
 
 ## Installation
 
@@ -24,7 +24,7 @@ Check out the instructions for [submitting custom metrics from AWS Lambda functi
 
 Once [installed](#installation), you should be able to view your function's traces in Datadog, and your function's logs should be automatically connected to the traces.
 
-For additional details on trace collection, take a look at [collecting traces from AWS Lambda functions](https://docs.datadoghq.com/integrations/amazon_lambda/?tab=python#trace-collection). 
+For additional details on trace collection, take a look at [collecting traces from AWS Lambda functions](https://docs.datadoghq.com/integrations/amazon_lambda/?tab=python#trace-collection).
 
 For additional details on trace and log connection, see [connecting logs and traces](https://docs.datadoghq.com/tracing/connect_logs_and_traces/python/).
 

datadog_lambda/tag_object.py

@@ -0,0 +1,59 @@
+# 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 2021 Datadog, Inc.
+
+import json
+import logging
+
+redactable_keys = ["authorization", "x-authorization", "password", "token"]
+max_depth = 10
+logger = logging.getLogger(__name__)
+
+
+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 _should_try_string(obj):
+        parsed = None
+        try:
+            parsed = json.loads(obj)
+            return tag_object(span, key, parsed, depth)
+        except ValueError:
+            redacted = _redact_val(key, obj[0:5000])
+            return span.set_tag(key, redacted)
+    if isinstance(obj, int) or isinstance(obj, float):
+        return span.set_tag(key, obj)
+    if isinstance(obj, list):
+        for k, v in enumerate(obj):
+            formatted_key = "{}.{}".format(key, k)
+            tag_object(span, formatted_key, v, depth)
+        return
+    if isinstance(obj, object):
+        for k in obj:
+            v = obj.get(k)
+            formatted_key = "{}.{}".format(key, k)
+            tag_object(span, formatted_key, v, depth)
+        return
+
+
+def _should_try_string(obj):
+    try:
+        if isinstance(obj, str) or isinstance(obj, unicode):
+            return True
+    except NameError:
+        if isinstance(obj, bytes):
+            return True
+
+    return False
+
+
+def _redact_val(k, v):
+    split_key = k.split(".").pop() or k
+    if split_key in redactable_keys:
+        return "redacted"
+    return v

datadog_lambda/tracing.py

@@ -261,7 +261,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
 
 def get_dd_trace_context():
     """
-    Return the Datadog trace context to be propogated on the outgoing requests.
+    Return the Datadog trace context to be propagated on the outgoing requests.
 
     If the Lambda function is invoked by a Datadog-traced service, a Datadog
     trace context may already exist, and it should be used. Otherwise, use the

datadog_lambda/wrapper.py

@@ -28,9 +28,13 @@
     create_function_execution_span,
 )
 from datadog_lambda.trigger import extract_trigger_tags, extract_http_status_code_tag
+from datadog_lambda.tag_object import tag_object
 
 logger = logging.getLogger(__name__)
 
+dd_capture_lambda_payload_enabled = (
+    os.environ.get("DD_CAPTURE_LAMBDA_PAYLOAD", "false").lower() == "true"
+)
 
 """
 Usage:
@@ -186,6 +190,10 @@ def _after(self, event, context):
                 flush_extension()
 
             if self.span:
+                if dd_capture_lambda_payload_enabled:
+                    tag_object(self.span, "function.request", event)
+                    tag_object(self.span, "function.response", self.response)
+
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)
                 self.span.finish()

tests/integration/serverless.yml

@@ -16,6 +16,7 @@ provider:
     DD_INTEGRATION_TEST: true
     DD_TRACE_ENABLED: true
     DD_API_KEY: ${env:DD_API_KEY}
+    DD_CAPTURE_LAMBDA_PAYLOAD: true
   lambdaHashingVersion: 20201221
   timeout: 15
   deploymentBucket: