import uvicorn import uuid import time import models.db as DB import nodes as F 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 # 1天有效期 now = time.time() expired_keys = [ token for token, (expires_at, _) in sessions.items() if expires_at + TTL < now ] 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() @app.post( "/run", tags=["Flow"], description="执行命令", response_model=BaseResponse ) async def flow_run(id: int): nodes = [] local_node = F.CmdNode() local_node.set_id("0") nodes.append(local_node) a800_node = F.CmdNode() a800_node.set_id("1") nodes.append(a800_node) # 设置依赖关系 # node0 - "true" >> node1 - "true" >> node2 for i in range(len(nodes) -1): # nodes[i] - "error" >> nodes[i + 1] pass flow = F.AsyncFlow(a800_node) flow.set_params( { "0": { "cmd": """ echo "1) 前端 - 拉取新代码" cd /Users/liruwei/Documents/code/project/dabai/dbpai-web git stash git pull echo "2) 后端 - 拉取新代码" cd /Users/liruwei/Documents/code/project/dabai/dbpai-server git stash echo "3) 后端 - 更新 dbpai-chat-widget.min.js" cp /Users/liruwei/Documents/code/project/dabai/dbpai-web/public/dbpai-chat-widget.min.js ./src/static git commit -am "feat: dbpai-chat-widget.min.js" git push echo "4) 打包" cd .. sh build-dbpai-jenkisn.sh all cd dbpai-jenkins git commit -am "build" git push echo "5) 还原" cd /Users/liruwei/Documents/code/project/dabai/dbpai-web git stash pop cd /Users/liruwei/Documents/code/project/dabai/dbpai-server git stash pop """ }, "1": { "cmd": """ cd /data/PAI/dbpai-beta/dbpai-jenkins git stash git pull origin v3.4.2 git stash pop docker logs ai-guard-server -n 50 """, "host": "58.214.239.10", "user": "dgx07", "password": "1f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, }, } ) shared = {} await flow.run_async(shared) logger.info(shared) return BaseResponse() return app app = create_app() if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)