Skip to content

Commit 9a14033

Browse files
michael-zhao459michael-zhao459
committed
removed .get, return only dd
_data, create carreir get function later
1 parent c505705 commit 9a14033

File tree

4 files changed

+56
-55
lines changed

4 files changed

+56
-55
lines changed

‎datadog_lambda/tracing.py‎

Lines changed: 5 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -265,16 +265,6 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
265265
if dd_json_data:
266266
dd_data = json.loads(dd_json_data)
267267

268-
data_streams_ctx = {}
269-
if config.data_streams_enabled:
270-
from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
271-
272-
data_streams_ctx = {
273-
PROPAGATION_KEY_BASE_64: dd_data[PROPAGATION_KEY_BASE_64]
274-
if PROPAGATION_KEY_BASE_64 in dd_data
275-
else {}
276-
}
277-
278268
if is_step_function_event(dd_data):
279269
try:
280270
return (
@@ -285,7 +275,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
285275
logger.debug(
286276
"Failed to extract Step Functions context from SQS/SNS event."
287277
)
288-
return propagator.extract(dd_data), data_streams_ctx.get
278+
return propagator.extract(dd_data), dd_data
289279
else:
290280
# Handle case where trace context is injected into attributes.AWSTraceHeader
291281
# example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -311,10 +301,10 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
311301
),
312302
None,
313303
)
314-
return extract_context_from_lambda_context(lambda_context), {}.get
304+
return extract_context_from_lambda_context(lambda_context), {}
315305
except Exception as e:
316306
logger.debug("The trace extractor returned with error %s", e)
317-
return extract_context_from_lambda_context(lambda_context), {}.get
307+
return extract_context_from_lambda_context(lambda_context), {}
318308

319309

320310
def _extract_context_from_eventbridge_sqs_event(event):
@@ -373,7 +363,6 @@ def extract_context_from_kinesis_event(event, lambda_context):
373363
"""
374364
Extract datadog trace context from a Kinesis Stream's base64 encoded data string
375365
"""
376-
data_streams_ctx = {}
377366
try:
378367
record = get_first_record(event)
379368
kinesis = record.get("kinesis")
@@ -389,20 +378,11 @@ def extract_context_from_kinesis_event(event, lambda_context):
389378
data_obj = json.loads(data_str)
390379
dd_ctx = data_obj.get("_datadog")
391380
if dd_ctx:
392-
if config.data_streams_enabled:
393-
from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
394-
395-
data_streams_ctx = {
396-
PROPAGATION_KEY_BASE_64: dd_ctx[PROPAGATION_KEY_BASE_64]
397-
if PROPAGATION_KEY_BASE_64 in dd_ctx
398-
else {}
399-
}
400-
401-
return propagator.extract(dd_ctx), data_streams_ctx.get
381+
return propagator.extract(dd_ctx), dd_ctx
402382
except Exception as e:
403383
logger.debug("The trace extractor returned with error %s", e)
404384

405-
return extract_context_from_lambda_context(lambda_context), data_streams_ctx.get
385+
return extract_context_from_lambda_context(lambda_context), {}
406386

407387

408388
def _deterministic_sha256_hash(s: str, part: str) -> int:

‎datadog_lambda/wrapper.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111

1212

1313
from datadog_lambda.asm import asm_start_response, asm_start_request
14-
from datadog_lambda.dsm import set_dsm_context
1514

1615
from datadog_lambda.extension import should_use_extension, flush_extension
1716
from datadog_lambda.cold_start import (
@@ -224,7 +223,7 @@ def _before(self, event, context):
224223
dd_context,
225224
trace_context_source,
226225
event_source,
227-
dsm_carrier,
226+
dd_data,
228227
) = extract_dd_trace_context(
229228
event,
230229
context,
@@ -256,6 +255,7 @@ def _before(self, event, context):
256255
or event_source.equals(EventTypes.KINESIS)
257256
):
258257
source_arn = extract_source_arn(event)
258+
dsm_carrier = _create_dsm_carrier_func(dd_data)
259259
set_dsm_checkpoint(
260260
dsm_carrier, event_source.to_string(), source_arn
261261
)
@@ -392,4 +392,13 @@ def extract_source_arn(event):
392392
return arn
393393

394394

395+
def _create_dsm_carrier_func(dd_data):
396+
"""Create a carrier function for DSM context extraction."""
397+
398+
def carrier_get(key):
399+
return dd_data.get(key) if dd_data else None
400+
401+
return carrier_get
402+
403+
395404
datadog_lambda_wrapper = _LambdaDecorator

‎tests/test_tracing.py‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2475,14 +2475,14 @@ def test_extract_context_from_sqs_event_with_datadog_context(self):
24752475
]
24762476
}
24772477

2478-
(context, dsm_carrier) = extract_context_from_sqs_or_sns_event_or_context(
2478+
(context, dd_data) = extract_context_from_sqs_or_sns_event_or_context(
24792479
sqs_event, lambda_ctx
24802480
)
24812481

24822482
self.assertEqual(context.trace_id, 456)
24832483
self.assertEqual(context.span_id, 789)
24842484
self.assertEqual(context.sampling_priority, 1)
2485-
self.assertEqual(dsm_carrier(PROPAGATION_KEY_BASE_64), "test-data")
2485+
self.assertEqual(dd_data, header)
24862486

24872487
@patch("datadog_lambda.config.Config.trace_enabled", True)
24882488
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
@@ -2511,14 +2511,14 @@ def test_extract_context_from_sns_event_with_datadog_context(self):
25112511
]
25122512
}
25132513

2514-
(context, dsm_carrier) = extract_context_from_sqs_or_sns_event_or_context(
2514+
(context, dd_data) = extract_context_from_sqs_or_sns_event_or_context(
25152515
sns_event, lambda_ctx
25162516
)
25172517

25182518
self.assertEqual(context.trace_id, 111)
25192519
self.assertEqual(context.span_id, 222)
25202520
self.assertEqual(context.sampling_priority, 2)
2521-
self.assertEqual(dsm_carrier(PROPAGATION_KEY_BASE_64), "test-data")
2521+
self.assertEqual(dd_data, expected_trace_data)
25222522

25232523
@patch("datadog_lambda.config.Config.trace_enabled", True)
25242524
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
@@ -2543,13 +2543,11 @@ def test_extract_context_from_kinesis_event_with_datadog_context(self):
25432543
]
25442544
}
25452545

2546-
context, dsm_carrier = extract_context_from_kinesis_event(
2547-
kinesis_event, lambda_ctx
2548-
)
2546+
context, dd_data = extract_context_from_kinesis_event(kinesis_event, lambda_ctx)
25492547

25502548
self.assertEqual(context.trace_id, 333)
25512549
self.assertEqual(context.span_id, 444)
2552-
self.assertEqual(dsm_carrier(PROPAGATION_KEY_BASE_64), "test-data")
2550+
self.assertEqual(dd_data, expected_trace_data)
25532551

25542552
@patch("datadog_lambda.config.Config.trace_enabled", True)
25552553
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
@@ -2587,11 +2585,11 @@ def test_extract_context_from_sns_to_sqs_event_with_datadog_context(self):
25872585
]
25882586
}
25892587

2590-
(context, dsm_carrier) = extract_context_from_sqs_or_sns_event_or_context(
2588+
(context, dd_data) = extract_context_from_sqs_or_sns_event_or_context(
25912589
sns_to_sqs_event, lambda_ctx
25922590
)
25932591

25942592
self.assertEqual(context.trace_id, 555)
25952593
self.assertEqual(context.span_id, 666)
25962594
self.assertEqual(context.sampling_priority, 1)
2597-
self.assertEqual(dsm_carrier(PROPAGATION_KEY_BASE_64), "test-data")
2595+
self.assertEqual(dd_data, expected_trace_data)

‎tests/test_wrapper.py‎

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
from datadog_lambda.thread_stats_writer import ThreadStatsWriter
1515
from datadog_lambda.trigger import EventTypes
1616
from ddtrace.trace import Span, tracer
17-
from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
1817
from ddtrace.internal.constants import MAX_UINT_64BITS
1918

2019
from tests.utils import get_mock_context, reset_xray_connection
@@ -566,12 +565,15 @@ def return_type_test(event, context):
566565

567566
@patch("datadog_lambda.config.Config.trace_enabled", True)
568567
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
569-
def test_set_dsm_checkpoint_called_when_DSM_and_tracing_enabled(self):
568+
@patch("datadog_lambda.wrapper._create_dsm_carrier_func")
569+
def test_set_dsm_checkpoint_called_when_DSM_and_tracing_enabled(
570+
self, mock_create_dsm_carrier_func
571+
):
570572
event_source = Mock()
571573
event_source.to_string.return_value = "sqs"
572574
event_source.equals.return_value = True
573575

574-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
576+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
575577
arn = "test-arn"
576578
self.mock_extract_dd_trace_context.return_value = (
577579
{},
@@ -589,15 +591,15 @@ def lambda_handler(event, context):
589591
result = lambda_handler({}, get_mock_context())
590592
self.assertEqual(result, "ok")
591593
self.mock_set_dsm_checkpoint.assert_called_once_with(
592-
data_streams_ctx, "sqs", arn
594+
mock_create_dsm_carrier_func.return_value, "sqs", arn
593595
)
594596

595597
@patch("datadog_lambda.config.Config.trace_enabled", False)
596598
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
597599
def test_set_dsm_checkpoint_not_called_when_only_DSM_enabled(self):
598600
event_source = Mock()
599601
event_source.to_string.return_value = "sqs"
600-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
602+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
601603
self.mock_extract_dd_trace_context.return_value = (
602604
{},
603605
None,
@@ -618,7 +620,7 @@ def lambda_handler(event, context):
618620
def test_set_dsm_checkpoint_not_called_when_only_tracing_enabled(self):
619621
event_source = Mock()
620622
event_source.to_string.return_value = "sqs"
621-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
623+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
622624
self.mock_extract_dd_trace_context.return_value = (
623625
{},
624626
None,
@@ -639,7 +641,7 @@ def lambda_handler(event, context):
639641
def test_set_dsm_checkpoint_not_called_when_tracing_and_DSM_disabled(self):
640642
event_source = Mock()
641643
event_source.to_string.return_value = "sqs"
642-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
644+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
643645
self.mock_extract_dd_trace_context.return_value = (
644646
{},
645647
None,
@@ -655,16 +657,19 @@ def lambda_handler(event, context):
655657
self.assertEqual(result, "ok")
656658
self.mock_set_dsm_checkpoint.assert_not_called()
657659

660+
@patch("datadog_lambda.wrapper._create_dsm_carrier_func")
658661
@patch("datadog_lambda.config.Config.trace_enabled", True)
659662
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
660-
def test_set_dsm_checkpoint_called_for_sqs_event(self):
663+
def test_set_dsm_checkpoint_called_for_sqs_event(
664+
self, mock_create_dsm_carrier_func
665+
):
661666
event_source = Mock()
662667
event_source.to_string.return_value = "sqs"
663668
event_source.equals.side_effect = (
664669
lambda event_type: event_type == EventTypes.SQS
665670
)
666671

667-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
672+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
668673
arn = "test-arn"
669674
self.mock_extract_dd_trace_context.return_value = (
670675
{},
@@ -682,19 +687,22 @@ def lambda_handler(event, context):
682687
result = lambda_handler({}, get_mock_context())
683688
self.assertEqual(result, "ok")
684689
self.mock_set_dsm_checkpoint.assert_called_once_with(
685-
data_streams_ctx, "sqs", arn
690+
mock_create_dsm_carrier_func.return_value, "sqs", arn
686691
)
687692

688693
@patch("datadog_lambda.config.Config.trace_enabled", True)
689694
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
690-
def test_set_dsm_checkpoint_called_for_sns_event(self):
695+
@patch("datadog_lambda.wrapper._create_dsm_carrier_func")
696+
def test_set_dsm_checkpoint_called_for_sns_event(
697+
self, mock_create_dsm_carrier_func
698+
):
691699
event_source = Mock()
692700
event_source.to_string.return_value = "sns"
693701
event_source.equals.side_effect = (
694702
lambda event_type: event_type == EventTypes.SNS
695703
)
696704

697-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
705+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
698706
arn = "test-arn"
699707
self.mock_extract_dd_trace_context.return_value = (
700708
{},
@@ -712,19 +720,22 @@ def lambda_handler(event, context):
712720
result = lambda_handler({}, get_mock_context())
713721
self.assertEqual(result, "ok")
714722
self.mock_set_dsm_checkpoint.assert_called_once_with(
715-
data_streams_ctx, "sns", arn
723+
mock_create_dsm_carrier_func.return_value, "sns", arn
716724
)
717725

718726
@patch("datadog_lambda.config.Config.trace_enabled", True)
719727
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
720-
def test_set_dsm_checkpoint_called_for_kinesis_event(self):
728+
@patch("datadog_lambda.wrapper._create_dsm_carrier_func")
729+
def test_set_dsm_checkpoint_called_for_kinesis_event(
730+
self, mock_create_dsm_carrier_func
731+
):
721732
event_source = Mock()
722733
event_source.to_string.return_value = "kinesis"
723734
event_source.equals.side_effect = (
724735
lambda event_type: event_type == EventTypes.KINESIS
725736
)
726737

727-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
738+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}.get
728739
arn = "test-arn"
729740
self.mock_extract_dd_trace_context.return_value = (
730741
{},
@@ -742,19 +753,22 @@ def lambda_handler(event, context):
742753
result = lambda_handler({}, get_mock_context())
743754
self.assertEqual(result, "ok")
744755
self.mock_set_dsm_checkpoint.assert_called_once_with(
745-
data_streams_ctx, "kinesis", arn
756+
mock_create_dsm_carrier_func.return_value, "kinesis", arn
746757
)
747758

748759
@patch("datadog_lambda.config.Config.trace_enabled", True)
749760
@patch("datadog_lambda.config.Config.data_streams_enabled", True)
750-
def test_set_dsm_checkpoint_not_called_for_unknown_event(self):
761+
@patch("datadog_lambda.wrapper._create_dsm_carrier_func")
762+
def test_set_dsm_checkpoint_not_called_for_unknown_event(
763+
self, mock_create_dsm_carrier_func
764+
):
751765
event_source = Mock()
752766
event_source.to_string.return_value = "kinesis"
753767
event_source.equals.side_effect = (
754768
lambda event_type: event_type == EventTypes.UNKNOWN
755769
)
756770

757-
data_streams_ctx = {PROPAGATION_KEY_BASE_64: "test-data"}.get
771+
data_streams_ctx = {"dd-pathway-ctx-base64": "test-data"}
758772
self.mock_extract_dd_trace_context.return_value = (
759773
{},
760774
None,

0 commit comments

Comments
 (0)