diff --git a/README.md b/README.md index 6e69df1d..c02e2385 100644 --- a/README.md +++ b/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 diff --git a/aws_xray_sdk/core/async_recorder.py b/aws_xray_sdk/core/async_recorder.py index 952c5762..c679fa0b 100644 --- a/aws_xray_sdk/core/async_recorder.py +++ b/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): diff --git a/aws_xray_sdk/core/models/segment.py b/aws_xray_sdk/core/models/segment.py index b36d4a17..d9cf2012 100644 --- a/aws_xray_sdk/core/models/segment.py +++ b/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 diff --git a/aws_xray_sdk/core/models/subsegment.py b/aws_xray_sdk/core/models/subsegment.py index 2cf5c1e2..f5227883 100644 --- a/aws_xray_sdk/core/models/subsegment.py +++ b/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. diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index fe579615..a33ab7f4 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/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): diff --git a/tests/test_async_recorder.py b/tests/test_async_recorder.py index fe018a5d..eba147f7 100644 --- a/tests/test_async_recorder.py +++ b/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 diff --git a/tests/test_recorder.py b/tests/test_recorder.py index 38de91ae..afe7f967 100644 --- a/tests/test_recorder.py +++ b/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