fix(fastapi): request body extraction (#324)

Author: maorlx

epsagon/runners/fastapi.py

@@ -35,12 +35,11 @@ class FastapiRunner(BaseEvent):
     RESOURCE_TYPE = 'fastapi'
     OPERATION = 'request'
 
-    def __init__(self, start_time, request, body):
+    def __init__(self, start_time, request):
         """
         Initialize.
         :param start_time: event's start time (epoch).
         :param request: the incoming request.
-        :param body: the body of request.
         """
 
         super(FastapiRunner, self).__init__(start_time)
@@ -62,13 +61,6 @@ def __init__(self, start_time, request, body):
                 query_params.items()
             )
 
-        if body:
-            add_data_if_needed(
-                self.resource['metadata'],
-                'Request Data',
-                body
-            )
-
         request_headers = dict(request.headers.items())
         if request_headers.get(EPSAGON_HEADER):
             self.resource['metadata']['http_trace_id'] = request_headers.get(
@@ -81,6 +73,17 @@ def __init__(self, start_time, request, body):
                 request_headers
             )
 
+    def update_request_body(self, body):
+        """
+        Adds request body to event
+        """
+        if body:
+            add_data_if_needed(
+                self.resource['metadata'],
+                'Request Data',
+                body
+            )
+
     def update_response(self, response):
         """
         Adds response data to event.

epsagon/wrappers/fastapi.py

@@ -9,6 +9,7 @@
 import warnings
 from fastapi.routing import APIRoute
 from fastapi import Request, Response
+from starlette.requests import ClientDisconnect
 
 import epsagon.trace
 from epsagon.runners.fastapi import FastapiRunner
@@ -52,11 +53,7 @@ async def custom_route_handler(request: Request) -> Response:
             runner = None
             response = None
             try:
-                body = await request.json()
-            except json.decoder.JSONDecodeError:
-                body = ''
-            try:
-                runner = FastapiRunner(time.time(), request, json.dumps(body))
+                runner = FastapiRunner(time.time(), request)
                 trace.set_runner(runner)
                 collect_container_metadata(runner.resource['metadata'])
             except Exception as exception: # pylint: disable=W0703
@@ -68,7 +65,20 @@ async def custom_route_handler(request: Request) -> Response:
                 raised_err = exception
                 traceback_data = get_traceback_data_from_exception(exception)
                 trace.runner.set_exception(exception, traceback_data)
-
+            try:
+                trace.runner.update_request_body(
+                    json.dumps(await request.json())
+                )
+            except json.decoder.JSONDecodeError:
+                print_debug('Could not JSON-decode request body')
+            except ClientDisconnect:
+                print_debug(
+                    'Could not extract request body - client is disconnected'
+                )
+            except Exception as exception: # pylint: disable=W0703
+                print_debug(
+                    'Could not extract request body: {}'.format(exception)
+                )
             try:
                 if not raised_err and response is not None and runner:
                     if ignore_request(