Add merge_xray_traces.

Author: purple4reina

datadog_lambda/config.py

@@ -42,6 +42,7 @@ class Config:
         "DD_INSTRUMENTATION_TELEMETRY_ENABLED", "false", as_bool
     )
     trace_enabled = _get_env("DD_TRACE_ENABLED", "true", as_bool)
+    merge_xray_traces = _get_env("DD_MERGE_XRAY_TRACES", "false", as_bool)
 
     @property
     def fips_mode_enabled(self):

datadog_lambda/wrapper.py

@@ -62,7 +62,6 @@
 logger = logging.getLogger(__name__)
 
 DD_LOGS_INJECTION = "DD_LOGS_INJECTION"
-DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES"
 AWS_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME"
 DD_LOCAL_TEST = "DD_LOCAL_TEST"
 DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR"
@@ -157,9 +156,6 @@ def __init__(self, func):
         """Executes when the wrapped function gets wrapped"""
         try:
             self.func = func
-            self.merge_xray_traces = (
-                os.environ.get(DD_MERGE_XRAY_TRACES, "false").lower() == "true"
-            )
             self.function_name = os.environ.get(AWS_LAMBDA_FUNCTION_NAME, "function")
             self.service = os.environ.get(DD_SERVICE, None)
             self.extractor_env = os.environ.get(DD_TRACE_EXTRACTOR, None)
@@ -311,7 +307,7 @@ def _before(self, event, context):
                 )
 
             if config.trace_enabled:
-                set_dd_trace_py_root(trace_context_source, self.merge_xray_traces)
+                set_dd_trace_py_root(trace_context_source, config.merge_xray_traces)
                 if self.make_inferred_span:
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, self.decode_authorizer_context
@@ -322,7 +318,7 @@ def _before(self, event, context):
                     is_cold_start=is_cold_start(),
                     is_proactive_init=is_proactive_init(),
                     trace_context_source=trace_context_source,
-                    merge_xray_traces=self.merge_xray_traces,
+                    merge_xray_traces=config.merge_xray_traces,
                     trigger_tags=self.trigger_tags,
                     parent_span=self.inferred_span,
                     span_pointers=calculate_span_pointers(event_source, event),

tests/test_config.py

@@ -115,6 +115,15 @@ def set_env(key, value):
     ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "1", True),  # CHANGED
     ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "0", False),
     ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "purple", False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", None, False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "", False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "true", True),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "TRUE", True),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "false", False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "FALSE", False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "1", True),  # CHANGED
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "0", False),
+    ("DD_MERGE_XRAY_TRACES", "merge_xray_traces", "purple", False),
 )