Use xray context span id when it exists (#401)

* Extract xray context span id when it exists

* Add test for no Parent in _X_AMZN_TRACE_ID.

---------

Co-authored-by: Daniel Ebrahimian <[email protected]>

Author: purple4reina

datadog_lambda/tracing.py

@@ -604,7 +604,7 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
         )
         if merge_xray_traces:
             xray_context = _get_xray_trace_context()
-            if xray_context.span_id:
+            if xray_context and xray_context.span_id:
                 context.span_id = xray_context.span_id
 
         tracer.context_provider.activate(context)

tests/test_tracing.py

@@ -14,6 +14,7 @@
 from datadog_lambda.constants import (
     SamplingPriority,
     TraceHeader,
+    TraceContextSource,
     XraySubsegment,
 )
 from datadog_lambda.tracing import (
@@ -854,6 +855,30 @@ def test_mixed_parent_context_when_merging(self):
         self.mock_activate.assert_called()
         self.mock_activate.assert_has_calls([call(expected_context)])
 
+    def test_set_dd_trace_py_root_no_span_id(self):
+        os.environ["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320"
+
+        lambda_ctx = get_mock_context()
+        ctx, source, event_type = extract_dd_trace_context(
+            {
+                "headers": {
+                    TraceHeader.TRACE_ID: "123",
+                    TraceHeader.PARENT_ID: "321",
+                    TraceHeader.SAMPLING_PRIORITY: "1",
+                }
+            },
+            lambda_ctx,
+        )
+        set_dd_trace_py_root(TraceContextSource.EVENT, True)
+
+        expected_context = Context(
+            trace_id=123,  # Trace Id from incomming context
+            span_id=321,  # Span Id from incoming context
+            sampling_priority=1,  # Sampling priority from incomming context
+        )
+        self.mock_activate.assert_called()
+        self.mock_activate.assert_has_calls([call(expected_context)])
+
 
 class TestAuthorizerInferredSpans(unittest.TestCase):
     def setUp(self):