fix(fastapi): custom responses

Author: maorlx

epsagon/events/dbapi.py

@@ -72,21 +72,14 @@ def __init__(
 
         # in case of pg instrumentation we extract data from the dsn property
         if hasattr(connection, 'dsn'):
-            print_debug('Parsing using dsn')
             dsn = parse_dsn(connection.dsn)
-            print_debug('Parsed dsn: {}'.format(dsn))
             db_name = dsn.get('dbname', '')
             host = dsn.get('host', 'local')
             query = cursor.query
         else:
             query = _args[0]
             host = connection.extract_hostname
             db_name = connection.extract_dbname
-        print_debug(
-            'Extracted db name - {}, host - {}, query - {}'.format(
-                db_name, host, query
-            )
-        )
 
         self.resource['name'] = db_name if db_name else host
 
@@ -104,11 +97,6 @@ def __init__(
             host,
             self.RESOURCE_TYPE
         )
-        print_debug(
-            '{}: resource name - {}, operation - {}'.format(
-                self.event_id, self.resource['name'], self.resource['operation']
-            )
-        )
         self.resource['metadata'] = {
             'Host': host,
             'Driver': connection.__class__.__module__.split('.')[0],
@@ -121,11 +109,6 @@ def __init__(
                 (not trace_factory.metadata_only)
         ):
             self.resource['metadata']['Query'] = query[:MAX_QUERY_SIZE]
-            print_debug(
-                '{}: collected query {}'.format(
-                    self.event_id, self.resource['metadata']['Query']
-                    )
-            )
 
         if exception is None:
             # Update response data
@@ -174,7 +157,6 @@ def create_event(wrapped, cursor_wrapper, args, kwargs, start_time,
         :param exception:
         :return:
         """
-        print_debug('Creating DBAPI event')
         event = DBAPIEvent(
             cursor_wrapper.connection_wrapper,
             cursor_wrapper,
@@ -183,5 +165,4 @@ def create_event(wrapped, cursor_wrapper, args, kwargs, start_time,
             start_time,
             exception,
         )
-        print_debug('Adding DBAPI event to trace')
         trace_factory.add_event(event)

epsagon/runners/fastapi.py

@@ -7,6 +7,7 @@
 import uuid
 import warnings
 from fastapi.responses import (
+    Response,
     JSONResponse,
     HTMLResponse,
     PlainTextResponse,
@@ -107,16 +108,16 @@ def update_response(self, response):
                     'Response Data',
                     body
                 )
+        if isinstance(response, Response):
+            response_headers = dict(response.headers.items())
+            if response.headers:
+                add_data_if_needed(
+                    self.resource['metadata'],
+                    'Response Headers',
+                    response_headers
+                )
 
-        response_headers = dict(response.headers.items())
-        if response.headers:
-            add_data_if_needed(
-                self.resource['metadata'],
-                'Response Headers',
-                response_headers
-            )
-
-        self.resource['metadata']['status_code'] = response.status_code
+            self.resource['metadata']['status_code'] = response.status_code
 
-        if response.status_code >= 500:
-            self.set_error()
+            if response.status_code >= 500:
+                self.set_error()

tests/wrappers/test_fastapi_wrapper.py

@@ -5,6 +5,7 @@
 import pytest
 import asynctest
 import asyncio
+from typing import List
 from httpx import AsyncClient
 from fastapi import FastAPI, APIRouter, Request
 from fastapi.responses import JSONResponse
@@ -22,6 +23,8 @@
 MULTIPLE_THREADS_ROUTE = f'/{MULTIPLE_THREADS_KEY}'
 MULTIPLE_THREADS_RETURN_VALUE = MULTIPLE_THREADS_KEY
 TEST_POST_DATA = {'post_test': '123'}
+CUSTOM_RESPONSE = ["A"]
+CUSTOM_RESPONSE_PATH = "/custom_response"
 
 def _get_response_data(key):
     return {key: key}
@@ -33,6 +36,9 @@ def _get_response(key):
 def handle():
     return _get_response(RETURN_VALUE)
 
+def handle_custom_response(response_model=List[str]):
+    return CUSTOM_RESPONSE
+
 def handle_given_request(request: Request):
     assert request.method == 'POST'
     loop = None
@@ -70,6 +76,7 @@ def handle_error():
 def fastapi_app():
     app = FastAPI()
     app.add_api_route("/", handle, methods=["GET"])
+    app.add_api_route(CUSTOM_RESPONSE_PATH, handle_custom_response, methods=["GET"])
     app.add_api_route(REQUEST_OBJ_PATH, handle_given_request, methods=["POST"])
     app.add_api_route("/a", handle_a, methods=["GET"])
     app.add_api_route("/b", handle_b, methods=["GET"])
@@ -100,6 +107,23 @@ async def test_fastapi_sanity(trace_transport, fastapi_app):
     assert not trace_factory.traces
 
 
+@pytest.mark.asyncio
+async def test_fastapi_custom_response(trace_transport, fastapi_app):
+    """Sanity test."""
+    request_path = f'{CUSTOM_RESPONSE_PATH}?x=testval'
+    async with AsyncClient(app=fastapi_app, base_url="http://test") as ac:
+        response = await ac.get(request_path)
+    response_data = response.json()
+    runner = trace_transport.last_trace.events[0]
+    assert isinstance(runner, FastapiRunner)
+    assert runner.resource['name'].startswith('127.0.0.1')
+    assert runner.resource['metadata']['Path'] == CUSTOM_RESPONSE_PATH
+    assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    assert response_data == CUSTOM_RESPONSE
+    # validating no `zombie` traces exist
+    assert not trace_factory.traces
+
+
 @pytest.mark.asyncio
 async def test_fastapi_given_request(trace_transport, fastapi_app):
     """Sanity test."""