Attach stack traces and exceptions to Flask traces (#151)

Author: wkiser

opencensus/trace/ext/flask/flask_middleware.py

@@ -12,16 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import flask
 import logging
+import sys
+
+import flask
+from google.rpc import code_pb2
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
-from opencensus.trace.propagation import google_cloud_format
+from opencensus.trace import stack_trace
+from opencensus.trace import status
+from opencensus.trace import tracer as tracer_module
 from opencensus.trace.exporters import print_exporter
 from opencensus.trace.ext import utils
+from opencensus.trace.propagation import google_cloud_format
 from opencensus.trace.samplers import always_on
-from opencensus.trace import tracer as tracer_module
 
 _FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
 
@@ -81,6 +86,7 @@ def __init__(self, app, blacklist_paths=None, sampler=None, exporter=None,
     def setup_trace(self):
         self.app.before_request(self._before_request)
         self.app.after_request(self._after_request)
+        self.app.teardown_request(self._teardown_request)
 
     def _before_request(self):
         """A function to be run before each request.
@@ -127,13 +133,38 @@ def _after_request(self, response):
             tracer.add_attribute_to_current_span(
                 HTTP_STATUS_CODE,
                 str(response.status_code))
+        except Exception:  # pragma: NO COVER
+            log.error('Failed to trace request', exc_info=True)
+        finally:
+            return response
+
+    def _teardown_request(self, exception):
+        # Do not trace if the url is blacklisted
+        if utils.disable_tracing_url(flask.request.url, self.blacklist_paths):
+            return
+
+        try:
+            tracer = execution_context.get_opencensus_tracer()
+
+            if exception is not None:
+                span = execution_context.get_current_span()
+                span.status = status.Status(
+                    code=code_pb2.UNKNOWN,
+                    message=str(exception)
+                )
+                # try attaching the stack trace to the span, only populated if
+                # the app has 'PROPAGATE_EXCEPTIONS', 'DEBUG', or 'TESTING'
+                # enabled
+                exc_type, _, exc_traceback = sys.exc_info()
+                if exc_traceback is not None:
+                    span.stack_trace = stack_trace.StackTrace.from_traceback(
+                        exc_traceback
+                    )
 
             tracer.end_span()
             tracer.finish()
         except Exception:  # pragma: NO COVER
             log.error('Failed to trace request', exc_info=True)
-        finally:
-            return response
 
 
 def get_flask_header():

tests/unit/trace/ext/flask/test_flask_middleware.py

@@ -19,8 +19,12 @@
 
 import flask
 import mock
+from google.rpc import code_pb2
 
 from opencensus.trace import execution_context
+from opencensus.trace import span_data
+from opencensus.trace import stack_trace
+from opencensus.trace import status
 from opencensus.trace.exporters import print_exporter
 from opencensus.trace.ext.flask import flask_middleware
 from opencensus.trace.propagation import google_cloud_format
@@ -44,6 +48,10 @@ def index():
         def health_check():
             return 'test health check'  # pragma: NO COVER
 
+        @app.route('/error')
+        def error():
+            raise Exception('error')
+
         return app
 
     def tearDown(self):
@@ -241,3 +249,36 @@ def test__after_request_blacklist(self):
 
         self.assertEqual(response.status_code, 200)
         assert isinstance(tracer, noop_tracer.NoopTracer)
+
+    def test_teardown_include_exception(self):
+        mock_exporter = mock.MagicMock()
+        app = self.create_app()
+        flask_middleware.FlaskMiddleware(app=app, exporter=mock_exporter)
+        response = app.test_client().get('/error')
+
+        self.assertEqual(response.status_code, 500)
+
+        exported_spandata = mock_exporter.export.call_args[0][0][0]
+        self.assertIsInstance(exported_spandata, span_data.SpanData)
+        self.assertIsInstance(exported_spandata.status, status.Status)
+        self.assertEqual(exported_spandata.status.code, code_pb2.UNKNOWN)
+        self.assertEqual(exported_spandata.status.message, 'error')
+
+    def test_teardown_include_exception_and_traceback(self):
+        mock_exporter = mock.MagicMock()
+        app = self.create_app()
+        app.config['TESTING'] = True
+        flask_middleware.FlaskMiddleware(app=app, exporter=mock_exporter)
+        with self.assertRaises(Exception):
+            app.test_client().get('/error')
+
+        exported_spandata = mock_exporter.export.call_args[0][0][0]
+        self.assertIsInstance(exported_spandata, span_data.SpanData)
+        self.assertIsInstance(exported_spandata.status, status.Status)
+        self.assertEqual(exported_spandata.status.code, code_pb2.UNKNOWN)
+        self.assertEqual(exported_spandata.status.message, 'error')
+        self.assertIsInstance(
+            exported_spandata.stack_trace, stack_trace.StackTrace
+        )
+        self.assertIsNotNone(exported_spandata.stack_trace.stack_trace_hash_id)
+        self.assertNotEqual(exported_spandata.stack_trace.stack_frames, [])