85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
from typing import Dict
|
|
from typing import List
|
|
from typing import NoReturn
|
|
from typing import Optional
|
|
|
|
import sentry_sdk
|
|
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
|
|
from aioredis import Redis
|
|
from aioredis import from_url
|
|
from fastapi import FastAPI
|
|
from starlette.requests import Request
|
|
from tortoise.contrib.fastapi import register_tortoise
|
|
|
|
from core.config import settings
|
|
|
|
sentry_sdk.init(settings.SENTRY_DSN)
|
|
|
|
|
|
def mysql_cdn():
|
|
"""mysql cdn 完整地址"""
|
|
return f'mysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@{settings.MYSQL_ADDRESS}/{settings.MYSQL_DB}'
|
|
|
|
|
|
def redis_cdn():
|
|
"""redis cdn"""
|
|
return f'redis://{settings.REDIS_ADDRESS}'
|
|
|
|
|
|
def register_middleware(app: FastAPI):
|
|
"""注册 日志错误收集 中间件"""
|
|
app.add_middleware(SentryAsgiMiddleware)
|
|
|
|
|
|
def register_orm(app: FastAPI, modules: Optional[Dict[str, List[str]]] = None) -> NoReturn:
|
|
"""注册orm"""
|
|
if modules is None:
|
|
modules = {"models": ["models"]}
|
|
register_tortoise(
|
|
app,
|
|
db_url=mysql_cdn(),
|
|
modules=modules
|
|
)
|
|
|
|
|
|
def redis_pool() -> Dict[int, Redis]:
|
|
"""redis 16个db 每个创建池"""
|
|
return {db: from_url(redis_cdn(), db=db, max_connections=10, encoding="utf-8", decode_responses=True)
|
|
for db in range(16)
|
|
}
|
|
|
|
|
|
def close_redis_pool(pool: Dict[int, Redis]):
|
|
for v in pool.values():
|
|
v.close()
|
|
|
|
|
|
def get_redis_db(request: Request, db: int = 0) -> Redis:
|
|
"""获取到FastAPI中的缓存redis 字典获取 指定 db"""
|
|
return request.app.state.redis_pool.get(db)
|
|
|
|
|
|
def register_redis(app: FastAPI):
|
|
"""redis 注册"""
|
|
|
|
@app.on_event("startup")
|
|
async def create_redis():
|
|
app.state.redis_pool = redis_pool()
|
|
|
|
@app.on_event("shutdown")
|
|
async def close_redis():
|
|
close_redis_pool(app.state.redis_pool)
|
|
|
|
|
|
# orm 迁移生成表
|
|
# https://tortoise.github.io/migration.html
|
|
TORTOISE_ORM = {
|
|
"connections": {"default": mysql_cdn()},
|
|
"apps": {
|
|
"models": {
|
|
"models": ["models", "aerich.models"],
|
|
"default_connection": "default",
|
|
},
|
|
}
|
|
}
|