use ddtracepy extractor #379

zARODz11z · 2023-10-16T21:54:50Z

What does this PR do?

Motivation

Testing Guidelines

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)

zARODz11z · 2023-10-16T21:55:02Z

datadog_lambda/constants.py

@@ -19,6 +19,11 @@ class TraceHeader(object):
    PARENT_ID = "x-datadog-parent-id"
    SAMPLING_PRIORITY = "x-datadog-sampling-priority"

+class TraceState(object):


purple4reina · 2023-10-16T22:42:31Z

datadog_lambda/tracing.py

-        ) = _extract_context_from_eventbridge_sqs_event(event)
-        return trace_id, parent_id, sampling_priority
+        context = _extract_context_from_eventbridge_sqs_event(event)
+        return context


We should probably check if None here.

datadog_lambda/tracing.py

purple4reina · 2023-11-06T18:43:12Z

datadog_lambda/tracing.py

@@ -175,8 +174,9 @@ def extract_context_from_lambda_context(lambda_context):
            trace_id = client_context.custom.get(TraceHeader.TRACE_ID)
            parent_id = client_context.custom.get(TraceHeader.PARENT_ID)
            sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY)


Consider what if the client_context.custom has w3c headers, not just dd headers.

datadog_lambda/tracing.py

purple4reina · 2023-11-06T19:01:41Z

datadog_lambda/tracing.py

-    span = tracer.trace("dummy.span")
-    span.trace_id = int(context[TraceHeader.TRACE_ID])
-    span.span_id = int(context[TraceHeader.PARENT_ID])
+    if isinstance(context, Context) and context is not None:


Let's try to make sure that context is always a Context object.

purple4reina · 2023-11-07T20:29:53Z

datadog_lambda/tracing.py

@@ -175,8 +174,9 @@ def extract_context_from_lambda_context(lambda_context):
            trace_id = client_context.custom.get(TraceHeader.TRACE_ID)
            parent_id = client_context.custom.get(TraceHeader.PARENT_ID)
            sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY)
-
-    return trace_id, parent_id, sampling_priority
+            context = Context(trace_id=trace_id, span_id=parent_id,sampling_priority=sampling_priority)


nit

Suggested change

context = Context(trace_id=trace_id, span_id=parent_id,sampling_priority=sampling_priority)

context = Context(trace_id=trace_id, span_id=parent_id, sampling_priority=sampling_priority)

purple4reina · 2023-11-07T20:32:26Z

datadog_lambda/tracing.py

-        return trace_id, parent_id, sampling_priority
+        context = propagator.extract(dd_data)
+        if context is not None:
+            return context


The if not None here seems unnecessary

purple4reina · 2023-11-07T20:33:43Z

datadog_lambda/tracing.py

@@ -328,10 +317,9 @@ def _extract_context_from_eventbridge_sqs_event(event):



Change

body_str = first_record.get("body", {})

to

body_str = first_record.get("body")

We already know the "body" key is in the dictionary. Plus, we're then calling json.loads on the returned value. That will fail if given a non-string like {}.

purple4reina · 2023-11-07T20:36:32Z

datadog_lambda/tracing.py

-        return trace_id, parent_id, sampling_priority
+        context = propagator.extract(dd_context)
+        if context is not None:
+            return context


The if is not None here is unnecessary

purple4reina · 2023-11-07T20:37:12Z

datadog_lambda/tracing.py

-        return trace_id, parent_id, sampling_priority
+        context = propagator.extract(dd_ctx)
+        if context is not None:
+            return context


The if is not None here is unnecessary

purple4reina · 2023-11-07T20:38:46Z

datadog_lambda/tracing.py

@@ -408,7 +393,8 @@ def extract_context_from_step_functions(event, lambda_context):
            execution_id + "#" + state_name + "#" + state_entered_time
        )
        sampling_priority = SamplingPriority.AUTO_KEEP
-        return trace_id, parent_id, sampling_priority
+        context = Context(trace_id=trace_id, parent_id= parent_id, sampling_priority=sampling_priority)
+        return context


or just

return Context(trace_id=trace_id, parent_id= parent_id, sampling_priority=sampling_priority)

purple4reina · 2023-11-07T20:48:30Z

datadog_lambda/tracing.py

@@ -660,18 +623,32 @@ def is_lambda_context():

 def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
    if trace_context_source == TraceContextSource.EVENT or merge_xray_traces:
-        context = dict(dd_trace_context)
+        global dd_trace_context


Changing this to a global alters how this function works. If merge xray traces is True, then you're mutating the global trace context object, which was not done before.

purple4reina · 2023-11-07T20:49:21Z

datadog_lambda/tracing.py

+                dd_trace_context.span_id = xray_context.span_id
+                dd_trace_context.trace_id = xray_context.trace_id
+                dd_trace_context.sampling_priority = xray_context.sampling_priority
+                dd_trace_context.dd_origin = xray_context.dd_origin


Previously we were only setting parent-id from xray context. We shouldn't change that functionality.

purple4reina · 2023-11-07T21:25:06Z

datadog_lambda/xray.py

+    #     dd_origin=xray_trace_entity.get("source")
+    # )
+    context = Context(trace_id=parent, span_id=parent, sampling_priority=sampled, dd_origin=TraceContextSource.XRAY)
+    return context


just

return Context(trace_id=parent, span_id=parent, sampling_priority=sampled, dd_origin=TraceContextSource.XRAY)

purple4reina · 2023-11-07T21:25:44Z

datadog_lambda/xray.py



 def generate_random_id():
    return binascii.b2a_hex(os.urandom(8)).decode("utf-8")


-def build_segment(context, key, metadata):
+def build_segment(context: Context, key, metadata):


I think we should all or nothing the python type checks.

purple4reina · 2023-11-07T21:26:27Z

tests/event_samples/lambda-url-w3c.json

+        "accept": "*/*",
+        "accept-encoding": "gzip, deflate",
+        "host": "szcqwshtpwosbcodv4hge5zj4a0bmker.lambda-url.sa-east-1.on.aws",
+        "traceparent": "00-00000000000000005f486460f346b80c-dc070c2bb8714351-01",


Suggested change

"traceparent": "00-00000000000000005f486460f346b80c-dc070c2bb8714351-01",

"traceparent": "01-00000000000000005f486460f346b80c-dc070c2bb8714351-01",

purple4reina · 2023-11-07T21:27:27Z

tests/test_tracing.py

@@ -42,7 +44,7 @@
    "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
 )
 fake_xray_header_value_parent_decimal = "10713633173203262661"
-fake_xray_header_value_root_decimal = "3995693151288333088"
+fake_xray_header_value_root_decimal = "1490261136348486853"


We shouldn't need to change this.

purple4reina · 2023-11-07T21:28:08Z

tests/test_tracing.py

            ctx,
-            {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"},
+            Context(trace_id=123, span_id=321, sampling_priority=1)


Probably best to do as you did previously, an assert each for trace id, span id, and sampling priority.

purple4reina · 2023-11-07T21:30:22Z

tests/test_tracing.py

+        self.assertEqual(
+            ctx.sampling_priority,
+            '2'
+        ) 


Go ahead and put this all on one line. We dont need this file longer than it already is.

purple4reina · 2023-11-07T21:30:59Z

tests/test_tracing.py

        self.mock_send_segment.assert_called_with(
            XraySubsegment.TRACE_KEY,
-            {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"},
+            actual_context


This is wrong. You should be doing assertions against the args returned from .call_args.

purple4reina · 2023-11-07T21:31:27Z

tests/test_tracing.py

-        )
+        self.assertEqual(str(ctx.trace_id), "666")
+        self.assertEqual(str(ctx.span_id), "777")
+        self.assertEqual(str(ctx.sampling_priority), "1")


Instead of casting to a string, just assert the values are ints.

purple4reina · 2023-11-07T21:31:42Z

tests/test_tracing.py

        self.mock_send_segment.assert_called_with(
            XraySubsegment.TRACE_KEY,
-            {"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
+            actual_context


This is wrong, see above.

purple4reina · 2023-11-07T21:32:05Z

tests/test_tracing.py

        self.mock_send_segment.assert_called_with(
            XraySubsegment.TRACE_KEY,
-            {"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
+            actual_context


This is wrong, see above.

purple4reina · 2023-11-07T21:33:07Z

tests/test_tracing.py

+        propagator = HTTPPropagator()
+        context = propagator.extract(headers)
+        print(context)
+        self.assertEqual(1,2)


This block should probably be removed.

purple4reina · 2023-11-07T21:35:25Z

tests/test_tracing.py

+        mock_span.trace_id = 123
+        mock_span.span_id = 456
+        mock_trace.return_value = mock_span
+        mock_current_span.return_value = mock_span


Using mocks is confusing and unnecessary. If you're worried about the random number generator creating different ids each time, try seeding the generator. Something like https://docs.python.org/3/library/random.html#random.seed

purple4reina · 2023-11-07T21:36:46Z

tests/test_tracing.py

-        self.assertEqual(context["sampling-priority"], "1")
+        self.assertEqual(context.trace_id, 7379586022458917877)
+        self.assertEqual(context.span_id, 2644033662113726488)
+        self.assertEqual(context.sampling_priority, 1)


purple4reina · 2023-11-07T21:39:26Z

tests/test_tracing.py

+        )
+        self.mock_activate.assert_called()
+        self.mock_activate.assert_has_calls([call(expected_context)])
+    def test_use_dd_trace_context(self):


This test doesn't seem to add anything we didn't already have.

purple4reina · 2023-11-07T21:40:36Z

tests/test_tracing.py

+        ctx = get_mock_context()
+        context, source, event_type = extract_dd_trace_context(event, ctx)
+        self.assertEqual(context._traceparent, "00-00000000000000005f486460f346b80c-dc070c2bb8714351-01")
+        self.assertEqual(context._tracestate, "dd=s:1;t.dm:-1")


Let's actually test the trace id, span id, and sampling priority on the context directly. Attributes starting with an _ in Python are implied to be private and could change at anytime.

purple4reina · 2023-11-07T21:41:09Z

tests/test_tracing.py

We will also need tests that check for w3c trace headers for all of the extract methods tested in this file.

purple4reina · 2023-11-07T21:42:52Z

datadog_lambda/tracing.py

-            parent_id = dd_data.get(TraceHeader.PARENT_ID)
-            sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
+            dd_data = client_context.custom.get("_datadog")
+            context = propagator.extract(dd_data)
        elif (
            TraceHeader.TRACE_ID in client_context.custom
            and TraceHeader.PARENT_ID in client_context.custom
            and TraceHeader.SAMPLING_PRIORITY in client_context.custom


We can't check for these headers exclusively. This would not support w3c headers in the client_context.custom.

purple4reina · 2023-11-07T21:43:46Z

datadog_lambda/tracing.py

-
-    return trace_id, parent_id, sampling_priority
+            context = propagator.extract(client_context.custom)
+        return context


This is a bug. If neither the if or the elif are hit, then the variable context is undefined.

use ddtracepy extractor

4ab9e36

zARODz11z requested a review from a team as a code owner October 16, 2023 21:54

zARODz11z commented Oct 16, 2023

View reviewed changes

purple4reina reviewed Oct 16, 2023

View reviewed changes

Andrew Rodriguez and others added 5 commits October 17, 2023 10:25

add None checks

4803490

fixed all tests

d6760da

clean up, just have a few TODOs

830e5cb

add lambda url w3c event sample and test

7416436

refactor some more, need to fix tests

072f95e