Give up when there is no open segment and LOG_ERROR is configured

CHANGELOG.rst

@@ -5,6 +5,7 @@ CHANGELOG
 unreleased
 ==========
 * feature: Added Sqlalchemy parameterized query capture. `PR34 <https://github.com/aws/aws-xray-sdk-python/pull/34>`_.
+* bugfix: Added new `raise_if_not_subsegment` parameter for Aiohttp Client tracing `PR58 <https://github.com/aws/aws-xray-sdk-python/pull/58>`_.
 
 1.0
 ===

aws_xray_sdk/ext/aiohttp/client.py

@@ -24,18 +24,30 @@
 async def begin_subsegment(session, trace_config_ctx, params):
     name = trace_config_ctx.name if trace_config_ctx.name else strip_url(str(params.url))
     subsegment = xray_recorder.begin_subsegment(name, REMOTE_NAMESPACE)
-    subsegment.put_http_meta(http.METHOD, params.method)
-    subsegment.put_http_meta(http.URL, params.url.human_repr())
-    inject_trace_header(params.headers, subsegment)
+
+    # No-op if subsegment is `None` due to `LOG_ERROR`.
+    if not subsegment:
+        trace_config_ctx.give_up = True
+    else:
+        trace_config_ctx.give_up = False
+        subsegment.put_http_meta(http.METHOD, params.method)
+        subsegment.put_http_meta(http.URL, params.url.human_repr())
+        inject_trace_header(params.headers, subsegment)
 
 
 async def end_subsegment(session, trace_config_ctx, params):
+    if trace_config_ctx.give_up:
+        return
+
     subsegment = xray_recorder.current_subsegment()
     subsegment.put_http_meta(http.STATUS, params.response.status)
     xray_recorder.end_subsegment()
 
 
 async def end_subsegment_with_exception(session, trace_config_ctx, params):
+    if trace_config_ctx.give_up:
+        return
+
     subsegment = xray_recorder.current_subsegment()
     subsegment.add_exception(
         params.exception,
@@ -54,10 +66,14 @@ def aws_xray_trace_config(name=None):
                  be used as identifier.
     :returns: TraceConfig.
     """
-    trace_config = aiohttp.TraceConfig(
-        trace_config_ctx_factory=lambda trace_request_ctx: SimpleNamespace(name=name,
-                                                                           trace_request_ctx=trace_request_ctx)
-    )
+
+    def _trace_config_ctx_factory(trace_request_ctx):
+        return SimpleNamespace(
+            name=name,
+            trace_request_ctx=trace_request_ctx
+        )
+
+    trace_config = aiohttp.TraceConfig(trace_config_ctx_factory=_trace_config_ctx_factory)
     trace_config.on_request_start.append(begin_subsegment)
     trace_config.on_request_end.append(end_subsegment)
     trace_config.on_request_exception.append(end_subsegment_with_exception)

tests/ext/aiohttp/test_client.py

@@ -3,6 +3,7 @@
 
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.async_context import AsyncContext
+from aws_xray_sdk.core.exceptions.exceptions import SegmentNotFoundException
 from aws_xray_sdk.ext.util import strip_url
 from aws_xray_sdk.ext.aiohttp.client import aws_xray_trace_config
 from aws_xray_sdk.ext.aiohttp.client import REMOTE_NAMESPACE, LOCAL_NAMESPACE
@@ -130,3 +131,27 @@ async def test_invalid_url(loop, recorder):
 
     exception = subsegment.cause['exceptions'][0]
     assert exception.type == 'ClientConnectorError'
+
+
+async def test_no_segment_raise(loop, recorder):
+    xray_recorder.configure(context_missing='RUNTIME_ERROR')
+    trace_config = aws_xray_trace_config()
+    status_code = 200
+    url = 'http://{}/status/{}?foo=bar'.format(BASE_URL, status_code)
+    with pytest.raises(SegmentNotFoundException):
+        async with ClientSession(loop=loop, trace_configs=[trace_config]) as session:
+            async with session.get(url):
+                pass
+
+
+async def test_no_segment_not_raise(loop, recorder):
+    xray_recorder.configure(context_missing='LOG_ERROR')
+    trace_config = aws_xray_trace_config()
+    status_code = 200
+    url = 'http://{}/status/{}?foo=bar'.format(BASE_URL, status_code)
+    async with ClientSession(loop=loop, trace_configs=[trace_config]) as session:
+        async with session.get(url) as resp:
+            status_received = resp.status
+
+    # Just check that the request was done correctly
+    assert status_received == status_code