Segment / Subsegment / Recorder.capture context managers (#97)

* Add initial code for {sub}segment context mangers

Add SubsegmentContextManager and SubsegmentContextManager classes that
implements the context manager protocol and on entering context returns
{sub}segment object to add metadata etc

Add in_{sub}segment function on recorder to manke use of the context
mangers

* Test basic context managers functionality

* Allow `recorder.capture` to be used as context manager

* Add async context managers support

* Add context managers usage to readme

README.md

@@ -66,6 +66,40 @@ xray_recorder.configure(
 
 ### Start a custom segment/subsegment
 
+Using context managers for implicit exceptions recording:
+
+```python
+from aws_xray_sdk.core import xray_recorder
+
+with xray_recorder.in_segment('segment_name') as segment:
+    # Add metadata or annotation here if necessary
+    segment.put_metadata('key', dict, 'namespace')
+    with xray_recorder.in_subsegment('subsegment_name') as subsegment:
+        subsegment.put_annotation('key', 'value')
+        # Do something here
+    with xray_recorder.in_subsegment('subsegment2') as subsegment:
+        subsegment.put_annotation('key2', 'value2')
+        # Do something else 
+```
+
+async versions of context managers:
+
+```python
+from aws_xray_sdk.core import xray_recorder
+
+async with xray_recorder.in_segment_async('segment_name') as segment:
+    # Add metadata or annotation here if necessary
+    segment.put_metadata('key', dict, 'namespace')
+    async with xray_recorder.in_subsegment_async('subsegment_name') as subsegment:
+        subsegment.put_annotation('key', 'value')
+        # Do something here
+    async with xray_recorder.in_subsegment_async('subsegment2') as subsegment:
+        subsegment.put_annotation('key2', 'value2')
+        # Do something else 
+```
+
+Default begin/end functions:
+
 ```python
 from aws_xray_sdk.core import xray_recorder
 
@@ -85,6 +119,8 @@ xray_recorder.end_segment()
 
 ### Capture
 
+As a decorator:
+
 ```python
 from aws_xray_sdk.core import xray_recorder
 
@@ -95,6 +131,19 @@ def myfunc():
 myfunc()
 ```
 
+or as a context manager:
+
+```python
+from aws_xray_sdk.core import xray_recorder
+
+with xray_recorder.capture('subsegment_name') as subsegment:
+    # Do something here
+    subsegment.put_annotation('mykey', val)
+    # Do something more
+```
+
+Async capture as decorator:
+
 ```python
 from aws_xray_sdk.core import xray_recorder
 
@@ -106,6 +155,17 @@ async def main():
     await myfunc()
 ```
 
+or as context manager:
+
+```python
+from aws_xray_sdk.core import xray_recorder
+
+async with xray_recorder.capture_async('subsegment_name') as subsegment:
+    # Do something here
+    subsegment.put_annotation('mykey', val)
+    # Do something more
+```
+
 ### Adding annotations/metadata using recorder
 
 ```python

aws_xray_sdk/core/async_recorder.py

@@ -4,6 +4,37 @@
 
 from aws_xray_sdk.core.recorder import AWSXRayRecorder
 from aws_xray_sdk.core.utils import stacktrace
+from aws_xray_sdk.core.models.subsegment import SubsegmentContextManager
+from aws_xray_sdk.core.models.segment import SegmentContextManager
+
+
+class AsyncSegmentContextManager(SegmentContextManager):
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return self.__exit__(exc_type, exc_val, exc_tb)
+
+class AsyncSubsegmentContextManager(SubsegmentContextManager):
+
+    @wrapt.decorator
+    async def __call__(self, wrapped, instance, args, kwargs):
+        func_name = self.name
+        if not func_name:
+            func_name = wrapped.__name__
+
+        return await self.recorder.record_subsegment_async(
+            wrapped, instance, args, kwargs,
+            name=func_name,
+            namespace='local',
+            meta_processor=None,
+        )
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return self.__exit__(exc_type, exc_val, exc_tb)
 
 
 class AsyncAWSXRayRecorder(AWSXRayRecorder):
@@ -15,23 +46,25 @@ def capture_async(self, name=None):
         params str name: The name of the subsegment. If not specified
         the function name will be used.
         """
+        return self.in_subsegment_async(name=name)
 
-        @wrapt.decorator
-        async def wrapper(wrapped, instance, args, kwargs):
-            func_name = name
-            if not func_name:
-                func_name = wrapped.__name__
+    def in_segment_async(self, name=None, **segment_kwargs):
+        """
+        Return a segment async context manger.
 
-            result = await self.record_subsegment_async(
-                wrapped, instance, args, kwargs,
-                name=func_name,
-                namespace='local',
-                meta_processor=None,
-            )
+        :param str name: the name of the segment
+        :param dict segment_kwargs: remaining arguments passed directly to `begin_segment`
+        """
+        return AsyncSegmentContextManager(self, name=name, **segment_kwargs)
 
-            return result
+    def in_subsegment_async(self, name=None, **subsegment_kwargs):
+        """
+        Return a subsegment async context manger.
 
-        return wrapper
+        :param str name: the name of the segment
+        :param dict segment_kwargs: remaining arguments passed directly to `begin_segment`
+        """
+        return AsyncSubsegmentContextManager(self, name=name, **subsegment_kwargs)
 
     async def record_subsegment_async(self, wrapped, instance, args, kwargs, name,
                                       namespace, meta_processor):

aws_xray_sdk/core/models/segment.py

@@ -1,4 +1,5 @@
 import copy
+import traceback
 
 from .entity import Entity
 from .traceid import TraceId
@@ -8,6 +9,37 @@
 ORIGIN_TRACE_HEADER_ATTR_KEY = '_origin_trace_header'
 
 
+class SegmentContextManager:
+    """
+    Wrapper for segment and recorder to provide segment context manager.
+    """
+
+    def __init__(self, recorder, name=None, **segment_kwargs):
+        self.name = name
+        self.segment_kwargs = segment_kwargs
+        self.recorder = recorder
+        self.segment = None
+
+    def __enter__(self):
+        self.segment = self.recorder.begin_segment(
+            name=self.name, **self.segment_kwargs)
+        return self.segment
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.segment is None:
+            return
+
+        if exc_type is not None:
+            self.segment.add_exception(
+                exc_val,
+                traceback.extract_tb(
+                    exc_tb,
+                    limit=self.recorder.max_trace_back,
+                )
+            )
+        self.recorder.end_segment()
+
+
 class Segment(Entity):
     """
     The compute resources running your application logic send data

aws_xray_sdk/core/models/subsegment.py

@@ -1,9 +1,56 @@
 import copy
+import traceback
+
+import wrapt
 
 from .entity import Entity
 from ..exceptions.exceptions import SegmentNotFoundException
 
 
+class SubsegmentContextManager:
+    """
+    Wrapper for segment and recorder to provide segment context manager.
+    """
+
+    def __init__(self, recorder, name=None, **subsegment_kwargs):
+        self.name = name
+        self.subsegment_kwargs = subsegment_kwargs
+        self.recorder = recorder
+        self.subsegment = None
+
+    @wrapt.decorator
+    def __call__(self, wrapped, instance, args, kwargs):
+        func_name = self.name
+        if not func_name:
+            func_name = wrapped.__name__
+
+        return self.recorder.record_subsegment(
+            wrapped, instance, args, kwargs,
+            name=func_name,
+            namespace='local',
+            meta_processor=None,
+        )
+
+    def __enter__(self):
+        self.subsegment = self.recorder.begin_subsegment(
+            name=self.name, **self.subsegment_kwargs)
+        return self.subsegment
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.subsegment is None:
+            return
+
+        if exc_type is not None:
+            self.subsegment.add_exception(
+                exc_val,
+                traceback.extract_tb(
+                    exc_tb,
+                    limit=self.recorder.max_trace_back,
+                )
+            )
+        self.recorder.end_subsegment()
+
+
 class Subsegment(Entity):
     """
     The work done in a single segment can be broke down into subsegments.

aws_xray_sdk/core/recorder.py

@@ -5,11 +5,9 @@
 import platform
 import time
 
-import wrapt
-
 from aws_xray_sdk.version import VERSION
-from .models.segment import Segment
-from .models.subsegment import Subsegment
+from .models.segment import Segment, SegmentContextManager
+from .models.subsegment import Subsegment, SubsegmentContextManager
 from .models.default_dynamic_naming import DefaultDynamicNaming
 from .models.dummy_entities import DummySegment, DummySubsegment
 from .emitters.udp_emitter import UDPEmitter
@@ -178,6 +176,24 @@ class to have your own implementation of the streaming process.
             self.sampler.load_settings(DaemonConfig(daemon_address),
                                        self.context, self._origin)
 
+    def in_segment(self, name=None, **segment_kwargs):
+        """
+        Return a segment context manger.
+
+        :param str name: the name of the segment
+        :param dict segment_kwargs: remaining arguments passed directly to `begin_segment`
+        """
+        return SegmentContextManager(self, name=name, **segment_kwargs)
+
+    def in_subsegment(self, name=None, **subsegment_kwargs):
+        """
+        Return a subsegment context manger.
+
+        :param str name: the name of the subsegment
+        :param dict segment_kwargs: remaining arguments passed directly to `begin_subsegment`
+        """
+        return SubsegmentContextManager(self, name=name, **subsegment_kwargs)
+
     def begin_segment(self, name=None, traceid=None,
                       parent_id=None, sampling=None):
         """
@@ -369,20 +385,7 @@ def capture(self, name=None):
         params str name: The name of the subsegment. If not specified
         the function name will be used.
         """
-        @wrapt.decorator
-        def wrapper(wrapped, instance, args, kwargs):
-            func_name = name
-            if not func_name:
-                func_name = wrapped.__name__
-
-            return self.record_subsegment(
-                wrapped, instance, args, kwargs,
-                name=func_name,
-                namespace='local',
-                meta_processor=None,
-            )
-
-        return wrapper
+        return self.in_subsegment(name=name)
 
     def record_subsegment(self, wrapped, instance, args, kwargs, name,
                           namespace, meta_processor):

tests/test_async_recorder.py

@@ -42,3 +42,16 @@ async def test_capture(loop):
     service = segment.service
     assert platform.python_implementation() == service.get('runtime')
     assert platform.python_version() == service.get('runtime_version')
+
+
+async def test_async_context_managers(loop):
+    xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop))
+
+    async with xray_recorder.in_segment_async('segment') as segment:
+        async with xray_recorder.capture_async('aio_capture') as subsegment:
+            assert segment.subsegments[0].name == 'aio_capture'
+        assert subsegment.in_progress is  False
+        async with xray_recorder.in_subsegment_async('in_sub') as subsegment:
+            assert segment.subsegments[1].name == 'in_sub'
+            assert subsegment.in_progress is  True
+        assert subsegment.in_progress is  False

tests/test_recorder.py

@@ -120,3 +120,48 @@ def test_first_begin_segment_sampled():
     segment = xray_recorder.begin_segment('name')
 
     assert segment.sampled
+
+
+def test_in_segment_closing():
+    xray_recorder = get_new_stubbed_recorder()
+    xray_recorder.configure(sampling=False)
+
+    with xray_recorder.in_segment('name') as segment:
+        assert segment.in_progress is True
+        segment.put_metadata('key1', 'value1')
+        segment.put_annotation('key2', 'value2')
+        with xray_recorder.in_subsegment('subsegment') as subsegment:
+            assert subsegment.in_progress is True
+
+        with xray_recorder.capture('capture') as subsegment:
+            assert subsegment.in_progress is True
+            assert subsegment.name == 'capture'
+
+    assert subsegment.in_progress is False
+    assert segment.in_progress is False
+    assert segment.annotations['key2'] == 'value2'
+    assert segment.metadata['default']['key1'] == 'value1'
+
+
+def test_in_segment_exception():
+    xray_recorder = get_new_stubbed_recorder()
+    xray_recorder.configure(sampling=False)
+
+    with pytest.raises(Exception):
+        with xray_recorder.in_segment('name') as segment:
+            assert segment.in_progress is True
+            assert 'exceptions' not in segment.cause
+            raise Exception('test exception')
+
+    assert segment.in_progress is False
+    assert segment.fault is True
+    assert len(segment.cause['exceptions']) == 1
+
+
+    with pytest.raises(Exception):
+        with xray_recorder.in_segment('name') as segment:
+            with xray_recorder.in_subsegment('name') as subsegment:
+                    assert subsegment.in_progress is True
+                    raise Exception('test exception')
+
+    assert len(subsegment.cause['exceptions']) == 1