feat: lambda support for DSM by michael-zhao459 · Pull Request #622 · DataDog/datadog-lambda-python

michael-zhao459 · 2025-06-20T14:18:25Z

What does this PR do?

This PR adds lambda support for Data Streams Monitoring (DSM) and reworks the original implementation.

DSM context is passed through the trace propagation headers, code is refactored to use existing extraction logic (deleted dsm.py, reinventing the wheel here).
If DSM is enabled, add custom DSM logic to the extracted context afterwards

Motivation

Remove redundant code. DSM customers wanted to have Lambda support, currently context is not propagated correctly with lambdas.

Testing Guidelines

Test case	Expected outcome	SQS	SNS	SNS -> SQS (arn from SQS is used)	Kinesis
Datadog Context propagated properly through stringValue	Data streams context propagated & data streams checkpoint set	✅	✅	✅	DNE
Datadog Context propagated properly through binaryValue	Data streams context propagated & data streams checkpoint set	✅	✅	✅	✅
No _datadog message attribute	Checkpoint set, no context propagation	✅	✅	✅	✅
Empty datadog message attribute	Checkpoint set, no context propagation	✅	✅	✅	✅
No data streams context in _datadog message attribute	Checkpoint set, no context propagation	✅	✅	✅	✅
Invalid datadog message attribute	Checkpoint set, no context propagation, debug logger called	✅	✅	✅	✅
source_arn is not found	No checkpoint set	✅	✅	✅	✅
Data streams disabled	No checkpoint set	✅	✅	✅	✅

The tests go through all of the SQS case first, then all of the SNS case, then all of the SNS->SQS case, then all of the Kinesis case

Additional Notes

Types of Changes

Bug fix
New feature
Breaking change
Misc (docs, refactoring, dependency upgrade, etc.)

Check all that apply

This PR's description is comprehensive
This PR contains breaking changes that are documented in the description
This PR introduces new APIs or parameters that are documented and unlikely to change in the foreseeable future
This PR impacts documentation, and it has been updated (or a ticket has been logged)
This PR's changes are covered by the automated tests
This PR collects user input/sensitive content into Datadog
This PR passes the integration tests (ask a Datadog member to run the tests)

michael-zhao459 · 2025-06-25T12:39:25Z

+ if config.data_streams_enabled:
+ from ddtrace.data_streams import PROPAGATION_KEY_BASE_64
+
+ data_streams_ctx = {


I know the else is redundant but datadog gets mad if i just do the if too many indents

piochelepiotr · 2025-06-25T12:42:11Z

 except Exception as e:
 logger.debug("The trace extractor returned with error %s", e)
- return extract_context_from_lambda_context(lambda_context)
+ return extract_context_from_lambda_context(lambda_context), None


we should not return None here

joeyzhao2018 · 2025-06-26T11:47:45Z


+ data_streams_ctx = {}
+ if config.data_streams_enabled:
+ from ddtrace.data_streams import PROPAGATION_KEY_BASE_64


My main concerns are
Creating dictionary objects and bound methods for every invocation is inefficient.
It is very hard to follow the logic and hard to maintain and may introduce unexpected behaviors that are hard to debug in the future.
May I suggest the following alternative implementation? Let me know what do you think.
def_create_dsm_carrier_func(dd_data): """Create a carrier function for DSM context extraction."""defcarrier_get(key): returndd_data.get(key) ifdd_dataelseNonereturncarrier_get# then in In the extraction functions:ifconfig.data_streams_enabled: dsm_carrier=_create_dsm_carrier_func(dd_data) # Pass the original dd_dataelse: dsm_carrier=None

I agree with the justifications you made for this change will change the code now!

joeyzhao2018

LGTM

joeyzhao2018

LGTM

piochelepiotr · 2025-07-08T15:39:20Z

 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=2.20.0,<4"
+ddtrace = ">=3.10.0"


this changes the major version. Is that what we want to do? Also, should we keep <4?

My mistake on the <4. The code will break without ddtrace version 3.10.0.

piochelepiotr · 2025-07-08T15:40:20Z

+ )
+
+ @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+ def test_sqs_incorrect_datadog_message_attribute(self, mock_dsm_set_checkpoint):


nit: incorrect -> invalid

piochelepiotr · 2025-07-08T15:42:31Z

+
+ @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+ @patch("datadog_lambda.tracing.logger")
+ def test_sqs_invalid_datadog_message_attribute_raises_exception(


nit: remove raises_exception from the test name. We care about what the test case tests, not the logic under the hood. Maybe in the future, the code won't raise an exception, and it's still OK because the function accepts invalid datadog message attributes.

piochelepiotr · 2025-07-08T15:44:33Z

+ event, self.lambda_context, parse_event_source(event)
+ )
+
+ mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "")


I don't understand, are we testing mock_dsm_set_checkpoint, or self.mock_checkpoint. We should test only one of them, and it should be consistent across all tests. If possible, it's better to test the lower level one (self.mock_checkpoint)

purple4reina · 2025-07-09T18:51:50Z

 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=2.20.0,<4"
+ddtrace = ">=3.10.0,<4"


We're currently working on a v3.10.2 release, once that's out, we should update this.

purple4reina · 2025-07-09T20:37:53Z

 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=2.20.0,<4"
+ddtrace = ">=3.10.2,<4"


Should be good to release this now that v3.10.2 is up in pypi.

datadog-datadog-prod-us1Bot reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

michael-zhao459 requested a review from piochelepiotr June 20, 2025 19:38

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threadtests/test_dsm.py Outdated

purple4reina reviewed Jun 20, 2025
View reviewed changes

Comment threadtests/test_dsm.py Outdated

piochelepiotr reviewed Jun 23, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

datadog-datadog-prod-us1Bot reviewed Jun 23, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

datadog-datadog-prod-us1Bot reviewed Jun 23, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

piochelepiotr reviewed Jun 24, 2025
View reviewed changes

Comment threaddatadog_lambda/wrapper.py Outdated
Comment threaddatadog_lambda/tracing.py Outdated
Comment threaddatadog_lambda/tracing.py Outdated

michael-zhao459 requested a review from piochelepiotr June 24, 2025 14:42

piochelepiotr reviewed Jun 24, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated
Comment threaddatadog_lambda/tracing.py Outdated
Comment threaddatadog_lambda/tracing.py Outdated

datadog-datadog-prod-us1Bot reviewed Jun 25, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated
Comment threaddatadog_lambda/tracing.py Outdated

michael-zhao459 commented Jun 25, 2025
View reviewed changes

piochelepiotr approved these changes Jun 25, 2025
View reviewed changes

michael-zhao459 marked this pull request as ready for review June 25, 2025 13:25

michael-zhao459 requested review from a team as code owners June 25, 2025 13:25

joeyzhao2018 reviewed Jun 26, 2025
View reviewed changes

michael-zhao459 requested a review from joeyzhao2018 June 26, 2025 12:34

joeyzhao2018 reviewed Jun 26, 2025
View reviewed changes

Comment threaddatadog_lambda/tracing.py Outdated

joeyzhao2018 self-requested a review June 26, 2025 14:02

joeyzhao2018 approved these changes Jun 26, 2025
View reviewed changes

purple4reina reviewed Jun 30, 2025
View reviewed changes

Comment threaddatadog_lambda/wrapper.py Outdated

remove not needed test
0f71f84

michael-zhao459 requested a review from purple4reina July 3, 2025 17:58

formatting fix
c4fa49b

michael-zhao459 force-pushed the michael.zhao/dsm-lambda branch from 66fdbb3 to c4fa49bCompare July 3, 2025 19:43

michael-zhao459 added 14 commits July 7, 2025 08:11

made tests clearer and more comprehensive
e467dd2

Merge branch 'main' into michael.zhao/dsm-lambda
018994b

fix
0665401

made comments more obvious on point of tests
647c447

update toml file
e390af5

edit build make more sense
5ffb0e7

made sure tests match with table
520a54f

fix
24821b8

removed not needed comments
5f838e2

addc comment
6625dfc

add some more kinesis tests
3bbdb88

add logger check in exception test cases
f0a7d88

fix comment
828a050

fix
48e07d1

piochelepiotr reviewed Jul 8, 2025
View reviewed changes

michael-zhao459 added 2 commits July 8, 2025 14:19

updated tests to use inner checkpoint check, keep build at <4
ae8ee9d

fix
19765e2

piochelepiotr approved these changes Jul 8, 2025
View reviewed changes

purple4reina reviewed Jul 9, 2025
View reviewed changes

change to 3.10.2
1cd76e4

michael-zhao459 requested a review from purple4reina July 9, 2025 20:12

purple4reina approved these changes Jul 9, 2025
View reviewed changes

purple4reina reviewed Jul 9, 2025
View reviewed changes

michael-zhao459 merged commit d2a195e into mainJul 9, 2025
61 checks passed

michael-zhao459 deleted the michael.zhao/dsm-lambda branch July 9, 2025 21:31

michael-zhao459 mentioned this pull request Jul 30, 2025
fix(dsm): set dsm checkpoint for all records in event #643
Draft
11 tasks

Uh oh!

Conversation

michael-zhao459 commented Jun 20, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Motivation

Testing Guidelines

Additional Notes

Types of Changes

Check all that apply

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

joeyzhao2018Jun 26, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

joeyzhao2018 left a comment

Choose a reason for hiding this comment

Uh oh!

joeyzhao2018 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

michael-zhao459 commented Jun 20, 2025•
edited
Loading

joeyzhao2018Jun 26, 2025•
edited
Loading