Give up when there is no open segment and LOG_ERROR is configured (#58)

* New flag `raise_if_not_subsegment` for Aiottp client

Without setting this flag as False, any HTTP call done with the Aiohttp
Client when there is no an active segment will fail with the
SegmentNotFoundException exception.

This issue can be annoying for this systems that do not create
explicitly the segements and relay on the sampled incomming requests,
which create automatically segments just for a sub set of the incomming
request.

* New entry within the CHANGELOG

* Refactorized a bit the code

* Give up when there is no open segment and LOG_ERROR is configured for context_missing

Author: pfreixes

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