coze_znjd/main.py

# -*- coding:utf-8 -*-
"""
@author: isaacqyang
@time: 2022/8/29
@desc:
"""

import argparse
import time
import traceback
import uuid

import uvicorn
from fastapi import FastAPI, Request, Path
from starlette.responses import RedirectResponse

from commom import f_doc_export, get_logger, request_id_context
from entitys.response import BaseResponse

logger = get_logger()
app = FastAPI()


def doc_export(
        token: str = Path(title='doc token', description="文档的token")
):
    try:
        logger.info(f"{token=}")
        f_doc_export(token)

        return BaseResponse.ofSuccess("")

    except Exception as msg:
        logger.error(traceback.format_exc())
        return BaseResponse.ofFailure(str(msg))


@app.middleware("http")
async def add_request_id_header(request: Request, call_next):
    request_id = request.headers.get("X-REQUEST-ID")
    if request_id is None or len(request_id) == 0:
        request_id = str(uuid.uuid4())
    request_id_context.set(request_id)
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-REQUEST-ID"] = request_id_context.get()
    response.headers["PROCESS-TIME"] = f"{process_time:.2f}"

    return response


async def document():
    return RedirectResponse(url="/docs")


def api_start(host, port):
    app.get("/", response_model=BaseResponse)(document)
    app.get("/znjd/doc/export/{token}", response_model=BaseResponse)(doc_export)
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=18070)
    args = parser.parse_args()
    api_start(args.host, args.port)