Add data_streams_enabled.

Author: purple4reina

datadog_lambda/config.py

@@ -89,6 +89,9 @@ class Config:
     profiling_enabled = _get_env("DD_PROFILING_ENABLED", "false", as_bool)
     llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
     exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool)
+    data_streams_enabled = _get_env(
+        "DD_DATA_STREAMS_ENABLED", "false", as_bool, depends_on_tracing=True
+    )
 
     is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-"))
 

datadog_lambda/wrapper.py

@@ -60,8 +60,6 @@
 
 DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
 DD_SERVICE = "DD_SERVICE"
-DD_DATA_STREAMS_ENABLED = "DD_DATA_STREAMS_ENABLED"
-
 
 init_timestamp_ns = time_ns()
 
@@ -121,26 +119,6 @@ def __init__(self, func):
             self.trace_extractor = None
             self.span = None
             self.inferred_span = None
-            depends_on_dd_tracing_enabled = (
-                lambda original_boolean: config.trace_enabled and original_boolean
-            )
-            self.cold_start_tracing = depends_on_dd_tracing_enabled(
-                os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
-            )
-            self.data_streams_enabled = (
-                os.environ.get(DD_DATA_STREAMS_ENABLED, "false").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:
@@ -261,7 +239,7 @@ def _before(self, event, context):
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, config.decode_authorizer_context
                     )
-                if self.data_streams_enabled:
+                if config.data_streams_enabled:
                     set_dsm_context(event, event_source)
                 self.span = create_function_execution_span(
                     context=context,
@@ -348,7 +326,7 @@ def _after(self, event, context):
                 LLMObs.flush()
 
             # Flush exception replay
-            if exception_replay_env_var:
+            if config.exception_replay_enabled:
                 LogsIntakeUploaderV1._instance.periodic()
 
             if config.encode_authorizer_context and is_authorizer_response(

tests/test_config.py

@@ -85,6 +85,7 @@ def _test_as_list(env_key, conf_key, default):
         "DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", default=False
     ),
     *_test_as_bool("DD_LOCAL_TEST", "local_test", default=False),
+    *_test_as_bool("DD_DATA_STREAMS_ENABLED", "data_streams_enabled", default=False),
     *_test_int(
         "DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", default=10
     ),

tests/test_wrapper.py

@@ -564,7 +564,7 @@ def return_type_test(event, context):
 
     def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self):
         os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
-        wrapper.dd_tracing_enabled = True
+        os.environ["DD_TRACE_ENABLED"] = "true"
 
         @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -578,7 +578,7 @@ def lambda_handler(event, context):
 
     def test_set_dsm_context_not_called_when_only_DSM_enabled(self):
         os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
-        wrapper.dd_tracing_enabled = False
+        os.environ["DD_TRACE_ENABLED"] = "false"
 
         @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -592,7 +592,7 @@ def lambda_handler(event, context):
 
     def test_set_dsm_context_not_called_when_only_tracing_enabled(self):
         os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
-        wrapper.dd_tracing_enabled = True
+        os.environ["DD_TRACE_ENABLED"] = "true"
 
         @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -606,7 +606,7 @@ def lambda_handler(event, context):
 
     def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self):
         os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
-        wrapper.dd_tracing_enabled = False
+        os.environ["DD_TRACE_ENABLED"] = "false"
 
         @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):