your-workflow/backend/app/main.py

248 lines
8.0 KiB
Python

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):
# 数据库
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=5)
scheduler.start()
logger.info("初始化定时器")
yield
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(
"/check_init",
tags=["System"],
description="系统是否初始化",
response_model=InitStatusResponse,
)
async def check_init():
return InitStatusResponse(data=await is_init())
@app.post(
"/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(
"/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()
).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(
"/logout", tags=["User"], description="退出登录", response_model=BaseResponse
)
async def logout(token: str):
sessions.pop(token, None)
logger.info(sessions)
return BaseResponse()
@app.post(
"/run", tags=["Flow"], description="执行命令", response_model=BaseResponse
)
async def flow_run(id: int):
nodes = []
local_node = F.CmdNode()
local_node.set_id("0")
nodes.append(local_node)
a800_node = F.CmdNode()
a800_node.set_id("1")
nodes.append(a800_node)
a800_git_node = F.GitNode()
a800_git_node.set_id("2")
nodes.append(a800_git_node)
# 设置依赖关系
# node0 - "true" >> node1 - "true" >> node2
for i in range(len(nodes) -1):
# nodes[i] - "error" >> nodes[i + 1]
pass
flow = F.AsyncFlow(a800_git_node)
flow.set_params(
{
"0": {
"cmd": """
echo "1) 前端 - 拉取新代码"
cd /Users/liruwei/Documents/code/project/dabai/dbpai-web
git stash
git pull
echo "2) 后端 - 拉取新代码"
cd /Users/liruwei/Documents/code/project/dabai/dbpai-server
git stash
echo "3) 后端 - 更新 dbpai-chat-widget.min.js"
cp /Users/liruwei/Documents/code/project/dabai/dbpai-web/public/dbpai-chat-widget.min.js ./src/static
git commit -am "feat: dbpai-chat-widget.min.js"
git push
echo "4) 打包"
cd ..
sh build-dbpai-jenkisn.sh all
cd dbpai-jenkins
git commit -am "build"
git push
echo "5) 还原"
cd /Users/liruwei/Documents/code/project/dabai/dbpai-web
git stash pop
cd /Users/liruwei/Documents/code/project/dabai/dbpai-server
git stash pop
"""
},
"1": {
"cmd": """
cd /data/PAI/dbpai-beta/dbpai-jenkins
git stash
git pull origin v3.4.2
git stash pop
docker logs ai-guard-server -n 50
""",
"host": "58.214.239.10",
"user": "dgx07",
"password": "1f=bHt/.7kRo:KIXq",
"port": 6378,
"sudo": True,
},
"2": {
"host": "58.214.239.10",
"user": "dgx07",
"password": "f=bHt/.7kRo:KIXq",
"port": 6378,
"sudo": True,
"git_path": "/data/PAI/demo/dbpai-web",
"git_user": "liruwei",
"git_password": "Wsad.8246",
"git_url": "https://tencentgit.dabby.com.cn/dbpai/dbpai-web.git",
"git_branch": "v3.4.2",
"git_clone_type": "http/https",
"git_clone_depth": 0,
},
}
)
shared = {}
await flow.run_async(shared)
logger.info(shared)
return BaseResponse()
return app
app = create_app()
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)