485 lines
20 KiB
Python
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)
|