import asyncio import datetime import json import os import uvicorn import uuid import time import traceback import models.db as DB import nodes as F from tool import format_content_uuid from tortoise.transactions import in_transaction from tortoise.functions import Max from fastapi import FastAPI, Depends, HTTPException, Query from models.schemas.system import InitStatusResponse, InitConfigRequest, BaseResponse, TableResponse from models.schemas.user import LoginRequest, LoginResponse from models.schemas.client import ClientRequest from models.schemas.work import WorkRequest 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 fastapi.responses import StreamingResponse 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 save_session_to_file(app: FastAPI): if sessions: login_json = {} for k, v in sessions.items(): login_json[k] = [v[0], v[1].id] logger.debug(login_json) with open(".login.json", "w", encoding="utf-8") as f: json.dump(login_json, f, indent=4, ensure_ascii=False) async def load_session_from_file(app: FastAPI): if os.path.exists(".login.json"): with open(".login.json", "r", encoding="utf-8") as f: login_json = json.load(f) if login_json: users = await DB.User.filter(is_deleted=0).all() for k, v in login_json.items(): u = next((u for u in users if u.id == v[1]), None) if u: sessions[k] = v[0], u logger.debug(sessions) 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 def pagination_query(page:int = Query(1), page_size:int = Query(20)): return {"page": page, "page_size": page_size} 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): 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("初始化数据库") # 读取上一次登录状态 await load_session_from_file(app) # 定时器 scheduler.add_job(clear_session, "interval", seconds=30, misfire_grace_time=10) scheduler.start() logger.info("初始化定时器") yield for _, v in tasks.items(): v.cancel() scheduler.shutdown() logger.info("关闭定时器") await close_tortoise() logger.info("关闭数据库") save_session_to_file(app) 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)): if not await is_init(): return LoginResponse(success=False) 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, user: DB.User = Depends(get_current_user)): 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 _append_host_config(detail: dict): result = {**detail} client = await DB.Client.filter(id=result.get("client", 0)).first() if client: result["host"] = client.host result["user"] = client.user result["password"] = client.password result["port"] = client.port return result 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, user: DB.User): work_log = DB.WorkLogs( user_id=user.id, work_id=work.id, status=DB.WorkLogsStatus.RUNNING, work_title=work.work_title, work_desc=work.work_desc, ) await work_log.save() try: start_time = int(datetime.datetime.now().timestamp()) work_content = await DB.WorkContent.filter(work_id=work.id).first() if not work_content: logger.warning(f"没有找到对应的内容: {work.uuid}") return # 找出从 “开始” 到最后的节点 node_dict = {x["uuid"]: x for x in work_content.content} node_list = [] current = "1" while current in node_dict: t = node_dict[current] node_list.append(t) if "edge" in t and "target" in t["edge"]: current = t["edge"]["target"] else: break # pocket flow 节点创建 nodes = [] params = {} shared = {} status = {} for x in node_list: node_id = x["uuid"] node_type = x["type"] if node_id == "1": continue params[node_id] = await _append_host_config(x["detail"]) status[node_id] = "wait" node_class_map = {"git": F.GitNode, "cmd": F.CmdNode, "status": F.StatusNode} node_cls = node_class_map[node_type] node = node_cls() node.set_info(node_id, lambda k, v: _update_node_status(task_id, k, v)) nodes.append(node) for i in range(len(nodes) - 1): nodes[i] - "true" >> nodes[i + 1] flow = F.AsyncFlow(nodes[0]) flow.set_params(params) 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) end_time = int(datetime.datetime.now().timestamp()) # 更新日志 work_log_result = [] for x in nodes: id = x.id node_info = node_dict[id] shared_info = shared.get(id) status = "" if shared_info and shared_info.get("err"): status = "error" elif shared_info: status = "success" work_log_result.append({ "id": id, "title": node_info.get("title", ""), "result": shared.get(id, {}), "status": status }) work_log.result = {"time": end_time - start_time, "logs": work_log_result} if next((x for x in shared.values() if x.get("err")), None): work_log.status = DB.WorkLogsStatus.FAILED else: work_log.status = DB.WorkLogsStatus.SUCCESS await work_log.save() except asyncio.CancelledError: work_log.status = DB.WorkLogsStatus.CANCELLED await work_log.save() logger.debug(f"任务被取消: {task_id}") app.state.tasks.pop(task_id, None) except Exception as e: traceback.print_exc() logger.error(e) work_log.status = DB.WorkLogsStatus.FAILED await work_log.save() app.state.tasks.pop(task_id, None) app.state.tasks[task_id] = asyncio.create_task(_task(task_id, work, user)) return BaseResponse(data=task_id) @app.post("/api/stop", tags=["Flow"], description="停止任务", response_model=BaseResponse) async def stop_run(work_uuid: str): task = app.state.tasks.pop(work_uuid, None) app.state.task_nodes.pop(work_uuid, None) if task: task.cancel() return BaseResponse() @app.get("/api/flows", tags=["Flow"], description="任务列表", response_model=TableResponse) async def flows_list(user: DB.User = Depends(get_current_user)): global_tasks = app.state.tasks global_status = app.state.task_nodes db_filter = { "is_deleted": 0 } if user.level != 0: work_ids = await DB.WorkUserRelation.filter(user_id=user.id).values_list("work_id", flat=True) db_filter["id__in": work_ids if work_ids else [0]] work_list = await DB.Work.filter(**db_filter).order_by("-id") work_ids = [x.id for x in work_list] work_content_list = await DB.WorkContent.filter(work_id__in=work_ids).all() sub_logs = await DB.WorkLogs.filter(work_id__in=work_ids).group_by('work_id').annotate(max_id=Max("id")).values_list("max_id", flat=True) work_logs_list = await DB.WorkLogs.filter(id__in=sub_logs).all() user_list = await DB.User.filter(id__in=list(set([x.user_id for x in work_content_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 [])] log: DB.WorkLogs = next((x for x in work_logs_list if x.work_id == work.id), DB.WorkLogs(user_id=-1)) log_user: DB.User = next((u for u in user_list if u.id == log.user_id), DB.User()) obj["last_user"] = log_user.username or "" obj["last_use_time"] = log.result.get("time", 0) obj["last_run_time"] = log.create_time.timestamp() if log.create_time else 0 obj["last_status"] = log.status 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)) @app.post("/api/flows", tags=["Flow"], description="任务创建", response_model=BaseResponse) async def create_flow(body: WorkRequest, user: DB.User = Depends(get_current_user)): work = DB.Work() work.uuid = uuid.uuid4().hex work.creator_id = user.id work.work_title = body.work_title work.work_desc = body.work_desc work_content = DB.WorkContent( user_id=work.creator_id, content=[], is_deleted=0, work_id=0, ) work_content.content = format_content_uuid(body.content) work_content.update_time = datetime.datetime.now() async with in_transaction(): await work.save() work_content.work_id = work.id await work_content.save() return BaseResponse(data=work.uuid) @app.put("/api/flows", tags=["Flow"], description="任务修改", response_model=BaseResponse) async def update_flow(work_uuid: str, body: WorkRequest, user: DB.User = Depends(get_current_user)): work = await DB.Work.filter(uuid=work_uuid).first() if not work: return BaseResponse(success=False, message="Not Found") work.work_title = body.work_title work.work_desc = body.work_desc work.update_time = datetime.datetime.now() work_content = await DB.WorkContent.filter(work_id=work.id).first() if not work_content: work_content = DB.WorkContent( user_id=work.creator_id, content=[], is_deleted=0, work_id=work.id, ) work_content.content = format_content_uuid(body.content) work_content.update_time = datetime.datetime.now() async with in_transaction(): await work.save() await work_content.save() return BaseResponse() @app.get("/api/flows/detail", tags=["Flow"], description="任务详情", response_model=BaseResponse) async def flows_detail(work_uuid: str, user: DB.User = Depends(get_current_user)): work = await DB.Work.filter(uuid=work_uuid).first() if not work: return BaseResponse(success=False, message="not found") work_content = await DB.WorkContent.filter(work_id=work.id).first() res = work.json() res["content"] = work_content.json() return BaseResponse(success=True, data=res) @app.get("/api/flows/status", tags=["Flow"], description="获取任务状态") async def flows_status(token: str = ""): # 如果出现长时间链接被中断,可以设置 nginx 的 proxy_read_timeout async def _event(): while True: try: if not token or token not in sessions: data = json.dumps({"code": 403}) yield f"data: {data}\n\n" break data = json.dumps(app.state.task_nodes) yield f"data: {data}\n\n" await asyncio.sleep(3) except: break return StreamingResponse(_event(), media_type="text/event-stream") @app.delete("/api/flows", tags=["Flow"], description="任务删除", response_model=BaseResponse) async def del_client(work_uuid: str, user: DB.User = Depends(get_current_user)): work = await DB.Work.filter(uuid=work_uuid).first() if not work: return BaseResponse() work_content = await DB.WorkContent.filter(work_id=work.id).first() async with in_transaction(): work.is_deleted = 1 work.update_time = datetime.datetime.now() await work.save() if work_content: work_content.is_deleted = 1 work_content.update_time = datetime.datetime.now() await work_content.save() return BaseResponse() @app.get("/api/clients", tags=["Client"], description="终端列表", response_model=TableResponse) async def clients_list(user: DB.User = Depends(get_current_user)): res = await DB.Client.filter(user_id=user.id).all() return TableResponse(data=[x.json() for x in res], total=len(res)) @app.post("/api/clients", tags=["Client"], description="终端创建", response_model=BaseResponse) async def create_client(body: ClientRequest, user: DB.User = Depends(get_current_user)): await DB.Client(**body.model_dump(), user_id=user.id).save() return BaseResponse() @app.post("/api/clients", tags=["Client"], description="终端创建", response_model=BaseResponse) async def create_client(body: ClientRequest, user: DB.User = Depends(get_current_user)): await DB.Client(**body.model_dump, user_id=user.id).save() return BaseResponse() @app.put("/api/clients", tags=["Client"], description="终端修改", response_model=BaseResponse) async def update_client(id: int, body: ClientRequest, user: DB.User = Depends(get_current_user)): obj = await DB.Client.filter(id=id).first() if obj: obj.update_from_dict(body.model_dump()) obj.update_time = datetime.datetime.now() await obj.save() return BaseResponse() @app.delete("/api/clients", tags=["Client"], description="终端删除", response_model=BaseResponse) async def del_client(id: int, user: DB.User = Depends(get_current_user)): await DB.Client.filter(id=id).delete() return BaseResponse() @app.get("/api/logs", tags=["Log"], description="日志列表", response_model=TableResponse) async def log_list(p:dict = Depends(pagination_query), user: DB.User = Depends(get_current_user)): page, page_size = p["page"], p["page_size"] users = await DB.User.all() res = await DB.WorkLogs.all().order_by('-id').limit(page_size).offset((page - 1) * page_size) count = await DB.WorkLogs.all().count() data = [] for l in res: u = next((x for x in users if x.id == l.user_id), None) obj = l.json() if u: obj["user"] = u.username data.append(obj) return TableResponse(data=data, total=count) return app app = create_app() if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)