69 lines
1.6 KiB
Python
69 lines
1.6 KiB
Python
from typing import Dict
|
|
from typing import List
|
|
from typing import NoReturn
|
|
from typing import Optional
|
|
|
|
|
|
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 common.config import settings
|
|
|
|
|
|
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=settings.MYSQL_DATABASE_URL,
|
|
modules=modules
|
|
)
|
|
|
|
|
|
def redis_pool() -> Dict[int, Redis]:
|
|
"""redis 16个db 每个创建池"""
|
|
return {db: from_url(settings.REDIS_URL, 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": settings.MYSQL_DATABASE_URL},
|
|
"apps": {
|
|
"models": {
|
|
"models": ["models", "aerich.models"],
|
|
"default_connection": "default",
|
|
},
|
|
}
|
|
}
|
|
|
|
|