Add trace_extractor.

Author: purple4reina

datadog_lambda/config.py

@@ -44,6 +44,7 @@ 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")
 
     @property
     def fips_mode_enabled(self):

datadog_lambda/wrapper.py

@@ -61,9 +61,7 @@
 
 logger = logging.getLogger(__name__)
 
-DD_LOGS_INJECTION = "DD_LOGS_INJECTION"
 DD_LOCAL_TEST = "DD_LOCAL_TEST"
-DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR"
 DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES"
 DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT"
 DD_DECODE_AUTHORIZER_CONTEXT = "DD_DECODE_AUTHORIZER_CONTEXT"
@@ -155,7 +153,6 @@ def __init__(self, func):
         """Executes when the wrapped function gets wrapped"""
         try:
             self.func = func
-            self.extractor_env = os.environ.get(DD_TRACE_EXTRACTOR, None)
             self.trace_extractor = None
             self.span = None
             self.inferred_span = None
@@ -194,8 +191,8 @@ def __init__(self, func):
             self.response = None
             if profiling_env_var:
                 self.prof = profiler.Profiler(env=env_env_var, service=config.service)
-            if self.extractor_env:
-                extractor_parts = self.extractor_env.rsplit(".", 1)
+            if config.trace_extractor:
+                extractor_parts = config.trace_extractor.rsplit(".", 1)
                 if len(extractor_parts) == 2:
                     (mod_name, extractor_name) = extractor_parts
                     modified_extractor_name = modify_module_name(mod_name)

tests/test_config.py

@@ -127,6 +127,9 @@ def set_env(key, value):
     ("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),
+    ("DD_TRACE_EXTRACTOR", "trace_extractor", None, None),
+    ("DD_TRACE_EXTRACTOR", "trace_extractor", "", ""),
+    ("DD_TRACE_EXTRACTOR", "trace_extractor", "my_extractor", "my_extractor"),
 )