Provide xray_recorder.capture() as context manager

[Dunedan]

This is a feature request to be able to use the functionality of AWSXRayRecorder.capture() as context manager.
The motivation for that is to have something to replace begin_subsegment() and end_subsegment() calls with a context manager which does proper exception handling. That's especially important when using nested subsegments, as subsegment-traces in one "branch" below the current segment will only be streamed to X-Ray if all subsegments in that branch got closed properly. When doing no proper exception handling manually to ensure to close all subsegments that quickly leads to missing data in X-Ray. What contributes to that is that AWSXRayRecorder.end_subsegment() doesn't close a specific subsegment, but rather the current open one and because of that AWSXRayRecorder isn't able to notice if one got missed.

Currently one would have to do something like that for every subsegment:

def some_function():
    print("code without custom subsegment")
    subsegment = xray_recorder.begin_subsegment("my_subsegment")
    exception = None
    stack = None
    try:
        print("code to trace as custom subsegment")
    except Exception as exception:
        stack = traceback.extract_stack(limit=self.max_trace_back)
        raise
    finally:
        if subsegment is not None:
            if exception:
                subsegment.add_exception(exception, stack)
            xray_recorder.end_subsegment()

With capture() available as context manager that would be as simple as:

def some_function():
    print("code without custom subsegment")
    with xray_recorder.capture("my_subsegment"):
        print("code to trace as custom subsegment")

As a workaround we currently use the following custom context manager, but it would be nicer if AWSXRayRecorder.capture() would support it directly of course:

class XRayCaptureContextManager:
    def __init__(self, name):
        self.name = name
        self.subsegment = None

    def __enter__(self):
        self.subsegment = xray_recorder.begin_subsegment(self.name)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.subsegment is not None:
            if exc_val:
                self.subsegment.add_exception(exc_val, traceback.extract_tb(exc_tb))
            xray_recorder.end_subsegment()

[haotianw465]

Thank you for the valuable feedback. I have one question that I would like to get some clarification. The capture decorator provided by the recorder is calling this function https://github.com/aws/aws-xray-sdk-python/blob/master/aws_xray_sdk/core/recorder.py#L364 which currrently does proper error handling. So based on my understanding you can use the following syntax (decorator style):

@xray_recorder.capture()
def some_function():
    print("code to trace as custom subsegment")

Is there anything missing on this code that we can improve so it can fit your use case? Or you mean to use the with style instead of a decorator?

[Dunedan]

The capture-decorator is fine as it is. It's perfect to get proper error handling when capturing whole functions. What we want to be able to do is to capture arbitrary parts of code inside of functions as well, with the same error handling properties as the capture-decorator does provide. That's why we would like to see the same functionality as context manager as well, as decorators can only be applied to functions, while context managers can be used inside of functions.

The example for the context manager I included above is a bit silly as the only code inside the function, which is not going to be traced, is a print function, but just imagine more complicated code flow. We e.g. have the case that we have a single function in which we want to trace multiple parts of code separately, something not possible with a decorator.

Does that answer your question?

[haotianw465]

Yes. Thank you for the clarification. We will cooperate this use case into future improvements and releases. A PR is always welcome.

[haotianw465]

This feature is released as part of 2.2.0: https://pypi.org/project/aws-xray-sdk/. Closing.