fix(flask): handle missing segment gracefully in after_request

aws_xray_sdk/ext/flask/middleware.py

@@ -67,10 +67,17 @@ def _before_request(self):
             segment.put_http_meta(http.CLIENT_IP, req.remote_addr)
 
     def _after_request(self, response):
-        if self.in_lambda_ctx:
-            segment = self._recorder.current_subsegment()
-        else:
-            segment = self._recorder.current_segment()
+        segment = None
+        try:
+            if self.in_lambda_ctx:
+                segment = self._recorder.current_subsegment()
+            else:
+                segment = self._recorder.current_segment()
+        except Exception:
+            pass
+        if not segment:
+            return response
+
         segment.put_http_meta(http.STATUS, response.status_code)
 
         origin_header = segment.get_origin_trace_header()

tests/ext/flask/test_flask.py

@@ -238,3 +238,20 @@ def subsegment():
     app_client.get(path)
     segment = recorder.emitter.pop()
     assert not segment  # Segment should be none because it's created and ended by the middleware
+
+
+def test_missing_segment_in_after_request():
+    new_recorder = get_new_stubbed_recorder()
+    new_recorder.configure(service='test', sampling=False)
+    new_app = Flask(__name__)
+
+    @new_app.before_request
+    def early_before_request():
+        return 'Early response', 400
+
+    XRayMiddleware(new_app, new_recorder)
+    app_client = new_app.test_client()
+
+    response = app_client.get('/early')
+    assert response.status_code == 400
+    assert response.data == b'Early response'