[flask]中间件

前言

中间件用于在请求到达视图函数之前或响应返回给客户端之前执行一些操作，比如身份验证、日志记录、错误处理等。

方式1，使用钩子函数

from flask import Flask, request, Response

def DemoMiddleWare(app: Flask):
    @app.before_request
    def before_request1():
        print("before_request")

    @app.after_request
    def after_request1(response: Response):
        print(f"after_request, method: {request.method}, path: {request.path}")
        return response
    
    return app

app = Flask(__name__)

@app.get("/")
def hello():
    return "Hello World!"

app = DemoMiddleWare(app)

if __name__ == "__main__":
    app.run(debug=False)

方式2，实现中间件类

from flask import Flask
from werkzeug.middleware.proxy_fix import ProxyFix
import time


class MyMiddleWare:
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        data = {
            "method": environ.get("REQUEST_METHOD", "NaN"),
            "scheme": environ.get("wsgi.url_scheme", "NaN"),
            "server_protocol": environ.get("SERVER_PROTOCOL", "NaN"),
            "remote_addr": environ.get("REMOTE_ADDR", "NaN"),
            "host": environ.get("HTTP_HOST", "NaN"),
            "path": environ.get("PATH_INFO", "NaN"),
            "user_agent": environ.get("HTTP_USER_AGENT", "NaN"),
        }

        # 自定义请求头都以 HTTP_* 开头
        if header_ha := environ.get("HTTP_HA", ""):
            data["header_ha"] = header_ha

        # 处理 cookie
        if cookie := environ.get("HTTP_COOKIE", ""):
            data["cookie"] = cookie

        start_time = time.time()

        response = self.app(environ, start_response)
        data["response_time"] = round((time.time() - start_time), 4)
        print(data)

        return response


class My2MiddleWare:
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        print("This is middleware2, before request")
        response = self.app(environ, start_response)
        print("This is middleware2, after request")

        return response


app = Flask(__name__)
# ProxyFix是werkzeug提供的一个中间件，当flask应用部署在反向代理器后面时，可以将反向代理器设置得X-Forwarded-*头信息
# 具体ProxyFix处理了哪些代理请求头可以参考ProxyFix的docstring
app.wsgi_app = ProxyFix(app.wsgi_app)
app.wsgi_app = MyMiddleWare(app.wsgi_app)
app.wsgi_app = My2MiddleWare(app.wsgi_app)  # 添加第2个中间件


@app.get("/")
def hello():
    return "Hello World!"


@app.get("/2")
def hello2():
    return "Hello World!"


if __name__ == "__main__":
    app.run(debug=False)

上面示例代码只处理了请求头等信息，有时候可能还需要处理响应，比如获取响应码、设置响应头等。

from flask import Flask
import time


class MyMiddleWare:
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        data = {
            "method": environ.get("REQUEST_METHOD", "NaN"),
            "scheme": environ.get("wsgi.url_scheme", "NaN"),
            "server_protocol": environ.get("SERVER_PROTOCOL", "NaN"),
            "remote_addr": environ.get("REMOTE_ADDR", "NaN"),
            "host": environ.get("HTTP_HOST", "NaN"),
            "path": environ.get("PATH_INFO", "NaN"),
            "user_agent": environ.get("HTTP_USER_AGENT", "NaN"),
        }

        # 自定义请求头都以 HTTP_* 开头
        if header_ha := environ.get("HTTP_HA", ""):
            data["header_ha"] = header_ha

        # 处理 cookie
        if cookie := environ.get("HTTP_COOKIE", ""):
            data["cookie"] = cookie

        resp_status_code = None
        resp_headers = []

        def catching_start_response(status, headers, exc_info=None):
            nonlocal resp_status_code, resp_headers
            resp_status_code = status
            headers.append(("X-Resp1", "OK"))  # 添加自定义响应头
            resp_headers = headers
            return start_response(status, headers, exc_info)

        start_time = time.time()
        response = self.app(environ, catching_start_response)
        data["status"] = resp_status_code
        data["response_time"] = round((time.time() - start_time), 4)
        print(data)
        print(f"response headers: {resp_headers}")

        return response


app = Flask(__name__)
app.wsgi_app = MyMiddleWare(app.wsgi_app)


@app.get("/")
def hello():
    return "Hello World!"


@app.get("/2")
def hello2():
    return "Hello World!"


if __name__ == "__main__":
    app.run(debug=False)

小结

两种方式实现中间件都行，不过用钩子函数会更方便点。如果想水代码行数，用中间件类也是不错的选择。

参考

• Flask入门（九）：Flask中间件
• Flask官方文档

来源链接：https://www.cnblogs.com/XY-Heruo/p/18859632