diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 13d290f1..f45291bc 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,6 +5,7 @@ CHANGELOG unreleased ========== * feature: Added Sqlalchemy parameterized query capture. `PR34 `_. +* bugfix: Added new `raise_if_not_subsegment` parameter for Aiohttp Client tracing `PR58 `_. 1.0 === diff --git a/aws_xray_sdk/ext/aiohttp/client.py b/aws_xray_sdk/ext/aiohttp/client.py index ed9fbf32..4843bfeb 100644 --- a/aws_xray_sdk/ext/aiohttp/client.py +++ b/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) diff --git a/tests/ext/aiohttp/test_client.py b/tests/ext/aiohttp/test_client.py index 1a0db0b7..36ffa138 100644 --- a/tests/ext/aiohttp/test_client.py +++ b/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