Flask: Add Status Code Overwrite Safety (#235)

* Only set status code if it hasnt already been set

* Yet another exceptin test

Author: pglombardo

instana/instrumentation/flask/vanilla.py

@@ -83,13 +83,13 @@ def teardown_request_with_instana(*argv, **kwargs):
     In the case of exceptions, after_request_with_instana isn't called
     so we capture those cases here.
     """
-    if hasattr(flask.g, 'scope'):
-        if flask.g.scope is not None:
-            if len(argv) > 0 and argv[0] is not None:
-                scope = flask.g.scope
-                scope.span.log_exception(argv[0])
+    if hasattr(flask.g, 'scope') and flask.g.scope is not None:
+        if len(argv) > 0 and argv[0] is not None:
+            scope = flask.g.scope
+            scope.span.log_exception(argv[0])
+            if ext.HTTP_STATUS_CODE not in scope.span.tags:
                 scope.span.set_tag(ext.HTTP_STATUS_CODE, 500)
-            scope.close()
+        flask.g.scope.close()
         flask.g.scope = None
 
 

instana/instrumentation/flask/with_blinker.py

@@ -126,12 +126,12 @@ def teardown_request_with_instana(*argv, **kwargs):
     In the case of exceptions, after_request_with_instana isn't called
     so we capture those cases here.
     """
-
     if hasattr(flask.g, 'scope') and flask.g.scope is not None:
         if len(argv) > 0 and argv[0] is not None:
             scope = flask.g.scope
             scope.span.log_exception(argv[0])
-            scope.span.set_tag(ext.HTTP_STATUS_CODE, 500)
+            if ext.HTTP_STATUS_CODE not in scope.span.tags:
+                scope.span.set_tag(ext.HTTP_STATUS_CODE, 500)
         flask.g.scope.close()
         flask.g.scope = None
 

tests/apps/flaskalino.py

@@ -39,6 +39,22 @@ def to_dict(self):
         return rv
 
 
+class NotFound(Exception):
+    status_code = 404
+
+    def __init__(self, message, status_code=None, payload=None):
+        Exception.__init__(self)
+        self.message = message
+        if status_code is not None:
+            self.status_code = status_code
+        self.payload = payload
+
+    def to_dict(self):
+        rv = dict(self.payload or ())
+        rv['message'] = self.message
+        return rv
+
+
 @app.route("/")
 def hello():
     return "<center><h1>🐍 Hello Stan! 🦄</h1></center>"
@@ -86,6 +102,11 @@ def fourhundred():
     return "Simulated Bad Request", 400
 
 
+@app.route("/custom-404")
+def custom404():
+    raise NotFound("My custom 404 message")
+
+
 @app.route("/405")
 def fourhundredfive():
     return "Simulated Method not allowed", 405
@@ -132,6 +153,7 @@ def response_headers():
     resp.headers['X-Capture-This'] = 'Ok'
     return resp
 
+
 @app.errorhandler(InvalidUsage)
 def handle_invalid_usage(error):
     logger.error("InvalidUsage error handler invoked")
@@ -140,6 +162,12 @@ def handle_invalid_usage(error):
     return response
 
 
+@app.errorhandler(404)
+@app.errorhandler(NotFound)
+def handle_not_found(e):
+    return "blah: %s" % str(e), 404
+
+
 if __name__ == '__main__':
     flask_server.request_queue_size = 20
     flask_server.serve_forever()

tests/test_flask.py

@@ -320,6 +320,74 @@ def test_301(self):
         # We should NOT have a path template for this route
         self.assertIsNone(wsgi_span.data["http"]["path_tpl"])
 
+    def test_custom_404(self):
+        with tracer.start_active_span('test'):
+            response = self.http.request('GET', testenv["wsgi_server"] + '/custom-404')
+
+        spans = self.recorder.queued_spans()
+
+        self.assertEqual(3, len(spans))
+
+        wsgi_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        assert response
+        self.assertEqual(404, response.status)
+
+        # assert('X-Instana-T' in response.headers)
+        # assert(int(response.headers['X-Instana-T'], 16))
+        # self.assertEqual(response.headers['X-Instana-T'], wsgi_span.t)
+        #
+        # assert('X-Instana-S' in response.headers)
+        # assert(int(response.headers['X-Instana-S'], 16))
+        # self.assertEqual(response.headers['X-Instana-S'], wsgi_span.s)
+        #
+        # assert('X-Instana-L' in response.headers)
+        # self.assertEqual(response.headers['X-Instana-L'], '1')
+        #
+        # assert('Server-Timing' in response.headers)
+        # server_timing_value = "intid;desc=%s" % wsgi_span.t
+        # self.assertEqual(response.headers['Server-Timing'], server_timing_value)
+
+        self.assertIsNone(tracer.active_span)
+
+        # Same traceId
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(test_span.t, wsgi_span.t)
+
+        # Parent relationships
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(wsgi_span.p, urllib3_span.s)
+
+        # Error logging
+        self.assertIsNone(test_span.ec)
+        self.assertEqual(None, urllib3_span.ec)
+        self.assertEqual(None, wsgi_span.ec)
+
+        # wsgi
+        self.assertEqual("wsgi", wsgi_span.n)
+        self.assertEqual('127.0.0.1:' + str(testenv['wsgi_port']), wsgi_span.data["http"]["host"])
+        self.assertEqual('/custom-404', wsgi_span.data["http"]["url"])
+        self.assertEqual('GET', wsgi_span.data["http"]["method"])
+        self.assertEqual(404, wsgi_span.data["http"]["status"])
+        self.assertIsNone(wsgi_span.data["http"]["error"])
+        self.assertIsNotNone(wsgi_span.stack)
+        self.assertEqual(2, len(wsgi_span.stack))
+
+        # urllib3
+        self.assertEqual("test", test_span.data["sdk"]["name"])
+        self.assertEqual("urllib3", urllib3_span.n)
+        self.assertEqual(404, urllib3_span.data["http"]["status"])
+        self.assertEqual(testenv["wsgi_server"] + '/custom-404', urllib3_span.data["http"]["url"])
+        self.assertEqual("GET", urllib3_span.data["http"]["method"])
+        self.assertIsNotNone(urllib3_span.stack)
+        self.assertTrue(type(urllib3_span.stack) is list)
+        self.assertTrue(len(urllib3_span.stack) > 1)
+
+        # We should NOT have a path template for this route
+        self.assertIsNone(wsgi_span.data["http"]["path_tpl"])
+
     def test_404(self):
         with tracer.start_active_span('test'):
             response = self.http.request('GET', testenv["wsgi_server"] + '/11111111111')