[POC] Add DD_TRACE_EXTRACTOR env var to customize Trace extraction customization (#112)

* add extractor and tests

* lint and other fixes

* fix lint

* Update tests

* Fix format

Co-authored-by: Tian Chu <tian.chu@datadoghq.com>

Author: Czechh

datadog_lambda/tracing.py

@@ -165,7 +165,20 @@ def extract_context_from_sqs_event_or_context(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
-def extract_dd_trace_context(event, lambda_context):
+def extract_context_custom_extractor(extractor, event, lambda_context):
+    """
+    Extract Datadog trace context using a custom trace extractor function
+    """
+    try:
+        (trace_id, parent_id, sampling_priority,) = extractor(event, lambda_context)
+        return trace_id, parent_id, sampling_priority
+    except Exception as e:
+        logger.debug("The trace extractor returned with error %s", e)
+
+        return None, None, None
+
+
+def extract_dd_trace_context(event, lambda_context, extractor=None):
     """
     Extract Datadog trace context from the Lambda `event` object.
 
@@ -175,7 +188,11 @@ def extract_dd_trace_context(event, lambda_context):
     global dd_trace_context
     trace_context_source = None
 
-    if "headers" in event:
+    if extractor is not None:
+        (trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor(
+            extractor, event, lambda_context
+        )
+    elif "headers" in event:
         (
             trace_id,
             parent_id,

datadog_lambda/wrapper.py

@@ -6,6 +6,7 @@
 import os
 import logging
 import traceback
+from importlib import import_module
 
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import set_cold_start, is_cold_start
@@ -15,6 +16,7 @@
     submit_invocations_metric,
     submit_errors_metric,
 )
+from datadog_lambda.module_name import modify_module_name
 from datadog_lambda.patch import patch_all
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
@@ -91,6 +93,16 @@ def __init__(self, func):
                 os.environ.get("DD_MERGE_XRAY_TRACES", "false").lower() == "true"
             )
             self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function")
+            self.extractor_env = os.environ.get("DD_TRACE_EXTRACTOR", None)
+            self.trace_extractor = None
+
+            if self.extractor_env:
+                extractor_parts = self.extractor_env.rsplit(".", 1)
+                if len(extractor_parts) == 2:
+                    (mod_name, extractor_name) = extractor_parts
+                    modified_extractor_name = modify_module_name(mod_name)
+                    extractor_module = import_module(modified_extractor_name)
+                    self.trace_extractor = getattr(extractor_module, extractor_name)
 
             # Inject trace correlation ids to logs
             if self.logs_injection:
@@ -125,7 +137,9 @@ def _before(self, event, context):
             set_cold_start()
             submit_invocations_metric(context)
             # Extract Datadog trace context and source from incoming requests
-            dd_context, trace_context_source = extract_dd_trace_context(event, context)
+            dd_context, trace_context_source = extract_dd_trace_context(
+                event, context, extractor=self.trace_extractor
+            )
             # Create a Datadog X-Ray subsegment with the trace context
             if dd_context and trace_context_source == TraceContextSource.EVENT:
                 create_dd_dummy_metadata_subsegment(

tests/integration/serverless-plugin.yml

@@ -6,6 +6,7 @@ provider:
     DD_INTEGRATION_TEST: true
     DD_API_KEY: ${env:DD_API_KEY}
     WITH_PLUGIN: true
+  lambdaHashingVersion: 20201221
 
 layers:
   python27:

tests/integration/serverless.yml

@@ -8,6 +8,7 @@ provider:
   environment:
     DD_INTEGRATION_TEST: true
     DD_API_KEY: ${env:DD_API_KEY}
+  lambdaHashingVersion: 20201221
 
 layers:
   python27:

tests/test_tracing.py

@@ -136,6 +136,70 @@ def test_with_complete_datadog_trace_headers(self):
             XraySubsegment.NAMESPACE,
         )
 
+    def test_with_extractor_function(self):
+        def extractor_foo(event, context):
+            foo = event.get("foo", {})
+            lowercase_foo = {k.lower(): v for k, v in foo.items()}
+
+            trace_id = lowercase_foo.get(TraceHeader.TRACE_ID)
+            parent_id = lowercase_foo.get(TraceHeader.PARENT_ID)
+            sampling_priority = lowercase_foo.get(TraceHeader.SAMPLING_PRIORITY)
+            return trace_id, parent_id, sampling_priority
+
+        lambda_ctx = get_mock_context()
+        ctx, ctx_source = extract_dd_trace_context(
+            {
+                "foo": {
+                    TraceHeader.TRACE_ID: "123",
+                    TraceHeader.PARENT_ID: "321",
+                    TraceHeader.SAMPLING_PRIORITY: "1",
+                }
+            },
+            lambda_ctx,
+            extractor=extractor_foo,
+        )
+        self.assertEquals(ctx_source, "event")
+        self.assertDictEqual(
+            ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",},
+        )
+        self.assertDictEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "123",
+                TraceHeader.PARENT_ID: "65535",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+            },
+        )
+
+    def test_graceful_fail_of_extractor_function(self):
+        def extractor_raiser(event, context):
+            raise Exception("kreator")
+
+        lambda_ctx = get_mock_context()
+        ctx, ctx_source = extract_dd_trace_context(
+            {
+                "foo": {
+                    TraceHeader.TRACE_ID: "123",
+                    TraceHeader.PARENT_ID: "321",
+                    TraceHeader.SAMPLING_PRIORITY: "1",
+                }
+            },
+            lambda_ctx,
+            extractor=extractor_raiser,
+        )
+        self.assertEquals(ctx_source, "xray")
+        self.assertDictEqual(
+            ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2",},
+        )
+        self.assertDictEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "4369",
+                TraceHeader.PARENT_ID: "65535",
+                TraceHeader.SAMPLING_PRIORITY: "2",
+            },
+        )
+
     def test_with_sqs_distributed_datadog_trace_data(self):
         lambda_ctx = get_mock_context()
         sqs_event = {

tests/test_wrapper.py

@@ -100,7 +100,7 @@ def lambda_handler(event, context):
         )
         self.mock_wrapper_lambda_stats.flush.assert_called()
         self.mock_extract_dd_trace_context.assert_called_with(
-            lambda_event, lambda_context
+            lambda_event, lambda_context, extractor=None
         )
         self.mock_set_correlation_ids.assert_called()
         self.mock_inject_correlation_ids.assert_called()