Purushothaman Kumaravel
5 days ago
committed by
GitHub
GitHub
4 changed files with 239 additions and 2 deletions
@ -0,0 +1,149 @@ |
|||
import time |
|||
from typing import Any, Awaitable, Callable, List, Optional, Set, Union |
|||
|
|||
from fastapi import APIRouter, FastAPI, Request, Response |
|||
from fastapi.responses import JSONResponse |
|||
from fastapi.routing import APIRoute |
|||
|
|||
|
|||
async def health_check(request: Request): |
|||
""" |
|||
Health check endpoint |
|||
""" |
|||
return Response(content="OK", status_code=200) |
|||
|
|||
|
|||
class TimedRoute(APIRoute): |
|||
def get_route_handler(self) -> Callable: |
|||
original_route_handler = super().get_route_handler() |
|||
|
|||
async def custom_route_handler(request: Request) -> Response: |
|||
before = time.time() |
|||
response: Response = await original_route_handler(request) |
|||
duration = time.time() - before |
|||
response.headers["X-Response-Time"] = str(duration) |
|||
print(f"{self.name} route duration: {duration}") |
|||
print(f"{self.name} route response: {response}") |
|||
print(f"{self.name} route response headers: {response.headers}") |
|||
return response |
|||
|
|||
return custom_route_handler |
|||
|
|||
|
|||
class AppRouter(APIRouter): |
|||
def __init__(self, prefix="", name="Global", tags: list = None, **kwargs): |
|||
self.name = name |
|||
tags = tags or [] |
|||
tags.insert(0, name) |
|||
super().__init__(prefix=prefix, tags=tags, **kwargs) |
|||
self._parent: Optional[AppRouter] = None |
|||
self._add_health_check() |
|||
|
|||
@property |
|||
def request_name_prefix(self): |
|||
return ( |
|||
f"{self._parent.request_name_prefix}.{self.name}" |
|||
if self._parent |
|||
else self.name |
|||
) |
|||
|
|||
def _add_health_check(self): |
|||
""" |
|||
Adding default health check route for all new routers |
|||
""" |
|||
self.add_api_route( |
|||
"/healthz", endpoint=health_check, methods=["GET"], name="health-check" |
|||
) |
|||
|
|||
def include_router(self, router: "AppRouter", **kwargs): |
|||
""" |
|||
Include another router into this router. |
|||
""" |
|||
router._parent = self |
|||
super().include_router(router, **kwargs) |
|||
|
|||
def add_api_route( |
|||
self, |
|||
path: str, |
|||
endpoint: Callable[..., Any], |
|||
methods: Union[Set[str], List[str]], # noqa |
|||
name: str, |
|||
**kwargs, |
|||
): |
|||
name = f"{self.request_name_prefix}.{name}" |
|||
return super().add_api_route( |
|||
path, |
|||
endpoint, |
|||
methods=methods, |
|||
name=name, |
|||
**kwargs, |
|||
) |
|||
|
|||
def add_route( |
|||
self, |
|||
path: str, |
|||
endpoint: Callable[[Request], Union[Awaitable[Response], Response]], |
|||
methods: Union[List[str], None] = None, |
|||
name: Union[str, None] = None, |
|||
include_in_schema: bool = True, |
|||
) -> None: |
|||
name = f"{self.request_name_prefix}.{name}" |
|||
return super().add_route( |
|||
path, |
|||
endpoint, |
|||
methods=methods, |
|||
name=name, |
|||
include_in_schema=include_in_schema, |
|||
) |
|||
|
|||
|
|||
app = FastAPI(route_class=TimedRoute, router_class=AppRouter) |
|||
model = AppRouter(prefix="/model", name="Model", route_class=TimedRoute) |
|||
item = AppRouter(prefix="/{model_id}/item", name="Item", route_class=TimedRoute) |
|||
|
|||
|
|||
async def create_model(request: Request): |
|||
""" |
|||
Create a model |
|||
""" |
|||
print("Model created") |
|||
route: TimedRoute = request.scope["route"] |
|||
router: AppRouter = request.scope["router"] |
|||
return JSONResponse( |
|||
{ |
|||
"route_class": route.__class__.__name__, |
|||
"route_name": route.name, |
|||
"router_class": router.__class__.__name__, |
|||
}, |
|||
status_code=200, |
|||
) |
|||
|
|||
|
|||
model.add_api_route( |
|||
path="/create", endpoint=create_model, methods=["POST"], name="create-model" |
|||
) |
|||
|
|||
|
|||
async def create_item(request: Request): |
|||
""" |
|||
Create an item |
|||
""" |
|||
print("Item created") |
|||
route: TimedRoute = request.scope["route"] |
|||
router: AppRouter = request.scope["router"] |
|||
return JSONResponse( |
|||
{ |
|||
"route_class": route.__class__.__name__, |
|||
"route_name": route.name, |
|||
"router_class": router.__class__.__name__, |
|||
}, |
|||
status_code=200, |
|||
) |
|||
|
|||
|
|||
item.add_api_route( |
|||
path="/create", endpoint=create_item, methods=["POST"], name="create-item" |
|||
) |
|||
|
|||
model.include_router(item) |
|||
app.include_router(model) |
|||
@ -0,0 +1,33 @@ |
|||
from fastapi.testclient import TestClient |
|||
|
|||
from docs_src.custom_api_router.tutorial001 import app |
|||
|
|||
client = TestClient(app) |
|||
|
|||
|
|||
def test_get_timed(): |
|||
response = client.get("/healthz") |
|||
assert response.text == "OK" |
|||
assert "X-Response-Time" in response.headers |
|||
assert float(response.headers["X-Response-Time"]) >= 0 |
|||
|
|||
|
|||
def test_route_class(): |
|||
response = client.post( |
|||
"/model/create", json={"name": "test", "description": "test"} |
|||
) |
|||
assert response.status_code == 200 |
|||
response_json = response.json() |
|||
assert response_json["route_name"] == "Global.Model.create-model" |
|||
assert response_json["route_class"] == "TimedRoute" |
|||
assert response_json["router_class"] == "AppRouter" |
|||
|
|||
|
|||
def test_route_name(): |
|||
response = client.post( |
|||
"/model/Model001/item/create", json={"name": "test", "description": "test"} |
|||
) |
|||
assert response.status_code == 200 |
|||
response_json = response.json() |
|||
assert response_json["route_name"] == "Global.Model.Item.create-item" |
|||
assert response_json["router_class"] == "AppRouter" |
|||
Loading…
Reference in new issue