first commit
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
from starlette.background import BackgroundTasks
|
||||
|
||||
from .ctx import CTX_BG_TASKS
|
||||
|
||||
|
||||
class BgTasks:
|
||||
"""后台任务统一管理"""
|
||||
|
||||
@classmethod
|
||||
async def init_bg_tasks_obj(cls):
|
||||
"""实例化后台任务，并设置到上下文"""
|
||||
bg_tasks = BackgroundTasks()
|
||||
CTX_BG_TASKS.set(bg_tasks)
|
||||
|
||||
@classmethod
|
||||
async def get_bg_tasks_obj(cls):
|
||||
"""从上下文中获取后台任务实例"""
|
||||
return CTX_BG_TASKS.get()
|
||||
|
||||
@classmethod
|
||||
async def add_task(cls, func, *args, **kwargs):
|
||||
"""添加后台任务"""
|
||||
bg_tasks = await cls.get_bg_tasks_obj()
|
||||
bg_tasks.add_task(func, *args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
async def execute_tasks(cls):
|
||||
"""执行后台任务，一般是请求结果返回之后执行"""
|
||||
bg_tasks = await cls.get_bg_tasks_obj()
|
||||
if bg_tasks.tasks:
|
||||
await bg_tasks()
|
||||
@@ -0,0 +1,163 @@
|
||||
# core/cache.py
|
||||
import json as _stdlib_json
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, Optional, Union
|
||||
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
# ===== 尝试使用你的 fff，否则回退到 stdlib =====
|
||||
try:
|
||||
import orjson as json
|
||||
except ImportError:
|
||||
json = _stdlib_json # type: ignore
|
||||
|
||||
|
||||
# ===== Redis 客户端单例（可替换为你自己的）=====
|
||||
class RedisClient:
|
||||
_instance: Optional["RedisClient"] = None
|
||||
_redis: Optional[Redis] = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
async def init_redis(self, url: str = "redis://localhost:6379/0") -> None:
|
||||
if self._redis is None:
|
||||
self._redis = Redis.from_url(url, decode_responses=False)
|
||||
|
||||
@property
|
||||
def client(self) -> Redis:
|
||||
if self._redis is None:
|
||||
raise RuntimeError("Redis not initialized. Call init_redis() first.")
|
||||
return self._redis
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._redis:
|
||||
await self._redis.close()
|
||||
self._redis = None
|
||||
|
||||
|
||||
redis_client = RedisClient()
|
||||
|
||||
|
||||
# ===== FastAPI 依赖注入 =====
|
||||
async def get_redis() -> Redis:
|
||||
"""FastAPI 依赖：获取 Redis 客户端"""
|
||||
return redis_client.client
|
||||
|
||||
async def invalidate_cache(key: str) -> bool:
|
||||
"""
|
||||
主动删除缓存键。
|
||||
返回是否成功删除（Redis delete 返回被删除的 key 数量）。
|
||||
"""
|
||||
try:
|
||||
result = await redis_client.client.delete(key)
|
||||
return result > 0
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to invalidate cache key: {key}", exc_info=e)
|
||||
return False
|
||||
|
||||
# ===== 缓存上下文管理器 =====
|
||||
class CacheResult:
|
||||
__slots__ = ("key", "hit", "value", "_to_set", "_set_called")
|
||||
|
||||
def __init__(self, key: str) -> None:
|
||||
self.key = key
|
||||
self.hit = False
|
||||
self.value: Any = None
|
||||
self._to_set: Any = None
|
||||
self._set_called = False
|
||||
|
||||
def set(self, value: Any) -> None:
|
||||
"""标记要缓存的值（可为 None）"""
|
||||
self._to_set = value
|
||||
self._set_called = True
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def cache_if(
|
||||
key: str,
|
||||
ttl: int = 3600,
|
||||
redis: Optional[Redis] = None,
|
||||
) -> AsyncGenerator[CacheResult, None]:
|
||||
"""
|
||||
异步缓存上下文管理器，支持 None 值缓存（防穿透）。
|
||||
|
||||
Args:
|
||||
key: 缓存键
|
||||
ttl: 正常值缓存时间（秒）
|
||||
redis: 可选 Redis 客户端（用于测试或自定义）
|
||||
|
||||
Usage:
|
||||
async with cache_if("report:123") as cache:
|
||||
if cache.hit:
|
||||
return cache.value
|
||||
result = await compute()
|
||||
cache.set(result) # result 可为 None
|
||||
"""
|
||||
result = CacheResult(key)
|
||||
client = redis or redis_client.client
|
||||
|
||||
# 尝试读缓存
|
||||
try:
|
||||
cached_val = await client.get(key)
|
||||
if cached_val is not None:
|
||||
# 解码
|
||||
if cached_val == b"__NULL__":
|
||||
result.hit = True
|
||||
result.value = None
|
||||
else:
|
||||
result.hit = True
|
||||
result.value = json.loads(cached_val)
|
||||
except Exception:
|
||||
# Redis 不可用，降级（不中断主流程）
|
||||
logger.exception(f"Redis get error for key: {key}")
|
||||
pass
|
||||
|
||||
yield result
|
||||
|
||||
# 写缓存（仅当调用了 set()）
|
||||
if result._set_called:
|
||||
try:
|
||||
if result._to_set is None:
|
||||
val = b"__NULL__"
|
||||
ex = 60 # 空值短 TTL
|
||||
else:
|
||||
# 注意：orjson.dumps 返回 bytes，stdlib 返回 str → 统一转 bytes
|
||||
serialized = json.dumps(result._to_set)
|
||||
val = serialized if isinstance(serialized, bytes) else serialized.encode("utf-8")
|
||||
ex = ttl
|
||||
await client.setex(key, ex, val)
|
||||
except Exception:
|
||||
# 写缓存失败，不影响主流程
|
||||
pass
|
||||
|
||||
|
||||
# ===== 装饰器版（可选补充）=====
|
||||
from functools import wraps
|
||||
import asyncio
|
||||
import hashlib
|
||||
|
||||
def cached(ttl: int = 3600):
|
||||
"""函数缓存装饰器（使用 cache_if）"""
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
# 生成 key（简单版，可替换为更 robust 的）
|
||||
key_data = str(args) + str(sorted(kwargs.items()))
|
||||
key = f"cached:{func.__name__}:{hashlib.md5(key_data.encode()).hexdigest()}"
|
||||
logger.debug(f"Cache key: {key}")
|
||||
|
||||
async with cache_if(key, ttl) as cache:
|
||||
if cache.hit:
|
||||
logger.debug(f"Cache hit for key: {key}")
|
||||
return cache.value
|
||||
result = await func(*args, **kwargs)
|
||||
cache.set(result)
|
||||
return result
|
||||
return wrapper
|
||||
return decorator
|
||||
@@ -0,0 +1,78 @@
|
||||
from typing import Any, Callable, Dict, Generic, List, NewType, Tuple, Type, TypeVar, Union
|
||||
|
||||
from pydantic import BaseModel
|
||||
from tortoise.expressions import Q
|
||||
from tortoise.models import Model
|
||||
|
||||
Total = NewType("Total", int)
|
||||
ModelType = TypeVar("ModelType", bound=Model)
|
||||
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
|
||||
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
|
||||
|
||||
|
||||
class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
|
||||
def __init__(self, model: Type[ModelType]):
|
||||
self.model = model
|
||||
|
||||
async def is_exist(self, **kwargs) -> bool:
|
||||
return await self.model.filter(**kwargs).first()
|
||||
|
||||
async def all(self, search: Q = Q()) -> List[ModelType]:
|
||||
return await self.model.filter(search).all()
|
||||
|
||||
async def get(self, id: int) -> ModelType:
|
||||
return await self.model.get(id=id)
|
||||
|
||||
async def list(self, page: int, page_size: int, search: Q = Q(), order: list = []) -> Tuple[Total, List[ModelType]]:
|
||||
query = self.model.filter(search)
|
||||
return await query.count(), await query.offset((page - 1) * page_size).limit(page_size).order_by(*order)
|
||||
|
||||
async def create(self, obj_in: CreateSchemaType) -> ModelType:
|
||||
if isinstance(obj_in, Dict):
|
||||
obj_dict = obj_in
|
||||
else:
|
||||
obj_dict = obj_in.model_dump()
|
||||
obj = self.model(**obj_dict)
|
||||
await obj.save()
|
||||
return obj
|
||||
|
||||
async def copy(self, id: int, handler: Callable[[Dict[str, Any]], Dict[str, Any]] = None) -> ModelType:
|
||||
obj = await self.get(id=id)
|
||||
obj_dict = await obj.to_dict()
|
||||
drop_keys = ["id", "created_at", "updated_at"]
|
||||
for key in drop_keys:
|
||||
obj_dict.pop(key, None)
|
||||
if handler:
|
||||
obj_dict = handler(obj_dict)
|
||||
obj = self.model(**obj_dict)
|
||||
await obj.save()
|
||||
return obj
|
||||
|
||||
async def update(self, id: int, obj_in: Union[UpdateSchemaType, Dict[str, Any]]) -> ModelType:
|
||||
if isinstance(obj_in, Dict):
|
||||
obj_dict = obj_in
|
||||
else:
|
||||
obj_dict = obj_in.model_dump(exclude_unset=True, exclude={"id"})
|
||||
obj = await self.get(id=id)
|
||||
obj = obj.update_from_dict(obj_dict)
|
||||
await obj.save()
|
||||
return obj
|
||||
|
||||
async def create_or_update(self, obj_in: CreateSchemaType, query_kwargs: Dict[str, Any], update_kwargs: Dict[str, Any] = None) -> ModelType:
|
||||
update_kwargs = update_kwargs or {}
|
||||
orm = await self.model.filter(**query_kwargs).first()
|
||||
if orm:
|
||||
need_update = False if update_kwargs else True
|
||||
for key, value in update_kwargs.items():
|
||||
if getattr(orm, key) != value:
|
||||
need_update = True
|
||||
break
|
||||
if need_update:
|
||||
return 'update', await self.update(orm.id, obj_in)
|
||||
else:
|
||||
return False, orm
|
||||
return 'create', await self.create(obj_in)
|
||||
|
||||
async def remove(self, id: int) -> None:
|
||||
obj = await self.get(id=id)
|
||||
await obj.delete()
|
||||
@@ -0,0 +1,17 @@
|
||||
import contextvars
|
||||
|
||||
from starlette.background import BackgroundTasks
|
||||
|
||||
CTX_USER_ID: contextvars.ContextVar[int] = contextvars.ContextVar("user_id", default=0)
|
||||
CTX_USER_NAME: contextvars.ContextVar[str] = contextvars.ContextVar("user_name", default="")
|
||||
CTX_BG_TASKS: contextvars.ContextVar[BackgroundTasks] = contextvars.ContextVar("bg_task", default=None)
|
||||
|
||||
def set_ctx_weixin_user(user_id: int, user_name: str):
|
||||
CTX_USER_ID.set(user_id)
|
||||
CTX_USER_NAME.set(user_name)
|
||||
|
||||
def get_ctx_weixin_user():
|
||||
return {
|
||||
"user_id": CTX_USER_ID.get(),
|
||||
"user_name": CTX_USER_NAME.get(),
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
from typing import Optional
|
||||
|
||||
import jwt
|
||||
from fastapi import Depends, Header, HTTPException, Request
|
||||
|
||||
from app.core.ctx import CTX_USER_ID
|
||||
from app.models import Role, User, WeixinUser
|
||||
from app.settings import settings
|
||||
|
||||
|
||||
class AuthControl:
|
||||
@classmethod
|
||||
async def is_authed(cls, token: str = Header(..., description="token验证")) -> Optional["User"]:
|
||||
try:
|
||||
if token == "dev":
|
||||
user = await User.filter().first()
|
||||
user_id = user.id
|
||||
else:
|
||||
decode_data = jwt.decode(token, settings.SECRET_KEY, algorithms=settings.JWT_ALGORITHM)
|
||||
user_id = decode_data.get("user_id")
|
||||
user = await User.filter(id=user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(status_code=401, detail="Authentication failed")
|
||||
CTX_USER_ID.set(int(user_id))
|
||||
return user
|
||||
except jwt.DecodeError:
|
||||
raise HTTPException(status_code=401, detail="无效的Token")
|
||||
except jwt.ExpiredSignatureError:
|
||||
raise HTTPException(status_code=401, detail="登录已过期")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"{repr(e)}")
|
||||
|
||||
@classmethod
|
||||
async def weixin_user(cls, request: Request) -> WeixinUser:
|
||||
path = request.url.path
|
||||
if path in [
|
||||
"/api/v1/msg/new_order",
|
||||
"/api/v1/event/feishu",
|
||||
"/api/v1/msg/get_order_user_days_before",
|
||||
]:
|
||||
return None
|
||||
|
||||
token = request.headers.get("token")
|
||||
user = await WeixinUser.filter(userid=token).first()
|
||||
assert user, "Authentication failed"
|
||||
CTX_USER_ID.set(int(user.id))
|
||||
return user # 返回的是 WeixinUser 实例
|
||||
|
||||
class PermissionControl:
|
||||
@classmethod
|
||||
async def has_permission(cls, request: Request, current_user: User = Depends(AuthControl.is_authed)) -> None:
|
||||
if current_user.is_superuser:
|
||||
return
|
||||
method = request.method
|
||||
path = request.url.path
|
||||
roles: list[Role] = await current_user.roles
|
||||
if not roles:
|
||||
raise HTTPException(status_code=403, detail="The user is not bound to a role")
|
||||
apis = [await role.apis for role in roles]
|
||||
permission_apis = list(set((api.method, api.path) for api in sum(apis, [])))
|
||||
# path = "/api/v1/auth/userinfo"
|
||||
# method = "GET"
|
||||
if (method, path) not in permission_apis:
|
||||
raise HTTPException(status_code=403, detail=f"Permission denied method:{method} path:{path}")
|
||||
|
||||
|
||||
DependAuth = Depends(AuthControl.is_authed)
|
||||
DependWeixinUser = Depends(AuthControl.weixin_user)
|
||||
DependPermisson = Depends(PermissionControl.has_permission)
|
||||
@@ -0,0 +1,150 @@
|
||||
from fastapi.exceptions import (
|
||||
HTTPException,
|
||||
RequestValidationError,
|
||||
ResponseValidationError,
|
||||
)
|
||||
from fastapi.requests import Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from tortoise.exceptions import DoesNotExist, IntegrityError
|
||||
|
||||
from app.http_base import HttpResp
|
||||
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class SettingNotFound(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class AssertException(Exception):
|
||||
"""断言异常"""
|
||||
def __init__(self, message: str, code: int = 1001):
|
||||
self.message = message
|
||||
self.code = code
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
async def DoesNotExistHandle(req: Request, exc: DoesNotExist) -> JSONResponse:
|
||||
"""处理对象不存在异常"""
|
||||
resp = HttpResp.REQUEST_404_ERROR
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=f"Object has not found, exc: {exc}, query_params: {req.query_params}",
|
||||
data=None
|
||||
)
|
||||
logger.warning(f"Object not found: {exc}, path: {req.url.path}")
|
||||
return JSONResponse(content=content, status_code=200)
|
||||
|
||||
|
||||
async def IntegrityHandle(req: Request, exc: IntegrityError) -> JSONResponse:
|
||||
"""处理数据完整性异常"""
|
||||
resp = HttpResp.SYSTEM_ERROR
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=f"IntegrityError: {str(exc)}",
|
||||
data=None
|
||||
)
|
||||
logger.error(f"Integrity error: {exc}, path: {req.url.path}", exc_info=True)
|
||||
return JSONResponse(content=content, status_code=200)
|
||||
|
||||
|
||||
async def HttpExcHandle(req: Request, exc: HTTPException) -> JSONResponse:
|
||||
"""处理 HTTP 异常"""
|
||||
content = dict(
|
||||
code=exc.status_code,
|
||||
msg=exc.detail,
|
||||
data=None
|
||||
)
|
||||
logger.warning(f"HTTP exception: {exc.status_code} - {exc.detail}, path: {req.url.path}")
|
||||
return JSONResponse(content=content, status_code=200)
|
||||
|
||||
|
||||
async def RequestValidationHandle(req: Request, exc: RequestValidationError) -> JSONResponse:
|
||||
"""处理请求验证异常"""
|
||||
resp = HttpResp.PARAMS_VALID_ERROR
|
||||
errors = []
|
||||
for error in exc.errors():
|
||||
errors.append({
|
||||
"loc": error["loc"],
|
||||
"msg": error["msg"],
|
||||
"type": error["type"]
|
||||
})
|
||||
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=f"{resp.msg}, errors: {errors}",
|
||||
data=errors
|
||||
)
|
||||
logger.warning(f"Request validation error: {errors}, path: {req.url.path}")
|
||||
return JSONResponse(content=content, status_code=200)
|
||||
|
||||
|
||||
async def ResponseValidationHandle(req: Request, exc: ResponseValidationError) -> JSONResponse:
|
||||
"""处理响应验证异常"""
|
||||
resp = HttpResp.SYSTEM_ERROR
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=f"Response validation error: {str(exc)}",
|
||||
data=None
|
||||
)
|
||||
logger.error(f"Response validation error: {exc}, path: {req.url.path}", exc_info=True)
|
||||
return JSONResponse(content=content, status_code=200)
|
||||
|
||||
|
||||
async def AssertExceptionHandle(request: Request, exc: AssertException) -> JSONResponse:
|
||||
"""处理断言异常"""
|
||||
resp = HttpResp.ASSERT_ARGUMENT_ERROR
|
||||
content = dict(
|
||||
code=exc.code,
|
||||
msg=exc.message,
|
||||
data=None
|
||||
)
|
||||
logger.warning(f"Assert exception: {exc.message}, path: {request.url.path}")
|
||||
return JSONResponse(
|
||||
content=content,
|
||||
status_code=200
|
||||
)
|
||||
|
||||
|
||||
async def AssertionErrorHandle(request: Request, exc: AssertionError) -> JSONResponse:
|
||||
"""处理 Python assert 异常"""
|
||||
resp = HttpResp.ASSERT_ARGUMENT_ERROR
|
||||
message = str(exc) if exc.args else resp.msg
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=message,
|
||||
data=None
|
||||
)
|
||||
logger.warning(f"Assertion error: {message}, path: {request.url.path}")
|
||||
return JSONResponse(
|
||||
content=content,
|
||||
status_code=200
|
||||
)
|
||||
|
||||
|
||||
async def GlobalExceptionHandler(request: Request, exc: Exception) -> JSONResponse:
|
||||
"""全局异常处理器"""
|
||||
resp = HttpResp.SYSTEM_ERROR
|
||||
content = dict(
|
||||
code=resp.code,
|
||||
msg=f"{resp.msg}: {str(exc)}",
|
||||
data=None
|
||||
)
|
||||
logger.error(f"Global exception: {str(exc)}, path: {request.url.path}", exc_info=True)
|
||||
return JSONResponse(
|
||||
content=content,
|
||||
status_code=200
|
||||
)
|
||||
|
||||
|
||||
# 便捷函数用于注册所有异常处理器
|
||||
def register_exception_handlers(app):
|
||||
"""注册所有异常处理器到 FastAPI 应用"""
|
||||
app.add_exception_handler(DoesNotExist, DoesNotExistHandle)
|
||||
app.add_exception_handler(IntegrityError, IntegrityHandle)
|
||||
app.add_exception_handler(HTTPException, HttpExcHandle)
|
||||
app.add_exception_handler(RequestValidationError, RequestValidationHandle)
|
||||
app.add_exception_handler(ResponseValidationError, ResponseValidationHandle)
|
||||
app.add_exception_handler(AssertException, AssertExceptionHandle)
|
||||
app.add_exception_handler(AssertionError, AssertionErrorHandle) # 捕获 Python assert
|
||||
app.add_exception_handler(Exception, GlobalExceptionHandler) # 全局异常处理器
|
||||
@@ -0,0 +1,273 @@
|
||||
import shutil
|
||||
|
||||
from aerich import Command
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware import Middleware
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from tortoise.expressions import Q
|
||||
|
||||
from app.api import api_router
|
||||
from app.controllers.api import api_controller
|
||||
from app.controllers.user import UserCreate, user_controller
|
||||
from app.core.exceptions import register_exception_handlers
|
||||
from app.log import logger
|
||||
from app.models.admin import Api, Menu, Role, Dept
|
||||
from app.models.automation import Task, Scenario
|
||||
from app.schemas.menus import MenuType
|
||||
from app.settings.config import settings
|
||||
from app.core.cache import redis_client
|
||||
|
||||
from .middlewares import BackGroundTaskMiddleware, HttpAuditLogMiddleware
|
||||
|
||||
|
||||
def make_middlewares():
|
||||
middleware = [
|
||||
Middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.CORS_ORIGINS,
|
||||
allow_credentials=settings.CORS_ALLOW_CREDENTIALS,
|
||||
allow_methods=settings.CORS_ALLOW_METHODS,
|
||||
allow_headers=settings.CORS_ALLOW_HEADERS,
|
||||
),
|
||||
Middleware(BackGroundTaskMiddleware),
|
||||
Middleware(
|
||||
HttpAuditLogMiddleware,
|
||||
methods=["GET", "POST", "PUT", "DELETE"],
|
||||
exclude_paths=[
|
||||
"/api/v1/msg/new_order",
|
||||
"/api/v1/weixin/jssdk-config",
|
||||
"/api/v1/base/access_token",
|
||||
"/docs",
|
||||
"/api-docs",
|
||||
"/openapi.json",
|
||||
"/static/*",
|
||||
"/gen/*",
|
||||
],
|
||||
),
|
||||
]
|
||||
return middleware
|
||||
|
||||
def mount_static_and_config_swagger(app: FastAPI):
|
||||
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.openapi.docs import get_swagger_ui_html
|
||||
|
||||
app.mount("/static", StaticFiles(directory="static"), name="静态文件")
|
||||
|
||||
@app.get(app.docs_url, include_in_schema=False)
|
||||
async def custom_swagger_ui_html():
|
||||
print('访问文档')
|
||||
return get_swagger_ui_html(
|
||||
openapi_url=app.openapi_url,
|
||||
title=app.title,
|
||||
swagger_js_url="/static/swagger-ui/swagger-ui-bundle.js",
|
||||
swagger_css_url="/static/swagger-ui/swagger-ui.css"
|
||||
)
|
||||
|
||||
def register_exceptions(app: FastAPI):
|
||||
register_exception_handlers(app)
|
||||
print('注册异常处理完成')
|
||||
|
||||
|
||||
def register_routers(app: FastAPI, prefix: str = "/api"):
|
||||
app.include_router(api_router, prefix=prefix)
|
||||
|
||||
|
||||
async def init_superuser():
|
||||
user = await user_controller.model.exists()
|
||||
if not user:
|
||||
await user_controller.create_user(
|
||||
UserCreate(
|
||||
username="admin",
|
||||
email="admin@admin.com",
|
||||
password="123456",
|
||||
is_active=True,
|
||||
is_superuser=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def init_menus():
|
||||
menus = await Menu.exists()
|
||||
if not menus:
|
||||
parent_menu = await Menu.create(
|
||||
menu_type=MenuType.CATALOG,
|
||||
name="系统管理",
|
||||
path="/system",
|
||||
order=1,
|
||||
parent_id=0,
|
||||
icon="carbon:gui-management",
|
||||
is_hidden=False,
|
||||
component="Layout",
|
||||
keepalive=False,
|
||||
redirect="/system/user",
|
||||
)
|
||||
children_menu = [
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="用户管理",
|
||||
path="user",
|
||||
order=1,
|
||||
parent_id=parent_menu.id,
|
||||
icon="material-symbols:person-outline-rounded",
|
||||
is_hidden=False,
|
||||
component="/system/user",
|
||||
keepalive=False,
|
||||
),
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="角色管理",
|
||||
path="role",
|
||||
order=2,
|
||||
parent_id=parent_menu.id,
|
||||
icon="carbon:user-role",
|
||||
is_hidden=False,
|
||||
component="/system/role",
|
||||
keepalive=False,
|
||||
),
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="菜单管理",
|
||||
path="menu",
|
||||
order=3,
|
||||
parent_id=parent_menu.id,
|
||||
icon="material-symbols:list-alt-outline",
|
||||
is_hidden=False,
|
||||
component="/system/menu",
|
||||
keepalive=False,
|
||||
),
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="API管理",
|
||||
path="api",
|
||||
order=4,
|
||||
parent_id=parent_menu.id,
|
||||
icon="ant-design:api-outlined",
|
||||
is_hidden=False,
|
||||
component="/system/api",
|
||||
keepalive=False,
|
||||
),
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="部门管理",
|
||||
path="dept",
|
||||
order=5,
|
||||
parent_id=parent_menu.id,
|
||||
icon="mingcute:department-line",
|
||||
is_hidden=False,
|
||||
component="/system/dept",
|
||||
keepalive=False,
|
||||
),
|
||||
Menu(
|
||||
menu_type=MenuType.MENU,
|
||||
name="审计日志",
|
||||
path="auditlog",
|
||||
order=6,
|
||||
parent_id=parent_menu.id,
|
||||
icon="ph:clipboard-text-bold",
|
||||
is_hidden=False,
|
||||
component="/system/auditlog",
|
||||
keepalive=False,
|
||||
),
|
||||
]
|
||||
await Menu.bulk_create(children_menu)
|
||||
await Menu.create(
|
||||
menu_type=MenuType.MENU,
|
||||
name="一级菜单",
|
||||
path="/top-menu",
|
||||
order=2,
|
||||
parent_id=0,
|
||||
icon="material-symbols:featured-play-list-outline",
|
||||
is_hidden=False,
|
||||
component="/top-menu",
|
||||
keepalive=False,
|
||||
redirect="",
|
||||
)
|
||||
|
||||
|
||||
async def init_apis():
|
||||
apis = await api_controller.model.exists()
|
||||
if not apis:
|
||||
await api_controller.refresh_api()
|
||||
|
||||
|
||||
async def init_db():
|
||||
command = Command(tortoise_config=settings.TORTOISE_ORM)
|
||||
try:
|
||||
await command.init_db(safe=True)
|
||||
except FileExistsError:
|
||||
pass
|
||||
|
||||
await command.init()
|
||||
try:
|
||||
await command.migrate()
|
||||
except AttributeError:
|
||||
logger.warning("unable to retrieve model history from database, model history will be created from scratch")
|
||||
shutil.rmtree("migrations")
|
||||
await command.init_db(safe=True)
|
||||
|
||||
await command.upgrade(run_in_transaction=True)
|
||||
|
||||
|
||||
async def init_roles():
|
||||
roles = await Role.exists()
|
||||
if not roles:
|
||||
admin_role = await Role.create(
|
||||
name="管理员",
|
||||
desc="管理员角色",
|
||||
)
|
||||
user_role = await Role.create(
|
||||
name="普通用户",
|
||||
desc="普通用户角色",
|
||||
)
|
||||
|
||||
await Dept.create(name="默认部门")
|
||||
|
||||
# 分配所有API给管理员角色
|
||||
all_apis = await Api.all()
|
||||
await admin_role.apis.add(*all_apis)
|
||||
# 分配所有菜单给管理员和普通用户
|
||||
all_menus = await Menu.all()
|
||||
await admin_role.menus.add(*all_menus)
|
||||
await user_role.menus.add(*all_menus)
|
||||
|
||||
# 为普通用户分配基本API
|
||||
basic_apis = await Api.filter(Q(method__in=["GET"]) | Q(tags="基础模块"))
|
||||
await user_role.apis.add(*basic_apis)
|
||||
|
||||
async def init_cache():
|
||||
await redis_client.init_redis(settings.REDIS_URL)
|
||||
|
||||
async def init_task():
|
||||
|
||||
task = await Task.exists(title="系统待办")
|
||||
if task: return
|
||||
|
||||
scenario = await Scenario.create(
|
||||
title="系统待办",
|
||||
visible=False,
|
||||
trigger={},
|
||||
actions=[],
|
||||
enabled=False,
|
||||
notes="用于为所有存量数据创建的待办任务",
|
||||
)
|
||||
task = await Task.create(
|
||||
title="系统待办",
|
||||
status='success',
|
||||
ui_schema={},
|
||||
source_scenario=scenario,
|
||||
notes="当系统有新的待办任务时关联动作",
|
||||
)
|
||||
|
||||
|
||||
async def init_data():
|
||||
await init_db()
|
||||
await init_task()
|
||||
await init_cache()
|
||||
await init_superuser()
|
||||
await init_menus()
|
||||
await init_apis()
|
||||
await init_roles()
|
||||
|
||||
async def tear_down():
|
||||
await redis_client.close()
|
||||
@@ -0,0 +1,192 @@
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime
|
||||
from typing import Any, AsyncGenerator
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import Response
|
||||
from fastapi.routing import APIRoute
|
||||
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||
from starlette.requests import Request
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from app.core.dependency import AuthControl
|
||||
from app.models.admin import AuditLog, User
|
||||
|
||||
from .bgtask import BgTasks
|
||||
|
||||
|
||||
class SimpleBaseMiddleware:
|
||||
def __init__(self, app: ASGIApp) -> None:
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
request = Request(scope, receive=receive)
|
||||
|
||||
response = await self.before_request(request) or self.app
|
||||
await response(request.scope, request.receive, send)
|
||||
await self.after_request(request)
|
||||
|
||||
async def before_request(self, request: Request):
|
||||
return self.app
|
||||
|
||||
async def after_request(self, request: Request):
|
||||
return None
|
||||
|
||||
|
||||
class BackGroundTaskMiddleware(SimpleBaseMiddleware):
|
||||
async def before_request(self, request):
|
||||
await BgTasks.init_bg_tasks_obj()
|
||||
|
||||
async def after_request(self, request):
|
||||
await BgTasks.execute_tasks()
|
||||
|
||||
|
||||
class HttpAuditLogMiddleware(BaseHTTPMiddleware):
|
||||
def __init__(self, app, methods: list[str], exclude_paths: list[str]):
|
||||
super().__init__(app)
|
||||
self.methods = methods
|
||||
self.exclude_paths = exclude_paths
|
||||
self.audit_log_paths = ["/api/v1/auditlog/list"]
|
||||
self.max_body_size = 1024 * 1024 # 1MB 响应体大小限制
|
||||
|
||||
async def get_request_args(self, request: Request) -> dict:
|
||||
args = {}
|
||||
# 获取查询参数
|
||||
for key, value in request.query_params.items():
|
||||
args[key] = value
|
||||
|
||||
# 判断是否为文件上传请求
|
||||
content_type = request.headers.get("content-type", "").lower()
|
||||
is_upload = (
|
||||
request.method in {"POST", "PUT"} # 上传通常用 POST/PUT
|
||||
and (
|
||||
"multipart/form-data" in content_type # 标准文件上传
|
||||
or "application/octet-stream" in content_type # 二进制流上传（较少见）
|
||||
)
|
||||
)
|
||||
|
||||
if is_upload:
|
||||
return {}
|
||||
|
||||
# 获取请求体
|
||||
if request.method in ["POST", "PUT", "PATCH"]:
|
||||
try:
|
||||
body = await request.json()
|
||||
args.update(body)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
body = await request.form()
|
||||
args.update(body)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return args
|
||||
|
||||
async def get_response_body(self, request: Request, response: Response) -> Any:
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
if not content_type.startswith("application/json"):
|
||||
return {"msg": "Non-JSON response (e.g., file download), skipped for audit"}
|
||||
|
||||
# 检查Content-Length
|
||||
content_length = response.headers.get("content-length")
|
||||
if content_length and int(content_length) > self.max_body_size:
|
||||
return {"code": 0, "msg": "Response too large to log", "data": None}
|
||||
|
||||
if hasattr(response, "body"):
|
||||
body = response.body
|
||||
else:
|
||||
body_chunks = []
|
||||
async for chunk in response.body_iterator:
|
||||
if not isinstance(chunk, bytes):
|
||||
chunk = chunk.encode(response.charset)
|
||||
body_chunks.append(chunk)
|
||||
|
||||
response.body_iterator = self._async_iter(body_chunks)
|
||||
body = b"".join(body_chunks)
|
||||
|
||||
if any(request.url.path.startswith(path) for path in self.audit_log_paths):
|
||||
try:
|
||||
data = self.lenient_json(body)
|
||||
# 只保留基本信息，去除详细的响应内容
|
||||
if isinstance(data, dict):
|
||||
data.pop("response_body", None)
|
||||
if "data" in data and isinstance(data["data"], list):
|
||||
for item in data["data"]:
|
||||
item.pop("response_body", None)
|
||||
return data
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
return self.lenient_json(body)
|
||||
|
||||
def lenient_json(self, v: Any) -> Any:
|
||||
if isinstance(v, (str, bytes)):
|
||||
try:
|
||||
return json.loads(v)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
return v
|
||||
|
||||
async def _async_iter(self, items: list[bytes]) -> AsyncGenerator[bytes, None]:
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
async def get_request_log(self, request: Request, response: Response) -> dict:
|
||||
"""
|
||||
根据request和response对象获取对应的日志记录数据
|
||||
"""
|
||||
data: dict = {"path": request.url.path, "status": response.status_code, "method": request.method}
|
||||
# 路由信息
|
||||
app: FastAPI = request.app
|
||||
for route in app.routes:
|
||||
if (
|
||||
isinstance(route, APIRoute)
|
||||
and route.path_regex.match(request.url.path)
|
||||
and request.method in route.methods
|
||||
):
|
||||
data["module"] = ",".join(route.tags)
|
||||
data["summary"] = route.summary
|
||||
# 获取用户信息
|
||||
try:
|
||||
token = request.headers.get("token")
|
||||
user_obj = None
|
||||
if token:
|
||||
user_obj: User = await AuthControl.is_authed(token)
|
||||
data["user_id"] = user_obj.id if user_obj else 0
|
||||
data["username"] = user_obj.username if user_obj else ""
|
||||
except Exception:
|
||||
data["user_id"] = 0
|
||||
data["username"] = ""
|
||||
return data
|
||||
|
||||
async def before_request(self, request: Request):
|
||||
request_args = await self.get_request_args(request)
|
||||
request.state.request_args = request_args
|
||||
|
||||
async def after_request(self, request: Request, response: Response, process_time: int):
|
||||
if request.method in self.methods:
|
||||
for path in self.exclude_paths:
|
||||
if re.search(path, request.url.path, re.I) is not None:
|
||||
return
|
||||
data: dict = await self.get_request_log(request=request, response=response)
|
||||
data["response_time"] = process_time
|
||||
|
||||
data["request_args"] = request.state.request_args
|
||||
data["response_body"] = await self.get_response_body(request, response)
|
||||
await AuditLog.create(**data)
|
||||
|
||||
return response
|
||||
|
||||
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||
start_time: datetime = datetime.now()
|
||||
await self.before_request(request)
|
||||
response = await call_next(request)
|
||||
end_time: datetime = datetime.now()
|
||||
process_time = int((end_time.timestamp() - start_time.timestamp()) * 1000)
|
||||
await self.after_request(request, response, process_time)
|
||||
return response
|
||||
Reference in New Issue
Block a user