From 58ce862567704ef1587150f1a2d0aa182c647492 Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Thu, 14 Aug 2025 12:02:01 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=BA=93=E7=9B=B8=E5=85=B3=E5=8A=9F=E8=83=BD=EF=BC=8C?= =?UTF-8?q?=E7=A7=BB=E9=99=A4=E6=97=A7=E7=9A=84=E6=95=B0=E6=8D=AE=E5=BA=93?= =?UTF-8?q?=E6=A8=A1=E5=9D=97=EF=BC=8C=E6=9B=B4=E6=96=B0API=E8=B7=AF?= =?UTF-8?q?=E7=94=B1=E4=BB=A5=E4=BD=BF=E7=94=A8=E6=96=B0=E7=9A=84=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=BA=93=E6=8E=A5=E5=8F=A3=EF=BC=8C=E5=A2=9E=E5=BC=BA?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E7=BB=93=E6=9E=84=E5=92=8C=E5=8F=AF=E7=BB=B4?= =?UTF-8?q?=E6=8A=A4=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- requirements.txt | 3 +- workflow_service/comfyui_client.py | 7 +- workflow_service/database.py | 253 ----------------- workflow_service/database/__init__.py | 47 ++++ workflow_service/database/api.py | 256 ++++++++++++++++++ workflow_service/database/connection.py | 42 +++ workflow_service/database/models.py | 111 ++++++++ workflow_service/main.py | 4 +- workflow_service/routes/run.py | 8 +- .../routes/runx/model_with_multi_dress.py | 5 +- workflow_service/routes/workflow.py | 15 +- 11 files changed, 479 insertions(+), 272 deletions(-) delete mode 100644 workflow_service/database.py create mode 100644 workflow_service/database/__init__.py create mode 100644 workflow_service/database/api.py create mode 100644 workflow_service/database/connection.py create mode 100644 workflow_service/database/models.py diff --git a/requirements.txt b/requirements.txt index c41e7c2..0d66211 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ pydantic-settings websockets aiohttp boto3 -aiosqlite \ No newline at end of file +aiosqlite +sqlalchemy \ No newline at end of file diff --git a/workflow_service/comfyui_client.py b/workflow_service/comfyui_client.py index 12a03c5..24a0d1b 100644 --- a/workflow_service/comfyui_client.py +++ b/workflow_service/comfyui_client.py @@ -6,22 +6,19 @@ import os import random import uuid import websockets -from collections import defaultdict from datetime import datetime -from typing import Dict, Any, Optional, List, Set, Union +from typing import Dict, Any import aiohttp from aiohttp import ClientTimeout from workflow_service.config import Settings, ComfyUIServer -from workflow_service.database import ( +from workflow_service.database.api import ( create_workflow_run, update_workflow_run_status, create_workflow_run_nodes, update_workflow_run_node_status, get_workflow_run, - get_pending_workflow_runs, - get_running_workflow_runs, get_workflow_run_nodes, ) diff --git a/workflow_service/database.py b/workflow_service/database.py deleted file mode 100644 index ded3de0..0000000 --- a/workflow_service/database.py +++ /dev/null @@ -1,253 +0,0 @@ -from datetime import datetime -from typing import Optional -import aiosqlite -import json -import re -import uuid - -from workflow_service.config import Settings - -DATABASE_FILE = Settings().DB_FILE - -async def init_db(): - async with aiosqlite.connect(DATABASE_FILE) as db: - await db.execute(""" - CREATE TABLE IF NOT EXISTS workflows ( - name TEXT PRIMARY KEY, - base_name TEXT NOT NULL, - version TEXT NOT NULL, - workflow_json TEXT NOT NULL - ) - """) - - # 新增 workflow_run 表 - await db.execute(""" - CREATE TABLE IF NOT EXISTS workflow_run ( - id TEXT PRIMARY KEY, - workflow_name TEXT NOT NULL, - prompt_id TEXT, - client_id TEXT, - status TEXT NOT NULL DEFAULT 'pending', - server_url TEXT, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - started_at TIMESTAMP, - completed_at TIMESTAMP, - error_message TEXT, - workflow_json TEXT NOT NULL, - api_spec TEXT NOT NULL, - request_data TEXT NOT NULL, - result TEXT - ) - """) - - # 新增 workflow_run_nodes 表，记录每个节点的运行状态 - await db.execute(""" - CREATE TABLE IF NOT EXISTS workflow_run_nodes ( - id TEXT PRIMARY KEY, - workflow_run_id TEXT NOT NULL, - node_id TEXT NOT NULL, - node_type TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - started_at TIMESTAMP, - completed_at TIMESTAMP, - output_data TEXT, - error_message TEXT, - FOREIGN KEY (workflow_run_id) REFERENCES workflow_run (id) - ) - """) - - await db.commit() - print(f"数据库 '{DATABASE_FILE}' 已准备就绪。") - -async def save_workflow(name: str, workflow_json: str): - version = datetime.now().strftime("%Y%m%d%H%M%S") - async with aiosqlite.connect(DATABASE_FILE) as db: - await db.execute( - "INSERT OR REPLACE INTO workflows (name, base_name, version, workflow_json) VALUES (?, ?, ?, ?)", - (f"{name} [{version}]", name, version, workflow_json) - ) - await db.commit() - -async def get_all_workflows() -> list[dict]: - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - # 按 base_name 分组，获取每个基础名称的最新版本记录 - cursor = await db.execute(""" - SELECT w1.name, w1.workflow_json, w1.base_name, w1.version - FROM workflows w1 - INNER JOIN ( - SELECT base_name, MAX(version) as max_version - FROM workflows - GROUP BY base_name - ) w2 ON w1.base_name = w2.base_name AND w1.version = w2.max_version - ORDER BY w1.base_name - """) - rows = await cursor.fetchall() - return [{"name": row["name"], "workflow": json.loads(row["workflow_json"])} for row in rows] - -async def get_latest_workflow_by_base_name(base_name: str) -> dict | None: - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute("SELECT * FROM workflows WHERE base_name = ? ORDER BY version DESC LIMIT 1", (base_name,)) - row = await cursor.fetchone() - return dict(row) if row else None - -async def get_workflow_by_version(base_name: str, version: str) -> dict | None: - name = f"{base_name} [{version}]" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute("SELECT * FROM workflows WHERE name = ?", (name,)) - row = await cursor.fetchone() - return dict(row) if row else None - -async def get_workflow(name: str, version: Optional[str] = None) -> dict | None: - 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 aiosqlite.connect(DATABASE_FILE) as db: - cursor = await db.execute("DELETE FROM workflows WHERE name = ?", (name,)) - await db.commit() - return cursor.rowcount > 0 - -# 新增 workflow_run 相关函数 -async def create_workflow_run( - workflow_run_id: str, - workflow_name: str, - workflow_json: str, - api_spec: str, - request_data: str -) -> str: - """创建新的工作流运行记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - await db.execute(""" - INSERT INTO workflow_run (id, workflow_name, workflow_json, api_spec, request_data) - VALUES (?, ?, ?, ?, ?) - """, (workflow_run_id, workflow_name, workflow_json, api_spec, request_data)) - await db.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 aiosqlite.connect(DATABASE_FILE) as db: - if status == 'running': - await db.execute(""" - UPDATE workflow_run - SET status = ?, server_url = ?, prompt_id = ?, client_id = ?, started_at = CURRENT_TIMESTAMP - WHERE id = ? - """, (status, server_url, prompt_id, client_id, workflow_run_id)) - elif status == 'completed': - await db.execute(""" - UPDATE workflow_run - SET status = ?, completed_at = CURRENT_TIMESTAMP, result = ? - WHERE id = ? - """, (status, result, workflow_run_id)) - elif status == 'failed': - await db.execute(""" - UPDATE workflow_run - SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP - WHERE id = ? - """, (status, error_message, workflow_run_id)) - else: - await db.execute(""" - UPDATE workflow_run - SET status = ? - WHERE id = ? - """, (status, workflow_run_id)) - await db.commit() - -async def create_workflow_run_nodes(workflow_run_id: str, nodes_data: list): - """创建工作流运行节点记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - for node in nodes_data: - node_id = str(uuid.uuid4()) - await db.execute(""" - INSERT INTO workflow_run_nodes (id, workflow_run_id, node_id, node_type) - VALUES (?, ?, ?, ?) - """, (node_id, workflow_run_id, node["id"], node["type"])) - await db.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 aiosqlite.connect(DATABASE_FILE) as db: - if status == 'running': - await db.execute(""" - UPDATE workflow_run_nodes - SET status = ?, started_at = CURRENT_TIMESTAMP - WHERE workflow_run_id = ? AND node_id = ? - """, (status, workflow_run_id, node_id)) - elif status == 'completed': - await db.execute(""" - UPDATE workflow_run_nodes - SET status = ?, output_data = ?, completed_at = CURRENT_TIMESTAMP - WHERE workflow_run_id = ? AND node_id = ? - """, (status, output_data, workflow_run_id, node_id)) - elif status == 'failed': - await db.execute(""" - UPDATE workflow_run_nodes - SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP - WHERE workflow_run_id = ? AND node_id = ? - """, (status, error_message, workflow_run_id, node_id)) - else: - await db.execute(""" - UPDATE workflow_run_nodes - SET status = ? - WHERE workflow_run_id = ? AND node_id = ? - """, (status, workflow_run_id, node_id)) - await db.commit() - -async def get_workflow_run(workflow_run_id: str) -> dict | None: - """获取工作流运行记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute(""" - SELECT * FROM workflow_run WHERE id = ? - """, (workflow_run_id,)) - row = await cursor.fetchone() - return dict(row) if row else None - -async def get_workflow_run_nodes(workflow_run_id: str) -> list[dict]: - """获取工作流运行节点记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute(""" - SELECT * FROM workflow_run_nodes WHERE workflow_run_id = ? - """, (workflow_run_id,)) - rows = await cursor.fetchall() - return [dict(row) for row in rows] - -async def get_pending_workflow_runs() -> list[dict]: - """获取所有待处理的工作流运行记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute(""" - SELECT * FROM workflow_run WHERE status = 'pending' ORDER BY created_at ASC - """) - rows = await cursor.fetchall() - return [dict(row) for row in rows] - -async def get_running_workflow_runs() -> list[dict]: - """获取所有正在运行的工作流运行记录""" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute(""" - SELECT * FROM workflow_run WHERE status = 'running' - """) - rows = await cursor.fetchall() - return [dict(row) for row in rows] \ No newline at end of file diff --git a/workflow_service/database/__init__.py b/workflow_service/database/__init__.py new file mode 100644 index 0000000..2378848 --- /dev/null +++ b/workflow_service/database/__init__.py @@ -0,0 +1,47 @@ +""" +数据库模块 +提供工作流和工作流运行的数据库操作接口 +""" + +from .models import Workflow, WorkflowRun, WorkflowRunNode +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", + + # 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", +] diff --git a/workflow_service/database/api.py b/workflow_service/database/api.py new file mode 100644 index 0000000..3396344 --- /dev/null +++ b/workflow_service/database/api.py @@ -0,0 +1,256 @@ +""" +数据库操作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 +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: + # 使用子查询获取每个base_name的最新版本 + stmt = ( + select(Workflow) + .where( + Workflow.version + == select(Workflow.version) + .where(Workflow.base_name == Workflow.base_name) + .order_by(Workflow.version.desc()) + .limit(1) + .scalar_subquery() + ) + .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[dict]: + """获取工作流运行记录""" + async with AsyncSessionLocal() as session: + workflow_run = await session.get(WorkflowRun, workflow_run_id) + return workflow_run.to_dict() if workflow_run else None + + +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] diff --git a/workflow_service/database/connection.py b/workflow_service/database/connection.py new file mode 100644 index 0000000..d7b0e26 --- /dev/null +++ b/workflow_service/database/connection.py @@ -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 workflow_service.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() diff --git a/workflow_service/database/models.py b/workflow_service/database/models.py new file mode 100644 index 0000000..404e61c --- /dev/null +++ b/workflow_service/database/models.py @@ -0,0 +1,111 @@ +""" +数据库模型定义 +定义工作流相关的所有数据表结构 +""" + +from datetime import datetime +from sqlalchemy import Column, String, Text, DateTime, ForeignKey +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, + } diff --git a/workflow_service/main.py b/workflow_service/main.py index bc9c873..d74e434 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -4,7 +4,7 @@ import uvicorn from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from workflow_service import database +from workflow_service.database.api import init_db from workflow_service.config import Settings from workflow_service.routes import service, workflow, run, runx @@ -30,7 +30,7 @@ web_app.include_router(runx.runx_router) @web_app.on_event("startup") async def startup_event(): """服务启动时，初始化数据库并为所有配置的服务器创建输入/输出目录。""" - await database.init_db() + await init_db() try: servers = settings.SERVERS print(f"检测到 {len(servers)} 个 ComfyUI 服务器配置。") diff --git a/workflow_service/routes/run.py b/workflow_service/routes/run.py index a2e4ef9..53c7b0f 100644 --- a/workflow_service/routes/run.py +++ b/workflow_service/routes/run.py @@ -1,11 +1,11 @@ import json from typing import Dict, Optional -from fastapi import APIRouter, Body, HTTPException, Request -from fastapi.openapi.models import RequestBody +from fastapi import APIRouter, Body, HTTPException from fastapi.responses import JSONResponse -from workflow_service import comfyui_client, database +from workflow_service import comfyui_client +from workflow_service.database.api import get_workflow run_router = APIRouter( prefix="/api/run", @@ -28,7 +28,7 @@ async def run_workflow( raise HTTPException(status_code=400, detail="`workflow_name` 字段是必需的") # 获取工作流定义 - workflow_data = await database.get_workflow(workflow_name, workflow_version) + workflow_data = await get_workflow(workflow_name, workflow_version) if not workflow_data: detail = ( f"工作流 '{workflow_name}'" diff --git a/workflow_service/routes/runx/model_with_multi_dress.py b/workflow_service/routes/runx/model_with_multi_dress.py index 8846d8d..e5395e8 100644 --- a/workflow_service/routes/runx/model_with_multi_dress.py +++ b/workflow_service/routes/runx/model_with_multi_dress.py @@ -3,7 +3,8 @@ from typing import Optional from fastapi import Body, HTTPException from fastapi.responses import JSONResponse -from workflow_service import comfyui_client, database +from workflow_service import comfyui_client +from workflow_service.database.api import get_workflow from ._base import runx_router RUNX_NAME = "model_with_multi_dress" @@ -21,7 +22,7 @@ async def model_with_multi_dress( workflow_name = RUNX_WORKFLOW_NAME # 获取工作流定义 - workflow_data = await database.get_workflow(workflow_name, workflow_version) + workflow_data = await get_workflow(workflow_name, workflow_version) if not workflow_data: detail = ( f"工作流 '{workflow_name}'" diff --git a/workflow_service/routes/workflow.py b/workflow_service/routes/workflow.py index 6c3932d..20c7e2c 100644 --- a/workflow_service/routes/workflow.py +++ b/workflow_service/routes/workflow.py @@ -5,7 +5,12 @@ from fastapi import APIRouter, Body, HTTPException, Path from fastapi.responses import JSONResponse from workflow_service import comfyui_client -from workflow_service import database +from workflow_service.database.api import ( + save_workflow, + get_all_workflows, + delete_workflow, + get_workflow, +) workflow_router = APIRouter( @@ -25,7 +30,7 @@ async def publish_workflow_endpoint(data: dict = Body(...)): raise HTTPException( status_code=400, detail="`name` and `workflow` fields are required." ) - await database.save_workflow(name, json.dumps(wf_json)) + await save_workflow(name, json.dumps(wf_json)) print(f"Workflow '{name}' published.") return JSONResponse( content={"status": "success", "message": f"Workflow '{name}' published."}, @@ -43,7 +48,7 @@ async def get_all_workflows_endpoint(): 获取所有工作流 """ try: - return await database.get_all_workflows() + return await get_all_workflows() except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to get workflows: {e}") @@ -59,7 +64,7 @@ async def delete_workflow_endpoint( 删除工作流 """ try: - success = await database.delete_workflow(workflow_name) + success = await delete_workflow(workflow_name) if success: return {"status": "deleted", "name": workflow_name} else: @@ -75,7 +80,7 @@ async def get_one_workflow_endpoint(base_name: str, version: Optional[str] = Non """ 获取工作流规范 """ - workflow_data = await database.get_workflow(base_name, version) + workflow_data = await get_workflow(base_name, version) if not workflow_data: detail = (