Skip to content

Commit 2957287

Browse files
committed
Add data_streams_enabled.
1 parent 20d733c commit 2957287

File tree

4 files changed

+10
-28
lines changed

4 files changed

+10
-28
lines changed

datadog_lambda/config.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,9 @@ class Config:
8989
profiling_enabled = _get_env("DD_PROFILING_ENABLED", "false", as_bool)
9090
llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
9191
exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool)
92+
data_streams_enabled = _get_env(
93+
"DD_DATA_STREAMS_ENABLED", "false", as_bool, depends_on_tracing=True
94+
)
9295

9396
is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-"))
9497

datadog_lambda/wrapper.py

Lines changed: 2 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,6 @@
6060

6161
DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
6262
DD_SERVICE = "DD_SERVICE"
63-
DD_DATA_STREAMS_ENABLED = "DD_DATA_STREAMS_ENABLED"
64-
6563

6664
init_timestamp_ns = time_ns()
6765

@@ -121,26 +119,6 @@ def __init__(self, func):
121119
self.trace_extractor = None
122120
self.span = None
123121
self.inferred_span = None
124-
depends_on_dd_tracing_enabled = (
125-
lambda original_boolean: config.trace_enabled and original_boolean
126-
)
127-
self.cold_start_tracing = depends_on_dd_tracing_enabled(
128-
os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
129-
)
130-
self.data_streams_enabled = (
131-
os.environ.get(DD_DATA_STREAMS_ENABLED, "false").lower() == "true"
132-
)
133-
self.cold_start_trace_skip_lib = [
134-
"ddtrace.internal.compat",
135-
"ddtrace.filters",
136-
]
137-
if DD_COLD_START_TRACE_SKIP_LIB in os.environ:
138-
try:
139-
self.cold_start_trace_skip_lib = os.environ[
140-
DD_COLD_START_TRACE_SKIP_LIB
141-
].split(",")
142-
except Exception:
143-
logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}")
144122
self.response = None
145123

146124
if config.profiling_enabled:
@@ -261,7 +239,7 @@ def _before(self, event, context):
261239
self.inferred_span = create_inferred_span(
262240
event, context, event_source, config.decode_authorizer_context
263241
)
264-
if self.data_streams_enabled:
242+
if config.data_streams_enabled:
265243
set_dsm_context(event, event_source)
266244
self.span = create_function_execution_span(
267245
context=context,
@@ -348,7 +326,7 @@ def _after(self, event, context):
348326
LLMObs.flush()
349327

350328
# Flush exception replay
351-
if exception_replay_env_var:
329+
if config.exception_replay_enabled:
352330
LogsIntakeUploaderV1._instance.periodic()
353331

354332
if config.encode_authorizer_context and is_authorizer_response(

tests/test_config.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ def _test_as_list(env_key, conf_key, default):
8585
"DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", default=False
8686
),
8787
*_test_as_bool("DD_LOCAL_TEST", "local_test", default=False),
88+
*_test_as_bool("DD_DATA_STREAMS_ENABLED", "data_streams_enabled", default=False),
8889
*_test_int(
8990
"DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", default=10
9091
),

tests/test_wrapper.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -564,7 +564,7 @@ def return_type_test(event, context):
564564

565565
def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self):
566566
os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
567-
wrapper.dd_tracing_enabled = True
567+
os.environ["DD_TRACE_ENABLED"] = "true"
568568

569569
@wrapper.datadog_lambda_wrapper
570570
def lambda_handler(event, context):
@@ -578,7 +578,7 @@ def lambda_handler(event, context):
578578

579579
def test_set_dsm_context_not_called_when_only_DSM_enabled(self):
580580
os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
581-
wrapper.dd_tracing_enabled = False
581+
os.environ["DD_TRACE_ENABLED"] = "false"
582582

583583
@wrapper.datadog_lambda_wrapper
584584
def lambda_handler(event, context):
@@ -592,7 +592,7 @@ def lambda_handler(event, context):
592592

593593
def test_set_dsm_context_not_called_when_only_tracing_enabled(self):
594594
os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
595-
wrapper.dd_tracing_enabled = True
595+
os.environ["DD_TRACE_ENABLED"] = "true"
596596

597597
@wrapper.datadog_lambda_wrapper
598598
def lambda_handler(event, context):
@@ -606,7 +606,7 @@ def lambda_handler(event, context):
606606

607607
def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self):
608608
os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
609-
wrapper.dd_tracing_enabled = False
609+
os.environ["DD_TRACE_ENABLED"] = "false"
610610

611611
@wrapper.datadog_lambda_wrapper
612612
def lambda_handler(event, context):

0 commit comments

Comments
 (0)