your-workflow/backend/app/main.py

485 lines
20 KiB
Python

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)