183 lines
7.3 KiB
Python
183 lines
7.3 KiB
Python
import os
|
|
import json
|
|
from typing import Dict, Any, Optional
|
|
from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import sessionmaker
|
|
from datetime import datetime
|
|
|
|
# 创建数据库基类
|
|
Base = declarative_base()
|
|
|
|
# 定义workflow_version表模型
|
|
class WorkflowVersion(Base):
|
|
__tablename__ = 'workflow_version'
|
|
|
|
id = Column(Integer, primary_key=True, autoincrement=True)
|
|
session_id = Column(String(255), nullable=False)
|
|
workflow_data = Column(Text, nullable=False) # JSON字符串 api格式
|
|
workflow_data_ui = Column(Text, nullable=True) # JSON字符串 ui格式
|
|
attributes = Column(Text, nullable=True) # JSON字符串，存储额外属性
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
def to_dict(self):
|
|
return {
|
|
'id': self.id,
|
|
'session_id': self.session_id,
|
|
'workflow_data': json.loads(self.workflow_data) if self.workflow_data else None,
|
|
'attributes': json.loads(self.attributes) if self.attributes else None,
|
|
'created_at': self.created_at.isoformat() if self.created_at else None
|
|
}
|
|
|
|
class DatabaseManager:
|
|
"""数据库管理器"""
|
|
|
|
def __init__(self, db_path: str = None):
|
|
if db_path is None:
|
|
# 默认数据库路径
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
db_dir = os.path.join(current_dir, '..', 'data')
|
|
os.makedirs(db_dir, exist_ok=True)
|
|
db_path = os.path.join(db_dir, 'workflow_debug.db')
|
|
|
|
self.db_path = db_path
|
|
self.engine = create_engine(f'sqlite:///{db_path}', echo=False)
|
|
self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
|
|
|
|
# 创建表
|
|
Base.metadata.create_all(bind=self.engine)
|
|
|
|
def get_session(self):
|
|
"""获取数据库会话"""
|
|
return self.SessionLocal()
|
|
|
|
def save_workflow_version(self, session_id: str, workflow_data: Dict[str, Any], workflow_data_ui: Dict[str, Any] = None, attributes: Optional[Dict[str, Any]] = None) -> int:
|
|
"""保存工作流版本，返回新版本的ID"""
|
|
session = self.get_session()
|
|
try:
|
|
workflow_version = WorkflowVersion(
|
|
session_id=session_id,
|
|
workflow_data=json.dumps(workflow_data),
|
|
workflow_data_ui=json.dumps(workflow_data_ui) if workflow_data_ui else None,
|
|
attributes=json.dumps(attributes) if attributes else None
|
|
)
|
|
session.add(workflow_version)
|
|
session.commit()
|
|
session.refresh(workflow_version)
|
|
return workflow_version.id
|
|
except Exception as e:
|
|
session.rollback()
|
|
raise e
|
|
finally:
|
|
session.close()
|
|
|
|
def get_current_workflow_data(self, session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取当前session的最新工作流数据（最大ID版本）"""
|
|
session = self.get_session()
|
|
try:
|
|
latest_version = session.query(WorkflowVersion)\
|
|
.filter(WorkflowVersion.session_id == session_id)\
|
|
.order_by(WorkflowVersion.id.desc())\
|
|
.first()
|
|
|
|
if latest_version:
|
|
return json.loads(latest_version.workflow_data)
|
|
return None
|
|
finally:
|
|
session.close()
|
|
|
|
def get_current_workflow_data_ui(self, session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取当前session的最新工作流数据（最大ID版本）"""
|
|
session = self.get_session()
|
|
try:
|
|
latest_version = session.query(WorkflowVersion)\
|
|
.filter(WorkflowVersion.session_id == session_id)\
|
|
.order_by(WorkflowVersion.id.desc())\
|
|
.first()
|
|
|
|
if latest_version:
|
|
return json.loads(latest_version.workflow_data_ui)
|
|
return None
|
|
finally:
|
|
session.close()
|
|
|
|
def get_workflow_version_by_id(self, version_id: int) -> Optional[Dict[str, Any]]:
|
|
"""根据版本ID获取工作流数据"""
|
|
session = self.get_session()
|
|
try:
|
|
version = session.query(WorkflowVersion)\
|
|
.filter(WorkflowVersion.id == version_id)\
|
|
.first()
|
|
|
|
if version:
|
|
result = version.to_dict()
|
|
# 添加UI格式的工作流数据
|
|
if version.workflow_data_ui:
|
|
result['workflow_data_ui'] = json.loads(version.workflow_data_ui)
|
|
return result
|
|
return None
|
|
finally:
|
|
session.close()
|
|
|
|
def update_workflow_version(self, version_id: int, workflow_data: Dict[str, Any], attributes: Optional[Dict[str, Any]] = None) -> bool:
|
|
"""更新指定版本的工作流数据"""
|
|
session = self.get_session()
|
|
try:
|
|
version = session.query(WorkflowVersion)\
|
|
.filter(WorkflowVersion.id == version_id)\
|
|
.first()
|
|
|
|
if version:
|
|
version.workflow_data = json.dumps(workflow_data)
|
|
if attributes:
|
|
version.attributes = json.dumps(attributes)
|
|
session.commit()
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
session.rollback()
|
|
raise e
|
|
finally:
|
|
session.close()
|
|
|
|
def update_workflow_ui(self, version_id: int, workflow_data_ui: Dict[str, Any]) -> bool:
|
|
"""只更新指定版本的workflow_data_ui字段，不影响其他字段"""
|
|
session = self.get_session()
|
|
try:
|
|
version = session.query(WorkflowVersion)\
|
|
.filter(WorkflowVersion.id == version_id)\
|
|
.first()
|
|
|
|
if version:
|
|
version.workflow_data_ui = json.dumps(workflow_data_ui)
|
|
session.commit()
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
session.rollback()
|
|
raise e
|
|
finally:
|
|
session.close()
|
|
|
|
# 全局数据库管理器实例
|
|
db_manager = DatabaseManager()
|
|
|
|
def get_workflow_data(session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取当前session的工作流数据的便捷函数"""
|
|
return db_manager.get_current_workflow_data(session_id)
|
|
|
|
def get_workflow_data_ui(session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取当前session的工作流数据的便捷函数"""
|
|
return db_manager.get_current_workflow_data_ui(session_id)
|
|
|
|
def save_workflow_data(session_id: str, workflow_data: Dict[str, Any], workflow_data_ui: Dict[str, Any] = None, attributes: Optional[Dict[str, Any]] = None) -> int:
|
|
"""保存工作流数据的便捷函数"""
|
|
return db_manager.save_workflow_version(session_id, workflow_data, workflow_data_ui, attributes)
|
|
|
|
def get_workflow_data_by_id(version_id: int) -> Optional[Dict[str, Any]]:
|
|
"""根据版本ID获取工作流数据的便捷函数"""
|
|
return db_manager.get_workflow_version_by_id(version_id)
|
|
|
|
def update_workflow_ui_by_id(version_id: int, workflow_data_ui: Dict[str, Any]) -> bool:
|
|
"""只更新指定版本的workflow_data_ui字段的便捷函数"""
|
|
return db_manager.update_workflow_ui(version_id, workflow_data_ui) |