diff --git a/docs_src/route_middleware/tutorial001.py b/docs_src/route_middleware/tutorial001.py index 1dca98dae4..7e4881c91c 100644 --- a/docs_src/route_middleware/tutorial001.py +++ b/docs_src/route_middleware/tutorial001.py @@ -1,13 +1,27 @@ -from fastapi import FastAPI, Request -from fastapi.route_middleware import log_route, route_middleware, verify_jwt +from fastapi import FastAPI, Request, route_middleware app = FastAPI() +# You can define your own middlewares +async def verify_jwt(req: Request): + if req.query_params.get("token") != "secret": + from fastapi import HTTPException, status + + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail="Invalid Token" + ) + req.scope["user"] = "admin" + + +def log_route(req: Request): + print(f"Request to: {req.url.path}") + + @app.post("/secure") @route_middleware(verify_jwt, log_route) async def secure_route(req: Request, is_true: bool): - return {"status": "ok", "is_true": is_true, "user": req.user} + return {"status": "ok", "is_true": is_true, "user": req.scope.get("user")} @app.post("/open") diff --git a/fastapi/__init__.py b/fastapi/__init__.py index 1176856891..a13a24667a 100644 --- a/fastapi/__init__.py +++ b/fastapi/__init__.py @@ -20,6 +20,7 @@ from .param_functions import Query as Query from .param_functions import Security as Security from .requests import Request as Request from .responses import Response as Response +from .route_middleware import route_middleware as route_middleware from .routing import APIRouter as APIRouter from .websockets import WebSocket as WebSocket from .websockets import WebSocketDisconnect as WebSocketDisconnect diff --git a/fastapi/route_middleware.py b/fastapi/route_middleware.py index cba0f1260a..949d92234f 100644 --- a/fastapi/route_middleware.py +++ b/fastapi/route_middleware.py @@ -1,13 +1,11 @@ from functools import wraps -from typing import Callable +from typing import Any, Callable -from fastapi import Request - -def route_middleware(*middlewares: Callable): - def decorator(route_func: Callable): +def route_middleware(*middlewares: Callable[..., Any]) -> Callable[..., Any]: + def decorator(route_func: Callable[..., Any]) -> Callable[..., Any]: @wraps(route_func) - async def wrapper(*args, **kwargs): + async def wrapper(*args: Any, **kwargs: Any) -> Any: req = kwargs.get("req") if req is None: raise ValueError("Route must have 'request: Request' parameter") @@ -22,17 +20,3 @@ def route_middleware(*middlewares: Callable): return wrapper return decorator - - -# Example middlewares -async def verify_jwt(req: Request): - # just a mock - if not (req.query_params.get("is_true") == "true"): - req.user = {"name": "xyz", "admin": True} - from fastapi import HTTPException, status - - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid JWT") - - -def log_route(req: Request): - print(f"[LOG] Path: {req.url.path}") diff --git a/tests/test_route_middleware.py b/tests/test_route_middleware.py index 3df3d10d1a..89971475f4 100644 --- a/tests/test_route_middleware.py +++ b/tests/test_route_middleware.py @@ -1,14 +1,29 @@ +import pytest from fastapi import FastAPI, Request -from fastapi.route_middleware import log_route, route_middleware, verify_jwt +from fastapi.route_middleware import route_middleware from fastapi.testclient import TestClient + +# Example middlewares for testing +async def verify_jwt(req: Request): + if not (req.query_params.get("is_true") == "true"): + req.scope["user"] = {"name": "xyz", "admin": True} + from fastapi import HTTPException, status + + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid JWT") + + +def log_route(req: Request): + print(f"[LOG] Path: {req.url.path}") + + app = FastAPI() @app.post("/secure") @route_middleware(verify_jwt, log_route) async def secure_route(req: Request, is_true: bool): - return {"status": "ok", "is_true": is_true, "user": req.user} + return {"status": "ok", "is_true": is_true, "user": req.scope.get("user")} @app.post("/open") @@ -32,3 +47,19 @@ def test_secure_route_fail(): def test_open_route(): response = client.post("/open?is_true=false") assert response.status_code == 200 + + +def test_route_middleware_wrong_param_name(): + # If the param name is not 'req', it should fail + # We use a separate app to avoid polluting the main 'app' + local_app = FastAPI() + + @local_app.get("/wrong-param") + @route_middleware(log_route) + async def wrong_param_handler(request: Request): + return {"msg": "ok"} # pragma: no cover + + with pytest.raises( + ValueError, match="Route must have 'request: Request' parameter" + ): + TestClient(local_app).get("/wrong-param")