feat(fastapi): adding support for exception handlers (#335)

Author: maorlx

epsagon/modules/fastapi.py

@@ -5,7 +5,11 @@
 from __future__ import absolute_import
 import wrapt
 from fastapi.routing import APIRoute
-from ..wrappers.fastapi import TracingAPIRoute
+from ..wrappers.fastapi import (
+    TracingAPIRoute,
+    exception_handler_wrapper,
+    server_call_wrapper,
+)
 from ..utils import is_lambda_env, print_debug
 
 def _wrapper(wrapped, _instance, args, kwargs):
@@ -30,6 +34,23 @@ def _wrapper(wrapped, _instance, args, kwargs):
     kwargs['route_class'] = TracingAPIRoute
     return wrapped(*args, **kwargs)
 
+def _exception_handler_wrapper(wrapped, _instance, args, kwargs):
+    """
+    Wraps the handler given to add_exception_handler.
+    :param wrapped: wrapt's wrapped
+    :param instance: wrapt's instance
+    :param args: wrapt's args
+    :param kwargs: wrapt's kwargs
+    """
+    # Skip on Lambda environment since it's not relevant and might be duplicate
+    if is_lambda_env():
+        return wrapped(*args, **kwargs)
+
+    if args and len(args) == 2:
+        args = list(args)
+        args[1] = exception_handler_wrapper(args[1])
+    return wrapped(*args, **kwargs)
+
 
 def patch():
     """
@@ -41,3 +62,13 @@ def patch():
         'APIRouter.__init__',
         _wrapper
     )
+    wrapt.wrap_function_wrapper(
+        'starlette.applications',
+        'Starlette.add_exception_handler',
+        _exception_handler_wrapper
+    )
+    wrapt.wrap_function_wrapper(
+        'starlette.middleware.errors',
+        'ServerErrorMiddleware.__call__',
+        server_call_wrapper
+    )

epsagon/modules/tornado.py

@@ -208,7 +208,7 @@ def thread_pool_submit(cls, func, _, args, kwargs):
 
         try:
             trace = epsagon.trace.trace_factory.get_trace()
-            if trace:
+            if trace and trace.unique_id in cls.RUNNERS:
                 unique_id = trace.unique_id
 
         except Exception as instrumentation_exception:  # pylint: disable=W0703

epsagon/runners/fastapi.py

@@ -109,10 +109,13 @@ def _update_raw_response_body(self, response, response_type):
         )
 
 
-    def _update_status_code(self, status_code):
+    def update_status_code(self, status_code, override=True):
         """
         Updates the event with given status code.
+        :param override: indicates whether to override existing status code
         """
+        if self.resource['metadata'].get('status_code') and not override:
+            return
         self.resource['metadata']['status_code'] = status_code
         if status_code and status_code >= 500:
             self.set_error()
@@ -134,7 +137,7 @@ def _update_raw_response(self, response):
                 'Response Headers',
                 response_headers
             )
-        self._update_status_code(response.status_code)
+        self.update_status_code(response.status_code)
 
     def update_response(self, response, status_code=None):
         """
@@ -154,4 +157,4 @@ def update_response(self, response, status_code=None):
                     'Could not json encode fastapi handler response data'
                 )
             if status_code:
-                self._update_status_code(status_code)
+                self.update_status_code(status_code)

epsagon/trace.py

@@ -247,7 +247,7 @@ def _get_thread_trace(self, should_create=False):
         :return: The trace.
         """
         # If multiple threads are used, then create a new trace for each thread
-        thread_id = self.get_trace_identifier()
+        thread_id = get_thread_id()
         if thread_id not in self.traces:
             if not should_create:
                 return None