core/registrar.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os

from contextlib import asynccontextmanager
from typing import AsyncGenerator

from asgi_correlation_id import CorrelationIdMiddleware
from fastapi import Depends, FastAPI
from fastapi_limiter import FastAPILimiter
from fastapi_pagination import add_pagination
from starlette.middleware.authentication import AuthenticationMiddleware
from starlette.staticfiles import StaticFiles

from common.exception.exception_handler import register_exception
from common.log import set_custom_logfile, setup_logging
from core.conf import settings
from core.path_conf import STATIC_DIR, UPLOAD_DIR
from utils.health_check import ensure_unique_route_names, http_limit_callback
from utils.openapi import simplify_operation_ids
from utils.serializers import MsgSpecJSONResponse


# @asynccontextmanager
# async def register_init(app: FastAPI) -> AsyncGenerator[None, None]:
#     """

#     :param app: FastAPI 
#     :return:
#     """
#     await create_table()
#     # 连接 redis
#     await redis_client.open()
#     await FastAPILimiter.init(
#         redis=redis_client,
#         prefix=settings.REQUEST_LIMITER_REDIS_PREFIX,
#         http_callback=http_limit_callback,
#     )

#     yield

#     await redis_client.close()
#     await FastAPILimiter.close()


def register_app() -> FastAPI:
    """ FastAPI """
    app = FastAPI(
        title=settings.FASTAPI_TITLE,
        version=settings.FASTAPI_VERSION,
        description=settings.FASTAPI_DESCRIPTION,
        docs_url=settings.FASTAPI_DOCS_URL,
        redoc_url=settings.FASTAPI_REDOC_URL,
        openapi_url=settings.FASTAPI_OPENAPI_URL,
        default_response_class=MsgSpecJSONResponse,
    )

    register_logger()
    register_static_file(app)
    register_middleware(app)
    register_router(app)
    register_page(app)
    register_exception(app)

    return app


def register_logger() -> None:
    setup_logging()
    set_custom_logfile()


def register_static_file(app: FastAPI) -> None:
    """

    :param app: FastAPI
    :return:
    """
    if not os.path.exists(UPLOAD_DIR):
        os.makedirs(UPLOAD_DIR)
    app.mount('/static/upload', StaticFiles(directory=UPLOAD_DIR), name='upload')

    if settings.FASTAPI_STATIC_FILES:
        app.mount('/static', StaticFiles(directory=STATIC_DIR), name='static')


def register_middleware(app: FastAPI) -> None:
    """
    :param app: FastAPI
    :return:
    """
    # Opera log
    # app.add_middleware(OperaLogMiddleware)

   

    # # Access log
    # if settings.MIDDLEWARE_ACCESS:
    #     from middleware.access_middleware import AccessMiddleware

    #     app.add_middleware(AccessMiddleware)

    # Trace ID 
    app.add_middleware(CorrelationIdMiddleware, validator=False)

    # CORS
    if settings.MIDDLEWARE_CORS:
        from fastapi.middleware.cors import CORSMiddleware

        app.add_middleware(
            CORSMiddleware,
            allow_origins=settings.CORS_ALLOWED_ORIGINS,
            allow_credentials=True,
            allow_methods=['*'],
            allow_headers=['*'],
            expose_headers=settings.CORS_EXPOSE_HEADERS,
        )


def register_router(app: FastAPI) -> None:
    """

    :param app: FastAPI
    :return:
    """
    from app.router import router

    app.include_router(router)

    ensure_unique_route_names(app)
    simplify_operation_ids(app)


def register_page(app: FastAPI) -> None:
    """

    :param app: FastAPI 
    :return:
    """
    add_pagination(app)