import asyncio from hashlib import md5 import json import uvicorn import uuid import time import models.db as DB import nodes as F from tortoise.transactions import in_transaction from fastapi import FastAPI, Depends, HTTPException from models.schemas.system import InitStatusResponse, InitConfigRequest, BaseResponse, TableResponse 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 async def _create_demo_records(app: FastAPI): """创建测试用的""" if await DB.Work.exists(): return async with in_transaction(): work = DB.Work( uuid=uuid.uuid4().hex, creator_id=1, work_title="pai-beta", work_desc="大白智护beta环境", ) await work.save() work_content = DB.WorkContent( user_id=work.creator_id, work_id=work.id, content=[ { "type": "git", "title": "同步前端仓库", "description": "clone 分支 pai-beta-build", "uuid": uuid.uuid4().hex, "detail": { "host": "58.214.239.10", "user": "dgx07", "password": "f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, "git_path": "/data/wrok-flow/dbpai-web", "git_user": "liruwei", "git_password": "Wsad.8246", "git_url": "https://tencentgit.dabby.com.cn/dbpai/dbpai-web.git", "git_branch": "pai-beta-build", "git_clone_type": "http/https", "git_clone_depth": 1, }, }, { "type": "git", "title": "同步后端仓库", "description": "clone 分支 v3.4.2", "uuid": uuid.uuid4().hex, "detail": { "host": "58.214.239.10", "user": "dgx07", "password": "f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, "git_path": "/data/wrok-flow/dbpai-server", "git_user": "liruwei", "git_password": "Wsad.8246", "git_url": "https://tencentgit.dabby.com.cn/dbpai/dbpai-server.git", "git_branch": "v3.4.2", "git_clone_type": "http/https", "git_clone_depth": 1, }, }, { "type": "git", "title": "同步项目集合仓库", "description": "clone 分支 v3.4.2", "uuid": uuid.uuid4().hex, "detail": { "host": "58.214.239.10", "user": "dgx07", "password": "f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, "git_path": "/data/wrok-flow/dbpai-jenkins", "git_user": "liruwei", "git_password": "Wsad.8246", "git_url": "https://tencentgit.dabby.com.cn/dbpai/dbpai-jenkins.git", "git_branch": "v3.4.2", "git_clone_type": "http/https", "git_clone_depth": 1, }, }, { "type": "cmd", "title": "更新项目集合仓库", "description": "更新 server/ html/ 内容", "uuid": uuid.uuid4().hex, "detail": { "host": "58.214.239.10", "user": "dgx07", "password": "f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, "cmd": ("cd /data/wrok-flow/dbpai-jenkins" " && rm -fr ./server" " && cp -r ../dbpai-server ./server" " && rm -fr ./html" " && cp -r ../dbpai-web/dist ./html" " && git add ." ' && git commit -m "build"' " && git push"), }, }, { "type": "cmd", "title": "部署", "description": "重新运行 pai-beta 分支 v3.4.2", "uuid": uuid.uuid4().hex, "detail": { "host": "58.214.239.10", "user": "dgx07", "password": "f=bHt/.7kRo:KIXq", "port": 6378, "sudo": True, "cmd": ("cd /data/PAI/dbpai-beta/dbpai-jenkins;" " git stash;" " git pull origin v3.4.2;" " git stash pop;" " docker-compose restart;"), }, }, ], ) await work_content.save() logger.debug("创建") @asynccontextmanager async def lifespan(app: FastAPI): tasks = {} task_nodes = {} app.state.tasks = tasks app.state.task_nodes = task_nodes # 数据库 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=30, misfire_grace_time=10) scheduler.start() logger.info("初始化定时器") await _create_demo_records(app) yield for _, v in tasks.items(): v.cancel() 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( "/api/check_init", tags=["System"], description="系统是否初始化", response_model=InitStatusResponse, ) async def check_init(): return InitStatusResponse(data=await is_init()) @app.post( "/api/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( "/api/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(), is_deleted=0).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( "/api/logout", tags=["User"], description="退出登录", response_model=BaseResponse, ) async def logout(token: str): sessions.pop(token, None) logger.info(sessions) return BaseResponse() @app.get( "/api/current_user", tags=["User"], description="当前用户信息", response_model=LoginResponse, ) async def current_user(user: DB.User = Depends(get_current_user)): return LoginResponse(data=user.json(del_columns=["password"])) @app.post( "/api/run", tags=["Flow"], description="执行命令", response_model=BaseResponse, ) async def flow_run(work_uuid: str): work = await DB.Work.filter(uuid=work_uuid).first() if not work: return BaseResponse(success=False, message="任务不存在") task_id = work_uuid if task_id in app.state.tasks: return BaseResponse(success=False, message="已有任务正在执行...") async def _update_node_status(task_id: str, node_id: str, status: str): task_nodes = app.state.task_nodes logger.info(f"task:{task_id} node:{node_id} status:{status}") if task_id in task_nodes and node_id in task_nodes[task_id]: task_nodes[task_id][node_id] = status async def _task(task_id: str, work: DB.Work): try: work_content = await DB.WorkContent.filter(work_id=work.id).first() if not work_content: logger.warning(f"没有找到对应的内容: {work.uuid}") return nodes = [] params = {} status = {} for x in work_content.content: node_id = x["uuid"] node_type = x["type"] node_info = x["detail"] params[node_id] = node_info status[node_id] = "wait" if node_type == "git": node = F.GitNode() node.set_info(node_id, lambda k, v: _update_node_status(task_id, k, v)) nodes.append(node) elif node_type == "cmd": node = F.CmdNode() node.set_info(node_id, lambda k, v: _update_node_status(task_id, k, v)) nodes.append(node) nodes.pop() for i in range(len(nodes) - 1): nodes[i] - "true" >> nodes[i + 1] flow = F.AsyncFlow(nodes[0]) flow.set_params(params) shared = {} app.state.task_nodes[task_id] = status await flow.run_async(shared) app.state.tasks.pop(task_id, None) app.state.task_nodes.pop(task_id, None) except asyncio.CancelledError: logger.debug(f"任务被取消: {task_id}") app.state.tasks.pop(task_id, None) app.state.tasks[task_id] = asyncio.create_task(_task(task_id, work)) return BaseResponse(data=task_id) @app.post("/api/stop", tags=["Flow"], description="停止执行指令", response_model=BaseResponse) async def stop_run(task_id: str): task = app.state.tasks.pop(task_id, None) app.state.task_nodes.pop(task_id, None) if task: task.cancel() return BaseResponse() @app.get( "/api/flows", tags=["Flow"], description="任务列表", response_model=TableResponse, ) async def task_list(): global_tasks = app.state.tasks global_status = app.state.task_nodes work_list = await DB.Work.filter(is_deleted=0).order_by('-id') work_content_list = await DB.WorkContent.filter(work_id__in=[x.id for x in work_list]).all() res = [] for work in work_list: obj = work.json() content = next((x for x in work_content_list if x.work_id == work.id), None) if content: obj["content"] = [{k: v for k, v in _c.items() if k != "detail"} for _c in (content.content or [])] obj["is_running"] = work.uuid in global_tasks obj["status"] = global_status[work.uuid] if work.uuid in global_status else {} res.append(obj) return TableResponse(data=res, total=len(res)) return app app = create_app() if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)