Python FastAPI series: Custom FastAPI middleware
-
- FastAPI middleware middleware execution logic
- Create FastAPI middleware middleware
-
- Create middleware using decorators
- Create middleware by inheriting BaseHTTPMiddleware
- Create middleware according to ASGI specifications
In some cases, we need to perform some common functions on all or part of the routes of the entire FastAPI application, such as authentication, logging, error handling, etc., which we can do through custom FastAPI middleware. FastAPI also comes with some commonly used middleware to complete request protocol restrictions, cross-domain submission, etc.
Generally speaking, when encountering the following demand scenarios, you can consider using FastAPI middleware to implement:
1. Authentication: Verify the identity of the request, such as checking a JWT token or using OAuth2 for verification.
2. Logging: Record request and response logs, including request method, URL, response status code and other information.
3. Error handling: Handle exceptions in the application, such as catching exceptions and returning custom error responses.
4. Request processing: Process the request, such as parsing request parameters, verifying request data, etc.
5. Caching: You can use middleware to implement caching functions. For example, in the middleware, you can check whether the requested response exists in the cache, and if so, return the cached response directly.
FastAPI middleware middleware execution logic
FastAPI middleware works between each Request request and Response. It can modify Request and Response. The detailed process is as follows:
1. Receive the Request request from the client;
2. Customize the operation for this Request request;
3. Then the Request request is sent back to the original route, and the business logic defined in the route continues to process the Request request;
4. After the original route business is processed, the FastAPI middleware will obtain the Response result generated by the route. At this time, you can customize the operation for the Response response result;
5. Finally, send the Response result to the client;
Create FastAPI middleware
Use decorators to create middleware
In FastAPI applications, you can use the app.middleware(“http”) decorator to create FastAPI middleware. In the following example, we record a pair of start and end times as the middleware time when the program enters the middleware, and record a pair of start and end times as the router time when the program enters the routing business logic. The router start time is the middleware that obtains the middleware and writes the request. Add 1 hour to the time, and the detailed code is as follows:
main.py
import time
from datetime import datetime, timedelta
import uvicorn as uvicorn
from fastapi import FastAPI
from starlette.requests import Request
from starlette.responses import Response
app = FastAPI()
# Format time into string
def _time2str(time_str):
return datetime.strftime(time_str, '%Y-%m-%d %H:%M:%S')
#Convert string to time
def _str2time(time_str):
return datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S')
#Define middleware
@app.middleware("http")
async def process_time_middleware(request: Request, call_next):
# Receive Request requests from the client;
headers = dict(request.scope['headers'])
#Define middleware start time
middleware_start_time = _time2str(datetime.now())
# Add the middleware start time to the headers of the request. Here request.headers is a readable and writable object, but its value is immutable, so here you need to convert request.headers into a dictionary, and then modify the dictionary. value, and finally convert the dictionary into a tuple and assign it to request.scope['headers'];
headers[b'middleware_start_time'] = middleware_start_time.encode('utf-8')
request.scope['headers'] = [(k, v) for k, v in headers.items()]
# Pass the Request request back to the original route
response = await call_next(request)
# In order to better observe the execution process of middleware, let middleware sleep for 1 second.
time.sleep(1)
# Receive the Response from the original route and add the middleware end time to the response headers.
response.headers["middleware_start_time"] = middleware_start_time
response.headers["middleware_end_time"] = _time2str(datetime.now())
return response
@app.get("/")
async def index(request: Request, response: Response):
# Get the middleware start time passed by the middleware through the request in the route
middleware_start_time = _str2time(request.headers.get("middleware_start_time"))
# Add 1 hour to middleware_start_time as the router start time, and add 2 hours as the router end time.
router_start_time = _str2time(middleware_start_time) + timedelta(hours=1)
router_end_time = _str2time(request.headers.get("middleware_start_time")) + timedelta(hours=2)
# Add the router start time and router end time to the response headers
response.headers["router_start_time"] = _time2str(router_start_time)
response.headers["router_end_time"] = _time2str(router_end_time)
return "test middleware"
if __name__ == '__main__':
uvicorn.run(app="main:app", port=8088, reload=True)
The picture below shows the execution result of FastAPI middleware:
Create middleware by inheriting BaseHTTPMiddleware
In actual projects, multiple FastAPI middlewares are often defined to achieve complete business requirements. In this case, we can create custom middlewares by inheriting the starlette.middleware.base.Middleware class. The following is the implementation of the above requirements by inheriting BaseHTTPMiddleware:
The project file can be split into main.py (main program), process_time_middleware.py (custom middleware class), utils.py (tool function)
main.py
from starlette.responses import Response
from fapi.process_time_middleware import ProcessTimeMiddleware
from fapi.utils import _str2time, _time2str
app = FastAPI()
# Add middleware to the main program
app.add_middleware(ProcessTimeMiddleware, header_namespace="middleware")
@app.get("/")
async def index(request: Request, response: Response):
# Get the middleware start time passed by the middleware through the request in the route
middleware_start_time = _str2time(request.headers.get("middleware_start_time"))
# Add 1 hour to middleware_start_time as the router start time, and add 2 hours as the router end time.
router_start_time = middleware_start_time + timedelta(hours=1)
router_end_time = middleware_start_time + timedelta(hours=2)
# Add the router start time and router end time to the response headers
response.headers["router_start_time"] = _time2str(router_start_time)
response.headers["router_end_time"] = _time2str(router_end_time)
return "test middleware"
if __name__ == '__main__':
uvicorn.run(app="main:app", port=8088, reload=True)
process_time_middleware.py
import time
from datetime import datetime
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from fapi.utils import _time2str
class ProcessTimeMiddleware(BaseHTTPMiddleware):
def __init__(self, app, header_namespace: str):
super().__init__(app)
# Custom parameters used to define the header namespace of middleware
self.header_namespace = header_namespace
async def dispatch(self, request: Request, call_next):
# Receive Request requests from the client;
headers = dict(request.scope['headers'])
#Define middleware start time
middleware_start_time = _time2str(datetime.now())
# Add the middleware start time to the headers of the request. Here request.headers is a readable and writable object, but its value is immutable, so here you need to convert request.headers into a dictionary, and then modify the dictionary. value, and finally convert the dictionary into a tuple and assign it to request.scope['headers'];
headers[b'middleware_start_time'] = middleware_start_time.encode('utf-8')
request.scope['headers'] = [(k, v) for k, v in headers.items()]
# Pass the Request request back to the original route
response = await call_next(request)
# In order to better observe the execution process of middleware, let middleware sleep for 1 second.
time.sleep(1)
# Receive the Response from the original route and add the middleware end time to the response headers.
response.headers[f"{<!-- -->self.header_namespace}_start_time"] = middleware_start_time
response.headers[f"{<!-- -->self.header_namespace}_end_time"] = _time2str(datetime.now())
return response
utils.py
from datetime import datetime
# Format time into string
def _time2str(time_str):
return datetime.strftime(time_str, '%Y-%m-%d %H:%M:%S')
#Convert string to time
def _str2time(time_str):
return datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S')
Create middleware according to ASGI specifications
Middleware created according to the ASGI specification can obtain lower-level functionality and enhance interoperability across frameworks and servers. The following code will achieve the needs in the above example by creating a pure ASGI class.
The project file can still be split into main.py (main program), process_time_asgi_middleware.py (custom ASGI middleware class), utils.py (tool function), where main.py and utils.py are consistent with the above:
main.py
from datetime import timedelta
import uvicorn as uvicorn
from fastapi import FastAPI
from starlette.requests import Request
from starlette.responses import Response
from fapi.process_time_asgi_middleware import ProcessTimeASGIMiddleware
from fapi.utils import _str2time, _time2str
app = FastAPI()
# Add ASGI middleware to the main program
app.add_middleware(ProcessTimeASGIMiddleware, header_namespace="middleware")
@app.get("/")
async def index(request: Request, response: Response):
# Get the middleware start time passed by the middleware through the request in the route
middleware_start_time = _str2time(request.headers.get("middleware_start_time"))
# Add 1 hour to middleware_start_time as the router start time, and add 2 hours as the router end time.
router_start_time = middleware_start_time + timedelta(hours=1)
router_end_time = middleware_start_time + timedelta(hours=2)
# Add the router start time and router end time to the response headers
response.headers["router_start_time"] = _time2str(router_start_time)
response.headers["router_end_time"] = _time2str(router_end_time)
return "test middleware"
if __name__ == '__main__':
uvicorn.run(app="main:app", port=8088, reload=True)
process_time_asgi_middleware.py
import time
from datetime import datetime
from fastapi import Request
from starlette.datastructures import MutableHeaders
from fapi.utils import _time2str
class ProcessTimeASGIMiddleware:
def __init__(self, app, header_namespace: str):
self.app = app
# Custom parameters used to define the header namespace of middleware
self.header_namespace = header_namespace
#ASGI middleware must be a callable object that accepts three parameters, namely scope, receive, and send;
async def __call__(self, scope, receive, send):
request = Request(scope)
# Receive Request requests from the client;
headers = dict(request.scope['headers'])
#Define middleware start time
middleware_start_time = _time2str(datetime.now())
headers[b'middleware_start_time'] = middleware_start_time.encode('utf-8')
request.scope['headers'] = [(k, v) for k, v in headers.items()]
# Define the Send function to add the middleware start time and middleware end time to the response headers
async def add_headers(message):
if message["type"] == "http.response.start":
new_headers = MutableHeaders(scope=message)
new_headers.append(f"{<!-- -->self.header_namespace}_start_time", middleware_start_time)
new_headers.append(f"{<!-- -->self.header_namespace}_end_time", _time2str(datetime.now()))
await send(message)
# Pass scope, receive, add_headers to the original ASGI application
return await self.app(scope, receive, add_headers)
utils.py
from datetime import datetime
# Format time into string
def _time2str(time_str):
return datetime.strftime(time_str, '%Y-%m-%d %H:%M:%S')
#Convert string to time
def _str2time(time_str):
return datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S')
In the code, ASGI middleware must be a callable object that accepts three parameters, namely scope, receive, and send. in
1. Scope is a dictionary that saves information about connections, where the type of scope[“type”] can be:
“http”: used for HTTP requests.
“websocket”: for WebSocket connections.
“lifespan”: used for ASGI lifecycle messages.
2. receive is used to exchange ASGI event messages with the ASGI server. The type and content of these messages depends on the scope type.
Of course, you can also use functions instead of pure ASGI middleware classes:
import functools
def asgi_middleware():
def asgi_decorator(app):
@functools.wraps(app)
async def wrapped_app(scope, receive, send):
await app(scope, receive, send)
return wrapped_app
return asgi_decorator
In summary, using middleware in FastAPI can provide more powerful functions and make the program more elegant and readable!