feat: Enable sqs -> lambda support for DSM #604
Conversation
There was a problem hiding this comment.
@michael-zhao459 the way your code works looks just fine. But as written right now, it's pretty inefficient, recalculates work already done, and needs some refactoring. Here's what I'm thinking.
We have a feature already implemented called "inferred spans", which similarly looks at the event type and does some stuff with the event payload based on the type. I'm thinking we'll want to create a mechanism similar to this.
Take a look in datadog_lambda/tracing.py at the create_inferred_span method and where it's called in datadog_lambda/wrapper.py. I'm thinking we'll want to add a method called process_dsm (or something like that) and call it from about the same spot we call create_inferred_span.
I can walk you through all of this during our pairing time.
tests/test_wrapper.py
Outdated
| @@ -563,6 +563,204 @@ def return_type_test(event, context): | |||
| self.assertEqual(result, test_result) | |||
| self.assertFalse(MockPrintExc.called) | |||
|
|
|||
| @patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}) | |||
| def test_datadog_lambda_wrapper_dsm_sqs_context_pathway_verification(self): | |||
There was a problem hiding this comment.
I think we can vastly simplify these tests. This might warrant another pairing session, but I'll let you take a stab at it on your own first. Feel free to schedule something with me if you wanna go over it together.
There are two different files which you changed, and therefore two different test files that will need updating: tests/test_wrapper.py and a new file tests/test_dsm.py.
test_wrapper.py
For the test_wrapper.py file, we simply need to test that, based on the env vars DD_DATA_STREAMS_ENABLED and DD_TRACE_ENABLED we either do or do not call set_dsm_context with the proper args. We'll push all of the verification of what happens inside of set_dsm_context to the test_dsm.py file.
I'm a bit fan of pytest, which isn't yet imported and used in this file, which allows you to reuse the same code over and over again to create "parametrized" tests. You can accomplish this same thing using unittest (as this file already uses), though it would mean creating 4 different test methods.
# test_wrapper.py
import pytest
_test_set_dsm_context = (
("true", "true", True),
("true", "false", False),
("false", "true", False),
("false", "false", False),
)
@pytest.mark.parametrize("trace_enabled,dsm_enabled,should_call", _test_set_dsm_context)
def test_set_dsm_context(trace_enabled, dsm_enabled, should_call, monkeypatch):
# use monkeypatch to set env vars DD_TRACE_ENABLED and DD_DATA_STREAMS_ENABLED
# use monkeypatch to create a mock for `set_dsm_context`, you can also use mock.patch
@wrapper.datadog_lambda_wrapper
def lambda_handler(event, context):
return "ok"
result = lambda_handler(sqs_event, get_mock_context())
assert result == "ok"
if should_call:
# not sure of the api here, so this is just made up
assert set_dsm_context_patch.called_with == (sqs_event, EventSource(EventSourceType.SQS))
else:
# again, not sure about api
assert set_dsm_context.not_calledtest_dsm.py
In the test_dsm.py file, this is where you'll assert to make sure that the set_dsm_context works as expected. You'll want to include several tests:
- Sending an event source of anything other than
SQSwill do nothing - Sending an event with no Records will do nothing
- For each Record in the event, dsm does the setting of context as expected
There was a problem hiding this comment.
Lemme take a stab at these today, I will definitely schedule a meeting if I get stuck on any of these parts
Co-authored-by: Rey Abolofia <[email protected]>
michael-zhao459
force-pushed
the
michael.zhao/lambda-support
branch
from
81c76b6 to
dda744e
Compare
| self.mock_set_dsm_context.assert_not_called() | ||
|
|
||
| del os.environ["DD_DATA_STREAMS_ENABLED"] | ||
|
|
There was a problem hiding this comment.
These look great! Super easy to read and follow.
tests/test_dsm.py
Outdated
|
|
||
| def test_non_sqs_event_source_does_nothing(self): | ||
| """Test that non-SQS event sources don't trigger DSM context setting""" | ||
| event = {"Records": [{"body": "test"}]} |
There was a problem hiding this comment.
nit: Just to make this a tad less confusing, this event object sure looks like an sqs event to me. The test of course passes bc we look at the event source, not the event for its type. That said, can we make this event something like {}, just to make sure that it also looks not sqs?
|
|
||
| for event in events_with_no_records: | ||
| _dsm_set_sqs_context(event) | ||
| self.mock_data_streams_processor.assert_not_called() |
tests/test_dsm.py
Outdated
| } | ||
|
|
||
| mock_event_source = MagicMock() | ||
| mock_event_source.equals.return_value = True |
There was a problem hiding this comment.
I'm thinking toward the future, where you'll be adding support for sns and kinesis, etc. Once those are in place, you're not gonna want event_source.equals(whatever) to always return True.
Basically, this code works just great right now. But your future self is gonna have to change it, so why not make it work right the first time.
Instead, it's pretty easy to create a workable event source object. Instead of a mock, what about something like:
event_source = _EventSource(EventTypes.SQS)
set_dsm_context(sqs_event, event_source)Make that change to remove all the mock_event_source objects in all these tests.
|
/merge |
|
View all feedbacks in Devflow UI.
This merge request is not mergeable yet, because of pending checks/missing approvals. It will be added to the queue as soon as checks pass and/or get approvals.
[email protected] unqueued this merge request |
|
/merge -c |
|
View all feedbacks in Devflow UI.
|
dd-devflow
bot
added
mergequeue-status: removed
and removed
mergequeue-status: waiting
labels
|
Running integration tests in this PR: #608 |
|
All tests passed on the other PR. Merging. |
What does this PR do?
Allow DSM to support sqs -> lambda
Motivation
Lambda support requested by users of DSM
Testing Guidelines
Wrote a unit test inside the test_wrapper.py code that ensures context is properly being set
throughout the entire pipeline
Additional Notes
In the test I wrote a patched version of get_datastreams_context, see DataDog/dd-trace-py#13526 for the change, forced to write patched function as no new release of tracer code yet. Will remove once there is a release of tracer code.
Types of Changes
Check all that apply