From bc4c4b5d817ea0d7b0fb36140a059ca8f98ce95f Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Wed, 27 Feb 2019 15:30:56 -0800 Subject: [PATCH] Flask & Django create subsegments only if in LambdaContext & Lambda Environment (#138) * Added unit tests to ensure these frameworks generate segmentby default and subsegments if in lambda context and environment. --- aws_xray_sdk/ext/django/middleware.py | 12 ++++++------ aws_xray_sdk/ext/flask/middleware.py | 18 +++++++++--------- tests/ext/django/test_middleware.py | 9 ++++++++- tests/ext/flask/test_flask.py | 26 ++++++++++++++++++++++++-- 4 files changed, 47 insertions(+), 18 deletions(-) diff --git a/aws_xray_sdk/ext/django/middleware.py b/aws_xray_sdk/ext/django/middleware.py index c2905d30..d3eeec7c 100644 --- a/aws_xray_sdk/ext/django/middleware.py +++ b/aws_xray_sdk/ext/django/middleware.py @@ -5,7 +5,7 @@ from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header -from aws_xray_sdk.core.lambda_launcher import check_in_lambda +from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext log = logging.getLogger(__name__) @@ -25,10 +25,10 @@ class XRayMiddleware(object): def __init__(self, get_response): self.get_response = get_response - self.in_lambda = False + self.in_lambda_ctx = False - if check_in_lambda(): - self.in_lambda = True + if check_in_lambda() and type(xray_recorder.context) == LambdaContext: + self.in_lambda_ctx = True # hooks for django version >= 1.10 def __call__(self, request): @@ -51,7 +51,7 @@ def __call__(self, request): sampling_req=sampling_req, ) - if self.in_lambda: + if self.in_lambda_ctx: segment = xray_recorder.begin_subsegment(name) else: segment = xray_recorder.begin_segment( @@ -83,7 +83,7 @@ def __call__(self, request): segment.put_http_meta(http.CONTENT_LENGTH, length) response[http.XRAY_HEADER] = prepare_response_header(xray_header, segment) - if self.in_lambda: + if self.in_lambda_ctx: xray_recorder.end_subsegment() else: xray_recorder.end_segment() diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 1a09913b..5d84ddf6 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -5,7 +5,7 @@ from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header -from aws_xray_sdk.core.lambda_launcher import check_in_lambda +from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext class XRayMiddleware(object): @@ -18,10 +18,10 @@ def __init__(self, app, recorder): self.app.before_request(self._before_request) self.app.after_request(self._after_request) self.app.teardown_request(self._handle_exception) - self.in_lambda = False + self.in_lambda_ctx = False - if check_in_lambda(): - self.in_lambda = True + if check_in_lambda() and type(self._recorder.context) == LambdaContext: + self.in_lambda_ctx = True _patch_render(recorder) @@ -44,7 +44,7 @@ def _before_request(self): sampling_req=sampling_req, ) - if self.in_lambda: + if self.in_lambda_ctx: segment = self._recorder.begin_subsegment(name) else: segment = self._recorder.begin_segment( @@ -67,7 +67,7 @@ def _before_request(self): segment.put_http_meta(http.CLIENT_IP, req.remote_addr) def _after_request(self, response): - if self.in_lambda: + if self.in_lambda_ctx: segment = self._recorder.current_subsegment() else: segment = self._recorder.current_segment() @@ -81,7 +81,7 @@ def _after_request(self, response): if cont_len: segment.put_http_meta(http.CONTENT_LENGTH, int(cont_len)) - if self.in_lambda: + if self.in_lambda_ctx: self._recorder.end_subsegment() else: self._recorder.end_segment() @@ -92,7 +92,7 @@ def _handle_exception(self, exception): return segment = None try: - if self.in_lambda: + if self.in_lambda_ctx: segment = self._recorder.current_subsegment() else: segment = self._recorder.current_segment() @@ -104,7 +104,7 @@ def _handle_exception(self, exception): segment.put_http_meta(http.STATUS, 500) stack = stacktrace.get_stacktrace(limit=self._recorder._max_trace_back) segment.add_exception(exception, stack) - if self.in_lambda: + if self.in_lambda_ctx: self._recorder.end_subsegment() else: self._recorder.end_segment() diff --git a/tests/ext/django/test_middleware.py b/tests/ext/django/test_middleware.py index cb36ddf9..a428d9e8 100644 --- a/tests/ext/django/test_middleware.py +++ b/tests/ext/django/test_middleware.py @@ -5,7 +5,7 @@ from aws_xray_sdk.core import xray_recorder, lambda_launcher from aws_xray_sdk.core.context import Context -from aws_xray_sdk.core.models import http, facade_segment +from aws_xray_sdk.core.models import http, facade_segment, segment from tests.util import get_new_stubbed_recorder import os @@ -132,3 +132,10 @@ def test_lambda_serverless(self): self.client.get(url) segment = new_recorder.emitter.pop() assert not segment + + def test_lambda_default_ctx(self): + # Track to make sure that Django will default to generating segments if context is not the lambda context + url = reverse('200ok') + self.client.get(url) + cur_segment = xray_recorder.emitter.pop() + assert type(cur_segment) == segment.Segment diff --git a/tests/ext/flask/test_flask.py b/tests/ext/flask/test_flask.py index b5bbf1f7..d8a7a671 100644 --- a/tests/ext/flask/test_flask.py +++ b/tests/ext/flask/test_flask.py @@ -5,7 +5,7 @@ from aws_xray_sdk.ext.flask.middleware import XRayMiddleware from aws_xray_sdk.core.context import Context from aws_xray_sdk.core import lambda_launcher -from aws_xray_sdk.core.models import http, facade_segment +from aws_xray_sdk.core.models import http, facade_segment, segment from tests.util import get_new_stubbed_recorder import os @@ -186,7 +186,7 @@ def trace_header(): return 'ok' middleware = XRayMiddleware(new_app, new_recorder) - middleware.in_lambda = True + middleware.in_lambda_ctx = True app_client = new_app.test_client() @@ -197,3 +197,25 @@ def trace_header(): path2 = '/trace_header' app_client.get(path2, headers={http.XRAY_HEADER: 'k1=v1'}) + + +def test_lambda_default_ctx(): + # Track to make sure that Flask will default to generating segments if context is not the lambda context + new_recorder = get_new_stubbed_recorder() + new_recorder.configure(service='test', sampling=False) + new_app = Flask(__name__) + + @new_app.route('/segment') + def subsegment(): + # Test in between request and make sure Lambda that uses default context generates a segment. + assert new_recorder.current_segment() + assert type(new_recorder.current_segment()) == segment.Segment + return 'ok' + + XRayMiddleware(new_app, new_recorder) + app_client = new_app.test_client() + + path = '/segment' + app_client.get(path) + segment = recorder.emitter.pop() + assert not segment # Segment should be none because it's created and ended by the middleware