bugfix: set_trace_entity() in lambda adds segment to thread local (#409)

Author: srprash

aws_xray_sdk/core/lambda_launcher.py

@@ -78,6 +78,18 @@ def put_subsegment(self, subsegment):
         current_entity.add_subsegment(subsegment)
         self._local.entities.append(subsegment)
 
+    def set_trace_entity(self, trace_entity):
+        """
+        For Lambda context, we additionally store the segment in the thread local.
+        """
+        if self._is_subsegment(trace_entity):
+            segment = trace_entity.parent_segment
+        else:
+            segment = trace_entity
+
+        setattr(self._local, 'segment', segment)
+        setattr(self._local, 'entities', [trace_entity])
+
     def get_trace_entity(self):
         self._refresh_context()
         if getattr(self._local, 'entities', None):

tests/test_lambda_context.py

@@ -83,3 +83,24 @@ def test_non_initialized():
     temp_context.put_subsegment(subsegment)
 
     assert temp_context.get_trace_entity() == subsegment
+
+
+def test_set_trace_entity():
+    segment = context.get_trace_entity()
+    subsegment = Subsegment('name', 'local', segment)
+
+    context. clear_trace_entities()
+
+    # should set the parent segment in thread local
+    context.set_trace_entity(subsegment)
+    tl = context._local
+    assert tl.__getattribute__('segment') == segment
+    assert context.get_trace_entity() == subsegment
+
+    context.clear_trace_entities()
+
+    # should set the segment in thread local
+    context.set_trace_entity(segment)
+    tl = context._local
+    assert tl.__getattribute__('segment') == segment
+    assert context.get_trace_entity() == segment