使用Pait快速构建gRPC Json网关

from flask import Response, request
# 生成的Python文件都存放在gRPC公有包里面
from grpc_example_common.protos.user import user_pb2 as user_message

from app_service.utils import g, get_uid_by_token, make_response


def create_user() -> Response:
    request_dict: dict = request.json
    g.user_grpc_service.create_user(
        uid=request_dict["uid"], user_name=request_dict["user_name"], password=request_dict["password"]
    )
    return make_response()


def delete_user() -> Response:
    request_dict: dict = request.json
    g.user_grpc_service.delete_user(uid=request_dict["uid"])
    return make_response()


def login_route() -> Response:
    request_dict: dict = request.json
    login_result: user_message.LoginUserResult = g.user_grpc_service.login_user(
        uid=request_dict["uid"], password=request_dict["password"]
    )
    return make_response({"token": login_result.token, "uid": login_result.uid})


def logout_route() -> Response:
    request_dict: dict = request.json
    if get_uid_by_token() == request_dict["uid"]:
        token: str = request.headers.get("token", "")
        g.user_grpc_service.logout_user(uid=request_dict["uid"], token=token)
        return make_response()
    else:
        raise RuntimeError("Uid ERROR")

from flask.blueprints import Blueprint

from app_service import user_route

user_bp: Blueprint = Blueprint("user_bp", __name__, url_prefix="/api/user")
user_bp.add_url_rule("/create", view_func=user_route.create_user, methods=["POST"])
user_bp.add_url_rule("/delete", view_func=user_route.delete_user, methods=["POST"])
user_bp.add_url_rule("/login", view_func=user_route.login_route, methods=["POST"])
user_bp.add_url_rule("/logout", view_func=user_route.logout_route, methods=["POST"])

from typing import Type
import grpc

from flask import Flask, jsonify, Response
from pydantic import BaseModel, Field

from pait.app.flask.grpc_route import GrpcGatewayRoute
from pait.app import set_app_attribute
from pait.model.response import PaitBaseResponseModel, PaitJsonResponseModel
from pait.util.grpc_inspect.stub import GrpcModel
from protobuf_to_pydantic import msg_to_pydantic_model

from grpc_example_common.protos.user import user_pb2_grpc


def gen_response_model_handle(grpc_model: GrpcModel) -> Type[PaitBaseResponseModel]:
    # 动态生成对应的响应模型
    class CustomerJsonResponseModel(PaitJsonResponseModel):
        class CustomerJsonResponseRespModel(BaseModel):
            code: int = Field(0, description="api code")
            msg: str = Field("success", description="api status msg")
            data: msg_to_pydantic_model(grpc_model.response) = Field(description="api response data")  # type: ignore

        name: str = grpc_model.response.DESCRIPTOR.name
        response_data: Type[BaseModel] = CustomerJsonResponseRespModel

    return CustomerJsonResponseModel


def add_grpc_gateway_route(app: Flask) -> None:

    def _make_response(resp_dict: dict) -> Response:
        return jsonify({"code": 0, "msg": "", "data": resp_dict})
    
    # 定义一个网关路由，该路由已经把gRPC的方法映射到指定的Flask应用中
    grpc_gateway_route: GrpcGatewayRoute = GrpcGatewayRoute(
        app,
        # 通过User.proto生成的Python代码中的`UserStub`类,Pait可以通过该类解析出User.proto对应的方法
        user_pb2_grpc.UserStub,
        # 指定对应路由的前缀
        prefix="/api/gateway",
        title="UserGrpc",
        gen_response_model_handle=gen_response_model_handle,  # 定义返回响应的OpenAPI模型，方便生成OpenAPI文档
        make_response=_make_response,   # 定义返回的响应格式
    )
    # 初始化gRPC的channel并传入`grpc_gateway_route`中
    grpc_gateway_route.init_channel(grpc.intercept_channel(grpc.insecure_channel("0.0.0.0:9001")))

class CustomerGrpcGatewayRoute(GrpcGatewayRoute):
    def gen_route(self, grpc_model: GrpcModel, request_pydantic_model_class: Type[BaseModel]) -> Callable:

        # 对于登录和创建用户接口不需要Token，所以采用原生的方法
        if grpc_model.method in ("/user.User/login_user", "/user.User/create_user"):
            return super().gen_route(grpc_model, request_pydantic_model_class)
        else:
            def _route(
                # 通过对应的grpc message生成的pydantic.BaseModel对象
                request_pydantic_model: request_pydantic_model_class,  # type: ignore
                # 添加token参数
                token: str = Header.i(description="User Token"),
                # 添加请求id参数，与需求无关
                req_id: str = Header.i(alias="X-Request-Id", default_factory=lambda: str(uuid4())),
            ) -> Any:
                func: Callable = self.get_grpc_func(grpc_model.method)
                request_dict: dict = request_pydantic_model.dict()  # type: ignore
                if grpc_model.method == "/user.User/logout_user":
                    # 退出登陆接口需要token
                    request_dict["token"] = token
                else:
                    # 其他接口不需要token参数，只需要校验token是否合法，如果不合法则抛出异常
                    result: user_pb2.GetUidByTokenResult = user_pb2_grpc.UserStub(self.channel).get_uid_by_token(
                        user_pb2.GetUidByTokenRequest(token=token)
                    )
                    if not result.uid:
                        raise RuntimeError(f"Not found user by token:{token}")
                # 生成调用参数后再调用gRPC方法，并把gRPC方法返回的数据返回给调用端
                request_msg: Message = self.get_msg_from_dict(grpc_model.request, request_dict)
                grpc_msg: Message = func(request_msg, metadata=[("req_id", req_id)])
                return self._make_response(self.get_dict_from_msg(grpc_msg))

            return _route

def create_app() -> Flask:
    app: Flask = Flask(__name__)
    app.json_encoder = CustomJSONEncoder
    app.register_blueprint(manager_book_bp)
    app.register_blueprint(social_book_bp)
    app.register_blueprint(user_bp)

    book_grpc_service: BookGrpcService = BookGrpcService("0.0.0.0", 9000)
    book_grpc_service.channel_ready_future(timeout=3)
    user_grpc_service: UserGrpcService = UserGrpcService("0.0.0.0", 9001)
    user_grpc_service.channel_ready_future(timeout=3)
    ContextMiddleware(app=app, book_grpc_service=book_grpc_service, user_grpc_service=user_grpc_service)

    # 映射gRPC服务对应的接口到app
    add_grpc_gateway_route(app)
    # 添加一个OpenAPI文档路由，从而可以查看被Pait托管路由的接口文档
    add_doc_route(app)
    app.errorhandler(Exception)(api_exception)
    return app


if __name__ == "__main__":
    # 屏蔽接口文档中`OPTIONS`和`HEAD`方法接口的展示
    config.init_config(apply_func_list=[apply_block_http_method_set({"OPTIONS", "HEAD"})])
    create_app().run("localhost", port=8000)