your-workflow/backend/app/main.py

377 lines
14 KiB
Python

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
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/tasks",
tags=["Flow"],
description="任务列表",
response_model=BaseResponse,
)
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 BaseResponse(data=res)
return app
app = create_app()
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)