feat(fastapi): Adding custom api routes support (#347)

epsagon/modules/fastapi.py

@@ -4,35 +4,13 @@
 
 from __future__ import absolute_import
 import wrapt
-from fastapi.routing import APIRoute
 from ..wrappers.fastapi import (
-    TracingAPIRoute,
     exception_handler_wrapper,
     server_call_wrapper,
+    route_class_wrapper,
 )
-from ..utils import is_lambda_env, print_debug
+from ..utils import is_lambda_env
 
-def _wrapper(wrapped, _instance, args, kwargs):
-    """
-    Adds TracingRoute into APIRouter (FastAPI).
-    :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)
-    route_class = kwargs.get('route_class', APIRoute)
-    if route_class != APIRoute:
-        # custom routes are not supported
-        print_debug(
-            f'Custom FastAPI route {route_class.__name__} is not supported'
-        )
-        return wrapped(*args, **kwargs)
-    kwargs['route_class'] = TracingAPIRoute
-    return wrapped(*args, **kwargs)
 
 def _exception_handler_wrapper(wrapped, _instance, args, kwargs):
     """
@@ -58,9 +36,9 @@ def patch():
     :return: None
     """
     wrapt.wrap_function_wrapper(
-        'fastapi',
-        'APIRouter.__init__',
-        _wrapper
+        'fastapi.routing',
+        'APIRoute.__init__',
+        route_class_wrapper
     )
     wrapt.wrap_function_wrapper(
         'starlette.applications',

epsagon/wrappers/fastapi.py

@@ -8,7 +8,6 @@
 import asyncio
 
 import warnings
-from fastapi.routing import APIRoute
 from fastapi import Request, Response
 from starlette.requests import ClientDisconnect
 from starlette.concurrency import run_in_threadpool
@@ -248,21 +247,27 @@ def wrapped_handler(*args, **kwargs):
     dependant.call = wrapped_handler
 
 
-class TracingAPIRoute(APIRoute):
+def route_class_wrapper(wrapped, instance, args, kwargs):
     """
-    Custom tracing route - traces each route request & response
+    Route class wrapper - traces each route request & response.
+    :param wrapped: wrapt's wrapped
+    :param instance: wrapt's instance
+    :param args: wrapt's args
+    :param kwargs: wrapt's kwargs
     """
+    result = wrapped(*args, **kwargs)
+    # Skip on Lambda environment since it's not relevant and might be duplicate
+    if not is_lambda_env() and instance:
+        try:
+            if instance.dependant and instance.dependant.call:
+                _wrap_handler(
+                    instance.dependant,
+                    kwargs.get('status_code', DEFAULT_SUCCESS_STATUS_CODE)
+                )
+        except Exception: # pylint: disable=broad-except
+            pass
 
-    def __init__(self, *args, **kwargs):
-        """
-        wraps the route endpoint with Epsagon wrapper
-        """
-        super().__init__(*args, **kwargs)
-        if self.dependant and self.dependant.call:
-            _wrap_handler(
-                self.dependant,
-                kwargs.pop('status_code', DEFAULT_SUCCESS_STATUS_CODE)
-            )
+    return result
 
 
 def exception_handler_wrapper(original_handler):

tests/wrappers/test_fastapi_wrapper.py

@@ -9,6 +9,7 @@
 from httpx import AsyncClient
 from pydantic import BaseModel
 from fastapi import FastAPI, APIRouter, Request
+from fastapi.routing import APIRoute
 from fastapi.responses import JSONResponse
 from fastapi.encoders import jsonable_encoder
 from epsagon import trace_factory
@@ -22,9 +23,12 @@
 
 RETURN_VALUE = 'testresponsedata'
 ROUTER_RETURN_VALUE = 'router-endpoint-return-data'
+CUSTOM_ROUTE_RETURN_VALUE = 'custom-route-endpoint-return-data'
 REQUEST_OBJ_PATH = '/given_request'
 TEST_ROUTER_PREFIX = '/test-router-path'
 TEST_ROUTER_PATH = '/test-router'
+TEST_CUSTOM_ROUTE_PATH = '/atest-custom-route'
+TEST_CUSTOM_ROUTE_PREFIX = '/test-custom-route-main'
 MULTIPLE_THREADS_KEY = "multiple_threads"
 MULTIPLE_THREADS_ROUTE = f'/{MULTIPLE_THREADS_KEY}'
 MULTIPLE_THREADS_RETURN_VALUE = MULTIPLE_THREADS_KEY
@@ -44,6 +48,9 @@
 class CustomBaseModel(BaseModel):
     data: List[str]
 
+class CustomRouteClass(APIRoute):
+    pass
+
 def _get_response_data(key):
     return {key: key}
 
@@ -87,6 +94,9 @@ def handle_b():
 def handle_router_endpoint():
     return _get_response(ROUTER_RETURN_VALUE)
 
+def handle_custom_route_endpoint():
+    return _get_response(CUSTOM_ROUTE_RETURN_VALUE)
+
 def multiple_threads_route():
     multiple_threads_handler()
     return _get_response(MULTIPLE_THREADS_RETURN_VALUE)
@@ -162,6 +172,9 @@ def _build_fastapi_app():
     router = APIRouter()
     router.add_api_route(TEST_ROUTER_PATH, handle_router_endpoint)
     app.include_router(router, prefix=TEST_ROUTER_PREFIX)
+    router_with_custom_route = APIRouter(route_class=CustomRouteClass)
+    router_with_custom_route.add_api_route(TEST_CUSTOM_ROUTE_PATH, handle_custom_route_endpoint)
+    app.include_router(router_with_custom_route, prefix=TEST_CUSTOM_ROUTE_PREFIX)
     return app
 
 @pytest.fixture(scope='function', autouse=False)
@@ -343,6 +356,25 @@ async def test_fastapi_custom_router(trace_transport, fastapi_app):
     assert response_data == expected_response_data
 
 
+@pytest.mark.asyncio
+async def test_fastapi_custom_api_route(trace_transport, fastapi_app):
+    """Custom api route sanity test."""
+    full_route_path= f'{TEST_CUSTOM_ROUTE_PREFIX}{TEST_CUSTOM_ROUTE_PATH}'
+    async with AsyncClient(app=fastapi_app, base_url="http://test") as ac:
+        response = await ac.get(full_route_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'] == full_route_path
+    assert runner.resource['metadata']['status_code'] == DEFAULT_SUCCESS_STATUS_CODE
+    expected_response_data = _get_response_data(CUSTOM_ROUTE_RETURN_VALUE)
+    assert runner.resource['metadata']['Response Data'] == (
+        expected_response_data
+    )
+    assert response_data == expected_response_data
+
+
 @pytest.mark.asyncio
 async def test_fastapi_exception(trace_transport, fastapi_app):
     """Test when the handler raises an exception."""