your-workflow/backend/app/main.py

144 lines
4.7 KiB
Python

import uvicorn
import models.db as DB
import uuid
import time
from fastapi import FastAPI, Depends, HTTPException
from models.schemas.system import InitStatusResponse, InitConfigRequest, BaseResponse
from models.schemas.user import LoginRequest, LoginResponse
from core.config import load_init_config, save_init_config
from core.db import init_tortoise, close_tortoise
from core.logger import logger
from contextlib import asynccontextmanager
from fastapi.security import OAuth2PasswordBearer
from apscheduler.schedulers.asyncio import AsyncIOScheduler
# 简单的登陆状态存储
sessions: dict[str, tuple[float, DB.User]] = {}
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
scheduler = AsyncIOScheduler()
# 定时清理过期session
def clear_session():
TTL = 60 * 60 * 24
now = time.time()
expired_keys = [token for token, (expires_at, _) in sessions.items() if expires_at + TTL < now]
logger.info(f"检测 session, {now} {sessions}")
for token in expired_keys:
sessions.pop(token, None)
if expired_keys:
logger.info(f"清理过期 now: {now} session: {expired_keys}")
def get_current_user(token:str = Depends(oauth2_scheme)) -> DB.User | None:
if not token or token not in sessions:
raise HTTPException(status_code=401, detail="用户未登陆")
_, user = sessions[token]
sessions[token] = time.time(), user
return user
async def is_init():
init_config = await load_init_config()
init_status = init_config.get("is_init")
return init_status
@asynccontextmanager
async def lifespan(app: FastAPI):
# 数据库
if await is_init():
init_config = await load_init_config()
await init_tortoise(init_config["db_type"], init_config, init_config.get("admin_password"))
logger.info("初始化数据库")
# 定时器
scheduler.add_job(clear_session, "interval", seconds=5)
scheduler.start()
logger.info("初始化定时器")
yield
scheduler.shutdown()
logger.info("关闭定时器")
await close_tortoise()
logger.info("关闭数据库")
def create_app():
app = FastAPI(title="Workflow API", version="1.0.0", lifespan=lifespan)
@app.get("/")
async def root():
return {"message": "hello"}
@app.get(
"/check_init",
tags=["System"],
description="系统是否初始化",
response_model=InitStatusResponse,
)
async def check_init():
return InitStatusResponse(data=await is_init())
@app.post(
"/init_system",
tags=["System"],
description="系统初始化",
response_model=InitStatusResponse,
)
async def init_system(request: InitConfigRequest):
logger.info("init_system")
try:
init_config = await load_init_config()
if init_config.get('is_init'):
return InitStatusResponse(data=True)
# 保存配置
init_config["db_type"] = request.db_type
init_config["db_host"] = request.db_host
init_config["db_port"] = request.db_port
init_config["db_user"] = request.db_user
init_config["db_password"] = request.db_password
init_config["db_name"] = request.db_name
init_config["admin_password"] = request.admin_password
init_config["is_init"] = True
init_success = await init_tortoise(request.db_type, init_config, request.admin_password)
if init_success:
await save_init_config(init_config)
return InitStatusResponse(data=True)
return InitStatusResponse(success=False, message="初始化数据库失败")
except Exception as e:
return InitStatusResponse(success=False, message=str(e))
@app.post("/login", tags=["User"], description="用户登录", response_model=LoginResponse)
async def login(login_info: LoginRequest):
user = await DB.User.filter(username=login_info.username, password=login_info.get_md5_password()).first()
if user:
token = uuid.uuid4().hex
sessions[token] = time.time(), user
logger.info(sessions)
return LoginResponse(data={
"token": token,
"user": user.json(del_columns=["password"])
})
return LoginResponse(success=False, message="帐号或密码错误")
@app.post("/logout", tags=["User"], description="退出登录", response_model=BaseResponse)
async def logout(token:str):
sessions.pop(token, None)
logger.info(sessions)
return BaseResponse()
return app
app = create_app()
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)