From 6e9b26bdf248a3f62b6ab9366ee456be5524b3b7 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Wed, 3 Feb 2021 17:48:50 -0500 Subject: [PATCH 1/2] always closing segment if present in teardown_request handler --- aws_xray_sdk/ext/flask/middleware.py | 10 +++++----- tests/ext/flask/test_flask.py | 18 +++++++++++++++++- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 53b45787..d773cca6 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -91,8 +91,6 @@ def _after_request(self, response): return response def _handle_exception(self, exception): - if not exception: - return segment = None try: if self.in_lambda_ctx: @@ -104,9 +102,11 @@ def _handle_exception(self, exception): if not segment: return - segment.put_http_meta(http.STATUS, 500) - stack = stacktrace.get_stacktrace(limit=self._recorder._max_trace_back) - segment.add_exception(exception, stack) + if 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_ctx: self._recorder.end_subsegment() else: diff --git a/tests/ext/flask/test_flask.py b/tests/ext/flask/test_flask.py index 4e587610..283c0e89 100644 --- a/tests/ext/flask/test_flask.py +++ b/tests/ext/flask/test_flask.py @@ -9,7 +9,6 @@ from tests.util import get_new_stubbed_recorder import os - # define a flask app for testing purpose app = Flask(__name__) @@ -29,6 +28,11 @@ def fault(): return {}['key'] +@app.route('/fault_no_exception') +def fault_no_exception(): + return "SomeException", 500 + + @app.route('/template') def template(): return render_template_string('hello template') @@ -108,6 +112,18 @@ def test_fault(): assert exception.type == 'KeyError' +def test_fault_no_exception(): + path = '/fault_no_exception' + app.get(path) + segment = recorder.emitter.pop() + assert not segment.in_progress + assert segment.fault + + response = segment.http['response'] + assert response['status'] == 500 + assert segment.cause == {} + + def test_render_template(): path = '/template' app.get(path) From 5a08b9ca59e91f225914bcb158dcdb645ac77a48 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 4 Feb 2021 09:59:48 -0500 Subject: [PATCH 2/2] close segment in teardown_request only. rename teardown method --- aws_xray_sdk/ext/flask/middleware.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index d773cca6..c0c1e226 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -17,7 +17,7 @@ def __init__(self, app, recorder): self._recorder = recorder self.app.before_request(self._before_request) self.app.after_request(self._after_request) - self.app.teardown_request(self._handle_exception) + self.app.teardown_request(self._teardown_request) self.in_lambda_ctx = False if check_in_lambda() and type(self._recorder.context) == LambdaContext: @@ -81,16 +81,9 @@ def _after_request(self, response): if cont_len: segment.put_http_meta(http.CONTENT_LENGTH, int(cont_len)) - if response.status_code >= 500: - return response - - if self.in_lambda_ctx: - self._recorder.end_subsegment() - else: - self._recorder.end_segment() return response - def _handle_exception(self, exception): + def _teardown_request(self, exception): segment = None try: if self.in_lambda_ctx: