274 lines
9.8 KiB
Python
274 lines
9.8 KiB
Python
import asyncio
|
|
import uvicorn
|
|
import uuid
|
|
import time
|
|
import models.db as DB
|
|
import nodes as F
|
|
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
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
tasks = {}
|
|
app.state.tasks = tasks
|
|
# 数据库
|
|
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("初始化定时器")
|
|
|
|
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(id: int):
|
|
|
|
task_id = uuid.uuid4().hex
|
|
|
|
async def _task(task_id):
|
|
try:
|
|
nodes = []
|
|
node_0 = F.GitNode()
|
|
node_0.set_id("01K3JHSEGBHHAVQ92RAH8M6GPX")
|
|
nodes.append(node_0)
|
|
node_1 = F.GitNode()
|
|
node_1.set_id("01K3JHSHKESMWYEGSY4VNCJY4D")
|
|
nodes.append(node_1)
|
|
node_2 = F.GitNode()
|
|
node_2.set_id("01K3T4GTWZJF3RT5FFH25D3T9C")
|
|
nodes.append(node_2)
|
|
node_3 = F.CmdNode()
|
|
node_3.set_id("01K3T4S2HN97T1B0TPFBCFSB00")
|
|
nodes.append(node_3)
|
|
|
|
logger.debug(app.state.tasks)
|
|
|
|
# 设置依赖关系: node0 - "true" >> node1 - "true" >> node2
|
|
for i in range(len(nodes) - 1):
|
|
nodes[i] - "true" >> nodes[i + 1]
|
|
|
|
flow = F.AsyncFlow(node_0)
|
|
flow.set_params(
|
|
{
|
|
"01K3JHSEGBHHAVQ92RAH8M6GPX": {
|
|
"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,
|
|
},
|
|
"01K3JHSHKESMWYEGSY4VNCJY4D": {
|
|
"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,
|
|
},
|
|
"01K3T4GTWZJF3RT5FFH25D3T9C": {
|
|
"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,
|
|
},
|
|
"01K3T4S2HN97T1B0TPFBCFSB00": {
|
|
"host": "58.214.239.10",
|
|
"user": "dgx07",
|
|
"password": "f=bHt/.7kRo:KIXq",
|
|
"port": 6378,
|
|
"sudo": True,
|
|
"cmd": "\n".join(
|
|
[
|
|
"cd /data/wrok-flow/dbpai-jenkins",
|
|
"git log",
|
|
]
|
|
),
|
|
},
|
|
}
|
|
)
|
|
shared = {}
|
|
await flow.run_async(shared)
|
|
app.state.tasks.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))
|
|
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)
|
|
if task:
|
|
task.cancel()
|
|
return BaseResponse()
|
|
|
|
return app
|
|
|
|
|
|
app = create_app()
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)
|