refactor: workflow_service改名app
This commit is contained in:
iHeyTang
48
app/database/__init__.py
Normal file
48
app/database/__init__.py
Normal file
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
数据库模块
|
||||
提供工作流和工作流运行的数据库操作接口
|
||||
"""
|
||||
|
||||
from .models import Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer
|
||||
from .api import (
|
||||
init_db,
|
||||
save_workflow,
|
||||
get_all_workflows,
|
||||
get_latest_workflow_by_base_name,
|
||||
get_workflow_by_version,
|
||||
get_workflow,
|
||||
delete_workflow,
|
||||
create_workflow_run,
|
||||
update_workflow_run_status,
|
||||
create_workflow_run_nodes,
|
||||
update_workflow_run_node_status,
|
||||
get_workflow_run,
|
||||
get_workflow_run_nodes,
|
||||
get_pending_workflow_runs,
|
||||
get_running_workflow_runs,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# 模型
|
||||
"Workflow",
|
||||
"WorkflowRun",
|
||||
"WorkflowRunNode",
|
||||
"ComfyUIServer",
|
||||
|
||||
# API函数
|
||||
"init_db",
|
||||
"save_workflow",
|
||||
"get_all_workflows",
|
||||
"get_latest_workflow_by_base_name",
|
||||
"get_workflow_by_version",
|
||||
"get_workflow",
|
||||
"delete_workflow",
|
||||
"create_workflow_run",
|
||||
"update_workflow_run_status",
|
||||
"create_workflow_run_nodes",
|
||||
"update_workflow_run_node_status",
|
||||
"get_workflow_run",
|
||||
"get_workflow_run_nodes",
|
||||
"get_pending_workflow_runs",
|
||||
"get_running_workflow_runs",
|
||||
]
|
||||
282
app/database/api.py
Normal file
282
app/database/api.py
Normal file
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
数据库操作API
|
||||
提供工作流相关的所有数据库操作函数
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Optional, List
|
||||
import json
|
||||
import uuid
|
||||
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy import delete
|
||||
|
||||
from .models import Base, Workflow, WorkflowRun, WorkflowRunNode, ComfyUIServer
|
||||
from .connection import async_engine, AsyncSessionLocal
|
||||
|
||||
|
||||
async def init_db():
|
||||
"""初始化数据库表结构"""
|
||||
async with async_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
print(f"数据库表结构已创建完成。")
|
||||
|
||||
|
||||
async def save_workflow(name: str, workflow_json: str):
|
||||
"""保存工作流"""
|
||||
version = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
workflow = Workflow(
|
||||
name=f"{name} [{version}]",
|
||||
base_name=name,
|
||||
version=version,
|
||||
workflow_json=workflow_json,
|
||||
)
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
await session.merge(workflow) # 使用merge实现INSERT OR REPLACE
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_all_workflows() -> List[dict]:
|
||||
"""获取所有工作流（最新版本）"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
# 使用SQLAlchemy ORM语法，通过子查询获取每个base_name的最新版本
|
||||
from sqlalchemy import func
|
||||
|
||||
# 子查询：获取每个base_name的最新版本
|
||||
latest_versions = (
|
||||
select(Workflow.base_name, func.max(Workflow.version).label("max_version"))
|
||||
.group_by(Workflow.base_name)
|
||||
.subquery()
|
||||
)
|
||||
|
||||
# 主查询：关联获取完整的工作流信息
|
||||
stmt = (
|
||||
select(Workflow)
|
||||
.join(
|
||||
latest_versions,
|
||||
(Workflow.base_name == latest_versions.c.base_name)
|
||||
& (Workflow.version == latest_versions.c.max_version),
|
||||
)
|
||||
.order_by(Workflow.base_name)
|
||||
)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
workflows = result.scalars().all()
|
||||
|
||||
return [
|
||||
{"name": wf.name, "workflow": json.loads(wf.workflow_json)}
|
||||
for wf in workflows
|
||||
]
|
||||
|
||||
|
||||
async def get_latest_workflow_by_base_name(base_name: str) -> Optional[dict]:
|
||||
"""根据基础名称获取最新版本的工作流"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = (
|
||||
select(Workflow)
|
||||
.where(Workflow.base_name == base_name)
|
||||
.order_by(Workflow.version.desc())
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
workflow = result.scalar_one_or_none()
|
||||
|
||||
return workflow.to_dict() if workflow else None
|
||||
|
||||
|
||||
async def get_workflow_by_version(base_name: str, version: str) -> Optional[dict]:
|
||||
"""根据版本获取工作流"""
|
||||
name = f"{base_name} [{version}]"
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = select(Workflow).where(Workflow.name == name)
|
||||
result = await session.execute(stmt)
|
||||
workflow = result.scalar_one_or_none()
|
||||
|
||||
return workflow.to_dict() if workflow else None
|
||||
|
||||
|
||||
async def get_workflow(name: str, version: Optional[str] = None) -> Optional[dict]:
|
||||
"""获取工作流"""
|
||||
if version:
|
||||
return await get_workflow_by_version(name, version)
|
||||
else:
|
||||
return await get_latest_workflow_by_base_name(name)
|
||||
|
||||
|
||||
async def delete_workflow(name: str) -> bool:
|
||||
"""删除工作流"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = delete(Workflow).where(Workflow.name == name)
|
||||
result = await session.execute(stmt)
|
||||
await session.commit()
|
||||
return result.rowcount > 0
|
||||
|
||||
|
||||
async def create_workflow_run(
|
||||
workflow_run_id: str,
|
||||
workflow_name: str,
|
||||
workflow_json: str,
|
||||
api_spec: str,
|
||||
request_data: str,
|
||||
) -> str:
|
||||
"""创建新的工作流运行记录"""
|
||||
workflow_run = WorkflowRun(
|
||||
id=workflow_run_id,
|
||||
workflow_name=workflow_name,
|
||||
workflow_json=workflow_json,
|
||||
api_spec=api_spec,
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
session.add(workflow_run)
|
||||
await session.commit()
|
||||
|
||||
return workflow_run_id
|
||||
|
||||
|
||||
async def update_workflow_run_status(
|
||||
workflow_run_id: str,
|
||||
status: str,
|
||||
server_url: str = None,
|
||||
prompt_id: str = None,
|
||||
client_id: str = None,
|
||||
error_message: str = None,
|
||||
result: str = None,
|
||||
):
|
||||
"""更新工作流运行状态"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
workflow_run = await session.get(WorkflowRun, workflow_run_id)
|
||||
if not workflow_run:
|
||||
return
|
||||
|
||||
workflow_run.status = status
|
||||
|
||||
if status == "running":
|
||||
workflow_run.server_url = server_url
|
||||
workflow_run.prompt_id = prompt_id
|
||||
workflow_run.client_id = client_id
|
||||
workflow_run.started_at = datetime.utcnow()
|
||||
elif status == "completed":
|
||||
workflow_run.completed_at = datetime.utcnow()
|
||||
workflow_run.result = result
|
||||
elif status == "failed":
|
||||
workflow_run.error_message = error_message
|
||||
workflow_run.completed_at = datetime.utcnow()
|
||||
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def create_workflow_run_nodes(workflow_run_id: str, nodes_data: List[dict]):
|
||||
"""创建工作流运行节点记录"""
|
||||
nodes = []
|
||||
for node in nodes_data:
|
||||
node_obj = WorkflowRunNode(
|
||||
id=str(uuid.uuid4()),
|
||||
workflow_run_id=workflow_run_id,
|
||||
node_id=node["id"],
|
||||
node_type=node["type"],
|
||||
)
|
||||
nodes.append(node_obj)
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
session.add_all(nodes)
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def update_workflow_run_node_status(
|
||||
workflow_run_id: str,
|
||||
node_id: str,
|
||||
status: str,
|
||||
output_data: str = None,
|
||||
error_message: str = None,
|
||||
):
|
||||
"""更新工作流运行节点状态"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = select(WorkflowRunNode).where(
|
||||
WorkflowRunNode.workflow_run_id == workflow_run_id,
|
||||
WorkflowRunNode.node_id == node_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
node = result.scalar_one_or_none()
|
||||
|
||||
if not node:
|
||||
return
|
||||
|
||||
node.status = status
|
||||
|
||||
if status == "running":
|
||||
node.started_at = datetime.utcnow()
|
||||
elif status == "completed":
|
||||
node.output_data = output_data
|
||||
node.completed_at = datetime.utcnow()
|
||||
elif status == "failed":
|
||||
node.error_message = error_message
|
||||
node.completed_at = datetime.utcnow()
|
||||
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_workflow_run(workflow_run_id: str) -> Optional[WorkflowRun]:
|
||||
"""获取工作流运行记录"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
workflow_run = await session.get(WorkflowRun, workflow_run_id)
|
||||
return workflow_run
|
||||
|
||||
|
||||
async def get_workflow_run_nodes(workflow_run_id: str) -> List[dict]:
|
||||
"""获取工作流运行节点记录"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = select(WorkflowRunNode).where(
|
||||
WorkflowRunNode.workflow_run_id == workflow_run_id
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
nodes = result.scalars().all()
|
||||
|
||||
return [node.to_dict() for node in nodes]
|
||||
|
||||
|
||||
async def get_pending_workflow_runs() -> List[dict]:
|
||||
"""获取所有待处理的工作流运行记录"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = (
|
||||
select(WorkflowRun)
|
||||
.where(WorkflowRun.status == "pending")
|
||||
.order_by(WorkflowRun.created_at.asc())
|
||||
)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
runs = result.scalars().all()
|
||||
|
||||
return [run.to_dict() for run in runs]
|
||||
|
||||
|
||||
async def get_running_workflow_runs() -> List[dict]:
|
||||
"""获取所有正在运行的工作流运行记录"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = select(WorkflowRun).where(WorkflowRun.status == "running")
|
||||
result = await session.execute(stmt)
|
||||
runs = result.scalars().all()
|
||||
|
||||
return [run.to_dict() for run in runs]
|
||||
|
||||
|
||||
async def get_workflow_runs_recent(
|
||||
start_time: datetime, end_time: datetime
|
||||
) -> List[dict]:
|
||||
"""获取指定时间范围内的最近工作流运行记录"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
stmt = (
|
||||
select(WorkflowRun)
|
||||
.where(
|
||||
WorkflowRun.created_at >= start_time, WorkflowRun.created_at <= end_time
|
||||
)
|
||||
.order_by(WorkflowRun.created_at.desc())
|
||||
)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
runs = result.scalars().all()
|
||||
|
||||
return [run.to_dict() for run in runs]
|
||||
42
app/database/connection.py
Normal file
42
app/database/connection.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
数据库连接配置
|
||||
管理数据库引擎、会话和连接池
|
||||
"""
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from app.config import Settings
|
||||
|
||||
# 数据库配置
|
||||
DATABASE_FILE = Settings().DB_FILE
|
||||
ASYNC_DATABASE_URL = f"sqlite+aiosqlite:///{DATABASE_FILE}"
|
||||
SYNC_DATABASE_URL = f"sqlite:///{DATABASE_FILE}"
|
||||
|
||||
# 创建异步引擎
|
||||
async_engine = create_async_engine(
|
||||
ASYNC_DATABASE_URL, echo=False, pool_pre_ping=True, pool_recycle=3600
|
||||
)
|
||||
|
||||
# 创建异步会话工厂
|
||||
AsyncSessionLocal = async_sessionmaker(
|
||||
async_engine, class_=AsyncSession, expire_on_commit=False
|
||||
)
|
||||
|
||||
# 创建同步引擎和会话（用于初始化）
|
||||
sync_engine = create_engine(
|
||||
SYNC_DATABASE_URL, echo=False, pool_pre_ping=True, pool_recycle=3600
|
||||
)
|
||||
SessionLocal = sessionmaker(bind=sync_engine)
|
||||
|
||||
|
||||
async def get_db_session() -> AsyncSession:
|
||||
"""获取数据库会话"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
return session
|
||||
|
||||
|
||||
def get_sync_session():
|
||||
"""获取同步数据库会话（用于初始化等场景）"""
|
||||
return SessionLocal()
|
||||
144
app/database/models.py
Normal file
144
app/database/models.py
Normal file
@@ -0,0 +1,144 @@
|
||||
"""
|
||||
数据库模型定义
|
||||
定义工作流相关的所有数据表结构
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Integer, Boolean
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import relationship
|
||||
|
||||
# 创建基类
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class Workflow(Base):
|
||||
"""工作流模型"""
|
||||
|
||||
__tablename__ = "workflows"
|
||||
|
||||
name = Column(String, primary_key=True)
|
||||
base_name = Column(String, nullable=False)
|
||||
version = Column(String, nullable=False)
|
||||
workflow_json = Column(Text, nullable=False)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"name": self.name,
|
||||
"base_name": self.base_name,
|
||||
"version": self.version,
|
||||
"workflow_json": self.workflow_json,
|
||||
}
|
||||
|
||||
|
||||
class WorkflowRun(Base):
|
||||
"""工作流运行记录模型"""
|
||||
|
||||
__tablename__ = "workflow_run"
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
workflow_name = Column(String, nullable=False)
|
||||
prompt_id = Column(String)
|
||||
client_id = Column(String)
|
||||
status = Column(String, nullable=False, default="pending")
|
||||
server_url = Column(String)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
started_at = Column(DateTime)
|
||||
completed_at = Column(DateTime)
|
||||
error_message = Column(Text)
|
||||
workflow_json = Column(Text, nullable=False)
|
||||
api_spec = Column(Text, nullable=False)
|
||||
request_data = Column(Text, nullable=False)
|
||||
result = Column(Text)
|
||||
|
||||
# 关联关系
|
||||
nodes = relationship(
|
||||
"WorkflowRunNode", back_populates="workflow_run", cascade="all, delete-orphan"
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"id": self.id,
|
||||
"workflow_name": self.workflow_name,
|
||||
"prompt_id": self.prompt_id,
|
||||
"client_id": self.client_id,
|
||||
"status": self.status,
|
||||
"server_url": self.server_url,
|
||||
"created_at": self.created_at.isoformat() if self.created_at else None,
|
||||
"started_at": self.started_at.isoformat() if self.started_at else None,
|
||||
"completed_at": (
|
||||
self.completed_at.isoformat() if self.completed_at else None
|
||||
),
|
||||
"error_message": self.error_message,
|
||||
"workflow_json": self.workflow_json,
|
||||
"api_spec": self.api_spec,
|
||||
"request_data": self.request_data,
|
||||
"result": self.result,
|
||||
}
|
||||
|
||||
|
||||
class WorkflowRunNode(Base):
|
||||
"""工作流运行节点模型"""
|
||||
|
||||
__tablename__ = "workflow_run_nodes"
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
workflow_run_id = Column(String, ForeignKey("workflow_run.id"), nullable=False)
|
||||
node_id = Column(String, nullable=False)
|
||||
node_type = Column(String, nullable=False)
|
||||
status = Column(String, nullable=False, default="pending")
|
||||
started_at = Column(DateTime)
|
||||
completed_at = Column(DateTime)
|
||||
output_data = Column(Text)
|
||||
error_message = Column(Text)
|
||||
|
||||
# 关联关系
|
||||
workflow_run = relationship("WorkflowRun", back_populates="nodes")
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"id": self.id,
|
||||
"workflow_run_id": self.workflow_run_id,
|
||||
"node_id": self.node_id,
|
||||
"node_type": self.node_type,
|
||||
"status": self.status,
|
||||
"started_at": self.started_at.isoformat() if self.started_at else None,
|
||||
"completed_at": (
|
||||
self.completed_at.isoformat() if self.completed_at else None
|
||||
),
|
||||
"output_data": self.output_data,
|
||||
"error_message": self.error_message,
|
||||
}
|
||||
|
||||
|
||||
class ComfyUIServer(Base):
|
||||
"""ComfyUI 服务器模型"""
|
||||
|
||||
__tablename__ = "comfyui_servers"
|
||||
|
||||
name = Column(String, primary_key=True)
|
||||
http_url = Column(String, nullable=False)
|
||||
ws_url = Column(String, nullable=False)
|
||||
status = Column(String, nullable=False, default="offline")
|
||||
last_health_check = Column(DateTime)
|
||||
current_tasks = Column(Integer, default=0)
|
||||
max_concurrent_tasks = Column(Integer, default=1)
|
||||
capabilities = Column(Text) # JSON 字符串
|
||||
server_metadata = Column(Text) # JSON 字符串
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"name": self.name,
|
||||
"http_url": self.http_url,
|
||||
"ws_url": self.ws_url,
|
||||
"status": self.status,
|
||||
"last_health_check": self.last_health_check.isoformat() if self.last_health_check else None,
|
||||
"current_tasks": self.current_tasks,
|
||||
"max_concurrent_tasks": self.max_concurrent_tasks,
|
||||
"capabilities": self.capabilities,
|
||||
"server_metadata": self.server_metadata,
|
||||
"created_at": self.created_at.isoformat() if self.created_at else None,
|
||||
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
||||
}
|
||||
Reference in New Issue
Block a user