Add cold_start_trace_skip_lib.

Author: purple4reina

datadog_lambda/config.py

@@ -37,13 +37,26 @@ def as_bool(val):
     return val.lower() == "true" or val == "1"
 
 
+def as_list(val):
+    if isinstance(val, str):
+        if not val:
+            return []
+        return [val.strip() for val in val.split(",") if val.strip()]
+    return val
+
+
 class Config:
 
     service = _get_env("DD_SERVICE")
     env = _get_env("DD_ENV")
 
     cold_start_tracing = _get_env("DD_COLD_START_TRACING", "true", as_bool)
     min_cold_start_trace_duration = _get_env("DD_MIN_COLD_START_DURATION", 3, int)
+    cold_start_trace_skip_lib = _get_env(
+            "DD_COLD_START_TRACE_SKIP_LIB", [
+                "ddtrace.internal.compat",
+                "ddtrace.filters",
+            ], as_list)
 
     capture_payload_max_depth = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", 10, int)
     capture_payload_enabled = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD", "false", as_bool)

datadog_lambda/wrapper.py

@@ -60,7 +60,6 @@
 DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT"
 DD_DECODE_AUTHORIZER_CONTEXT = "DD_DECODE_AUTHORIZER_CONTEXT"
 DD_COLD_START_TRACING = "DD_COLD_START_TRACING"
-DD_COLD_START_TRACE_SKIP_LIB = "DD_COLD_START_TRACE_SKIP_LIB"
 DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
 DD_SERVICE = "DD_SERVICE"
 DD_ENV = "DD_ENV"
@@ -139,17 +138,6 @@ def __init__(self, func):
             self.cold_start_tracing = depends_on_dd_tracing_enabled(
                 os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
             )
-            self.cold_start_trace_skip_lib = [
-                "ddtrace.internal.compat",
-                "ddtrace.filters",
-            ]
-            if DD_COLD_START_TRACE_SKIP_LIB in os.environ:
-                try:
-                    self.cold_start_trace_skip_lib = os.environ[
-                        DD_COLD_START_TRACE_SKIP_LIB
-                    ].split(",")
-                except Exception:
-                    logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}")
             self.response = None
             if config.profiling_enabled:
                 self.prof = profiler.Profiler(env=config.env, service=config.service)
@@ -334,7 +322,7 @@ def _after(self, event, context):
                         following_span.start_ns,
                         trace_ctx,
                         config.min_cold_start_trace_duration,
-                        self.cold_start_trace_skip_lib,
+                        config.cold_start_trace_skip_lib,
                     ).trace()
                 except Exception as e:
                     logger.debug("Failed to create cold start spans. %s", e)

tests/test_config.py

@@ -192,6 +192,16 @@ def set_env(key, value):
     ("DD_LOCAL_TEST", "local_test", "1", True),  # CHANGED
     ("DD_LOCAL_TEST", "local_test", "0", False),
     ("DD_LOCAL_TEST", "local_test", "purple", False),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", None, ["ddtrace.internal.compat", "ddtrace.filters"]),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "", []),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", " ", []),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", ",", []),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", " , ", []),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "a", ["a"]),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "a,", ["a"]),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "a, ", ["a"]),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "a,b", ["a", "b"]),
+    ("DD_COLD_START_TRACE_SKIP_LIB", "cold_start_trace_skip_lib", "a, b", ["a", "b"]),
 )
 
 

tests/test_wrapper.py

@@ -641,39 +641,40 @@ def handler(event, context):
 
 
 class TestLambdaWrapperFlushExtension(unittest.TestCase):
-    def setUp(self):
-        self.orig_environ = os.environ
+    @patch("datadog_lambda.config.Config.local_test", True)
+    @patch("datadog_lambda.wrapper.should_use_extension", True)
+    def test_local_test_true_flushing(self):
+        flushes = []
+        lambda_event = {}
+        lambda_context = get_mock_context()
 
-    def tearDown(self):
-        os.environ = self.orig_environ
+        def flush():
+            flushes.append(1)
 
+        @patch("datadog_lambda.wrapper.flush_extension", flush)
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            pass
+
+        lambda_handler(lambda_event, lambda_context)
+
+        self.assertEqual(len(flushes), 1)
+
+    @patch("datadog_lambda.config.Config.local_test", False)
     @patch("datadog_lambda.wrapper.should_use_extension", True)
-    def test_local_test_envvar_flushing(self):
+    def test_local_test_false_flushing(self):
         flushes = []
         lambda_event = {}
         lambda_context = get_mock_context()
 
         def flush():
             flushes.append(1)
 
-        for environ, flush_called in (
-            ({"DD_LOCAL_TEST": "True"}, True),
-            ({"DD_LOCAL_TEST": "true"}, True),
-            ({"DD_LOCAL_TEST": "1"}, True),
-            ({"DD_LOCAL_TEST": "False"}, False),
-            ({"DD_LOCAL_TEST": "false"}, False),
-            ({"DD_LOCAL_TEST": "0"}, False),
-            ({"DD_LOCAL_TEST": ""}, False),
-            ({}, False),
-        ):
-            os.environ = environ
-            flushes.clear()
-
-            @patch("datadog_lambda.wrapper.flush_extension", flush)
-            @wrapper.datadog_lambda_wrapper
-            def lambda_handler(event, context):
-                pass
+        @patch("datadog_lambda.wrapper.flush_extension", flush)
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            pass
 
-            lambda_handler(lambda_event, lambda_context)
+        lambda_handler(lambda_event, lambda_context)
 
-            self.assertEqual(flush_called, len(flushes) == 1)
+        self.assertEqual(len(flushes), 0)