""" 统一数据迁移脚本 功能： - 阶段化执行迁移，支持单独运行某个阶段 - 详细的日志输出和进度追踪 - 支持预览、回滚和验证 - 数据完整性检查 使用方式： # 预览所有迁移 python scripts/migrate_all.py --dry-run # 执行所有迁移 python scripts/migrate_all.py --execute # 只迁移业务数据 (SQLite -> PostgreSQL) python scripts/migrate_all.py --execute --stage business # 只迁移知识库元数据 (JSON -> PostgreSQL) python scripts/migrate_all.py --execute --stage knowledge # 验证迁移结果 python scripts/migrate_all.py --verify # 回滚所有迁移 python scripts/migrate_all.py --rollback # 回滚指定阶段 python scripts/migrate_all.py --rollback --stage business """ import argparse import asyncio import glob import json import os import sys from dataclasses import dataclass, field from datetime import datetime, UTC from typing import Any, Callable # 确保路径正确 sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) os.environ.setdefault("YUXI_SKIP_APP_INIT", "1") from sqlalchemy import Column, DateTime, Integer, String, Text, create_engine, select, text from sqlalchemy.orm import declarative_base, sessionmaker from src import config from src.storage.postgres.manager import pg_manager from src.storage.postgres.models_business import ( Department, User, Conversation, Message, ToolCall, ConversationStats, OperationLog, MessageFeedback, MCPServer, AgentConfig, TaskRecord, ) from src.storage.postgres.models_knowledge import ( KnowledgeBase, KnowledgeFile, EvaluationBenchmark, EvaluationResult, EvaluationResultDetail, ) from src.utils import logger # ============================================================ # 迁移阶段定义 # ============================================================ @dataclass class MigrationStage: """迁移阶段""" name: str # 阶段名称 description: str # 阶段描述 migrate_fn: Callable # 迁移函数 rollback_fn: Callable | None = None # 回滚函数 verify_fn: Callable | None = None # 验证函数 depends_on: list[str] = field(default_factory=list) # 依赖阶段 @dataclass class MigrationResult: """迁移结果""" stage_name: str success: bool dry_run: bool records_total: int = 0 records_migrated: int = 0 records_skipped: int = 0 error: str | None = None duration_ms: float = 0.0 # ============================================================ # SQLite 模型定义 (仅用于迁移) # ============================================================ SQLiteBase = declarative_base() class SqliteDepartment(SQLiteBase): __tablename__ = "departments" id = Column(Integer, primary_key=True) name = Column(String(100), nullable=False) description = Column(Text) created_at = Column(DateTime) class SqliteUser(SQLiteBase): __tablename__ = "users" id = Column(Integer, primary_key=True) username = Column(String(50), unique=True, nullable=False) user_id = Column(String(50), unique=True) phone_number = Column(String(20)) avatar = Column(String(500)) password_hash = Column(String(255)) role = Column(String(20), default="user") department_id = Column(Integer) created_at = Column(DateTime) last_login = Column(DateTime) login_failed_count = Column(Integer, default=0) last_failed_login = Column(DateTime) login_locked_until = Column(DateTime) is_deleted = Column(Integer, default=0) deleted_at = Column(DateTime) class SqliteConversation(SQLiteBase): __tablename__ = "conversations" id = Column(Integer, primary_key=True) thread_id = Column(String(50), unique=True) user_id = Column(String(64), nullable=False) agent_id = Column(String(50)) title = Column(String(255)) status = Column(String(20), default="active") created_at = Column(DateTime) updated_at = Column(DateTime) extra_metadata = Column(Text) class SqliteMessage(SQLiteBase): __tablename__ = "messages" id = Column(Integer, primary_key=True) conversation_id = Column(Integer, nullable=False) role = Column(String(20), nullable=False) content = Column(Text) message_type = Column(String(20), default="text") created_at = Column(DateTime) token_count = Column(Integer) extra_metadata = Column(Text) image_content = Column(Text) class SqliteToolCall(SQLiteBase): __tablename__ = "tool_calls" id = Column(Integer, primary_key=True) message_id = Column(Integer, nullable=False) langgraph_tool_call_id = Column(String(100)) tool_name = Column(String(100)) tool_input = Column(Text) tool_output = Column(Text) status = Column(String(20), default="pending") error_message = Column(Text) created_at = Column(DateTime) class SqliteConversationStats(SQLiteBase): __tablename__ = "conversation_stats" id = Column(Integer, primary_key=True) conversation_id = Column(Integer, nullable=False) message_count = Column(Integer, default=0) total_tokens = Column(Integer, default=0) model_used = Column(String(100)) user_feedback = Column(String(20)) created_at = Column(DateTime) updated_at = Column(DateTime) class SqliteOperationLog(SQLiteBase): __tablename__ = "operation_logs" id = Column(Integer, primary_key=True) user_id = Column(Integer) operation = Column(String(100)) details = Column(Text) ip_address = Column(String(50)) timestamp = Column(DateTime) class SqliteMessageFeedback(SQLiteBase): __tablename__ = "message_feedbacks" id = Column(Integer, primary_key=True) message_id = Column(Integer, nullable=False) user_id = Column(String(64), nullable=False) rating = Column(String(20)) reason = Column(Text) created_at = Column(DateTime) class SqliteMCPServer(SQLiteBase): __tablename__ = "mcp_servers" # 注意：SQLite 中 mcp_servers 表没有 id 列，主键是 name name = Column(String(100), unique=True, nullable=False, primary_key=True) description = Column(Text) transport = Column(String(20), default="sse") url = Column(String(500)) command = Column(String(255)) args = Column(Text) headers = Column(Text) timeout = Column(Integer) sse_read_timeout = Column(Integer) tags = Column(Text) icon = Column(String(500)) enabled = Column(Integer, default=1) disabled_tools = Column(Text) created_by = Column(String(100), nullable=False) updated_by = Column(String(100), nullable=False) created_at = Column(DateTime) updated_at = Column(DateTime) # ============================================================ # 工具函数 # ============================================================ def _utc_dt(value: Any) -> datetime | None: """转换各种 datetime 格式为 naive UTC datetime""" if not value: return None if isinstance(value, datetime): if value.tzinfo is None: return value return value.astimezone(UTC).replace(tzinfo=None) if isinstance(value, (int, float)): return datetime.fromtimestamp(value, tz=UTC).replace(tzinfo=None) if isinstance(value, str): v = value.strip() if not v: return None try: dt_val = datetime.fromisoformat(v.replace("Z", "+00:00")) if dt_val.tzinfo is None: return dt_val return dt_val.astimezone(UTC).replace(tzinfo=None) except ValueError: return None return None def _load_json(path: str) -> dict[str, Any]: """加载 JSON 文件""" if not os.path.exists(path): return {} with open(path, encoding="utf-8") as f: return json.load(f) def _parse_json(value: Any) -> Any: """解析 JSON 字符串或返回原值""" if value is None: return None if isinstance(value, dict): return value if isinstance(value, str): try: return json.loads(value) except json.JSONDecodeError: return None return value def _log_separator(title: str = "", char: str = "=", width: int = 60) -> str: """生成分隔线""" if title: return f"{char * ((width - len(title) - 2) // 2)} {title} {char * ((width - len(title) - 2) // 2)}" return char * width # ============================================================ # SQLite 读取器 # ============================================================ class SQLiteReader: """SQLite 数据读取器""" def __init__(self): db_path = os.path.join(config.save_dir, "database", "server.db") if not os.path.exists(db_path): raise FileNotFoundError(f"SQLite 数据库不存在: {db_path}") self.engine = create_engine(f"sqlite:///{db_path}") self.Session = sessionmaker(bind=self.engine) def get_session(self): return self.Session() def count_table(self, table_name: str) -> int: with self.get_session() as session: result = session.execute(text(f"SELECT COUNT(*) FROM {table_name}")) return result.scalar() or 0 def read_all(self, model): with self.get_session() as session: return session.execute(select(model)).scalars().all() # ============================================================ # 迁移阶段实现 # ============================================================ class MigrationRunner: """迁移运行器""" def __init__(self, dry_run: bool = False): self.dry_run = dry_run self.results: list[MigrationResult] = [] self.start_time: datetime | None = None def log(self, message: str, level: str = "INFO"): """带时间戳的日志输出""" now = datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S") prefix = { "INFO": "ℹ️", "WARN": "⚠️", "ERROR": "❌", "SUCCESS": "✅", "STAGE": "🔄", }.get(level, "ℹ️") logger.info(f"[{now}] {prefix} {message}") async def run_stage(self, stage: MigrationStage) -> MigrationResult: """执行单个迁移阶段""" start = datetime.now() result = MigrationResult(stage_name=stage.name, success=False, dry_run=self.dry_run) self.log(_log_separator(f"阶段: {stage.name}")) self.log(stage.description) try: if self.dry_run: self.log("[DRY-RUN] 预览模式，跳过实际迁移") result.success = True else: result = await stage.migrate_fn(result) result.success = True except Exception as e: result.error = str(e) result.success = False self.log(f"迁移失败: {e}", level="ERROR") result.duration_ms = (datetime.now() - start).total_seconds() * 1000 self.results.append(result) # 输出阶段结果 status = "✅ 成功" if result.success else "❌ 失败" self.log(f"阶段完成: {status} ({result.duration_ms:.1f}ms)") if result.records_total > 0: self.log(f" 记录: {result.records_migrated}/{result.records_total} 迁移, {result.records_skipped} 跳过") return result # ----- 业务数据迁移阶段 ----- async def migrate_business_departments(self, result: MigrationResult) -> MigrationResult: """迁移部门数据""" sqlite_reader = SQLiteReader() sqlite_depts = sqlite_reader.read_all(SqliteDepartment) result.records_total = len(sqlite_depts) if self.dry_run: for d in sqlite_depts: self.log(f"[DRY-RUN] 将创建部门: {d.name}") return result async with pg_manager.get_async_session_context() as session: for sqlite_dept in sqlite_depts: existing = await session.execute(select(Department).where(Department.id == sqlite_dept.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue dept = Department( id=sqlite_dept.id, name=sqlite_dept.name, description=sqlite_dept.description, created_at=_utc_dt(sqlite_dept.created_at), ) session.add(dept) result.records_migrated += 1 return result async def migrate_business_users(self, result: MigrationResult) -> MigrationResult: """迁移用户数据""" sqlite_reader = SQLiteReader() sqlite_users = sqlite_reader.read_all(SqliteUser) result.records_total = len(sqlite_users) if self.dry_run: for u in sqlite_users: self.log(f"[DRY-RUN] 将创建用户: {u.username} ({u.user_id})") return result async with pg_manager.get_async_session_context() as session: for sqlite_user in sqlite_users: existing = await session.execute(select(User).where(User.id == sqlite_user.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue user = User( id=sqlite_user.id, username=sqlite_user.username, user_id=sqlite_user.user_id, phone_number=sqlite_user.phone_number, avatar=sqlite_user.avatar, password_hash=sqlite_user.password_hash, role=sqlite_user.role, department_id=sqlite_user.department_id, created_at=_utc_dt(sqlite_user.created_at), last_login=_utc_dt(sqlite_user.last_login), login_failed_count=sqlite_user.login_failed_count, last_failed_login=_utc_dt(sqlite_user.last_failed_login), login_locked_until=_utc_dt(sqlite_user.login_locked_until), is_deleted=sqlite_user.is_deleted, deleted_at=_utc_dt(sqlite_user.deleted_at), ) session.add(user) result.records_migrated += 1 return result async def migrate_business_conversations(self, result: MigrationResult) -> MigrationResult: """迁移对话数据""" sqlite_reader = SQLiteReader() sqlite_convs = sqlite_reader.read_all(SqliteConversation) result.records_total = len(sqlite_convs) if self.dry_run: for c in sqlite_convs: self.log(f"[DRY-RUN] 将创建对话: {c.thread_id}") return result async with pg_manager.get_async_session_context() as session: for sqlite_conv in sqlite_convs: existing = await session.execute(select(Conversation).where(Conversation.id == sqlite_conv.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue title = sqlite_conv.title if title and len(title) > 255: title = title[:255] conv = Conversation( id=sqlite_conv.id, thread_id=sqlite_conv.thread_id, user_id=sqlite_conv.user_id, agent_id=sqlite_conv.agent_id, title=title, status=sqlite_conv.status, created_at=_utc_dt(sqlite_conv.created_at), updated_at=_utc_dt(sqlite_conv.updated_at), extra_metadata=_parse_json(sqlite_conv.extra_metadata), ) session.add(conv) result.records_migrated += 1 return result async def migrate_business_messages(self, result: MigrationResult) -> MigrationResult: """迁移消息数据""" sqlite_reader = SQLiteReader() sqlite_msgs = sqlite_reader.read_all(SqliteMessage) result.records_total = len(sqlite_msgs) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(sqlite_msgs)} 条消息") return result async with pg_manager.get_async_session_context() as session: for sqlite_msg in sqlite_msgs: existing = await session.execute(select(Message).where(Message.id == sqlite_msg.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue msg = Message( id=sqlite_msg.id, conversation_id=sqlite_msg.conversation_id, role=sqlite_msg.role, content=sqlite_msg.content, message_type=sqlite_msg.message_type, created_at=_utc_dt(sqlite_msg.created_at), token_count=sqlite_msg.token_count, extra_metadata=_parse_json(sqlite_msg.extra_metadata), image_content=sqlite_msg.image_content, ) session.add(msg) result.records_migrated += 1 return result async def migrate_business_tool_calls(self, result: MigrationResult) -> MigrationResult: """迁移工具调用数据""" sqlite_reader = SQLiteReader() sqlite_calls = sqlite_reader.read_all(SqliteToolCall) result.records_total = len(sqlite_calls) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(sqlite_calls)} 个工具调用") return result async with pg_manager.get_async_session_context() as session: for sqlite_call in sqlite_calls: existing = await session.execute(select(ToolCall).where(ToolCall.id == sqlite_call.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue call = ToolCall( id=sqlite_call.id, message_id=sqlite_call.message_id, langgraph_tool_call_id=sqlite_call.langgraph_tool_call_id, tool_name=sqlite_call.tool_name, tool_input=_parse_json(sqlite_call.tool_input), tool_output=sqlite_call.tool_output, status=sqlite_call.status, error_message=sqlite_call.error_message, created_at=_utc_dt(sqlite_call.created_at), ) session.add(call) result.records_migrated += 1 return result async def migrate_business_stats(self, result: MigrationResult) -> MigrationResult: """迁移对话统计数据""" sqlite_reader = SQLiteReader() sqlite_stats = sqlite_reader.read_all(SqliteConversationStats) result.records_total = len(sqlite_stats) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(sqlite_stats)} 条对话统计") return result async with pg_manager.get_async_session_context() as session: for sqlite_stat in sqlite_stats: existing = await session.execute( select(ConversationStats).where(ConversationStats.id == sqlite_stat.id) ) if existing.scalar_one_or_none(): result.records_skipped += 1 continue stat = ConversationStats( id=sqlite_stat.id, conversation_id=sqlite_stat.conversation_id, message_count=sqlite_stat.message_count, total_tokens=sqlite_stat.total_tokens, model_used=sqlite_stat.model_used, user_feedback=sqlite_stat.user_feedback, created_at=_utc_dt(sqlite_stat.created_at), updated_at=_utc_dt(sqlite_stat.updated_at), ) session.add(stat) result.records_migrated += 1 return result async def migrate_business_operation_logs(self, result: MigrationResult) -> MigrationResult: """迁移操作日志""" sqlite_reader = SQLiteReader() sqlite_logs = sqlite_reader.read_all(SqliteOperationLog) result.records_total = len(sqlite_logs) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(sqlite_logs)} 条操作日志") return result async with pg_manager.get_async_session_context() as session: for sqlite_log in sqlite_logs: existing = await session.execute(select(OperationLog).where(OperationLog.id == sqlite_log.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue log = OperationLog( id=sqlite_log.id, user_id=sqlite_log.user_id, operation=sqlite_log.operation, details=sqlite_log.details, ip_address=sqlite_log.ip_address, timestamp=_utc_dt(sqlite_log.timestamp), ) session.add(log) result.records_migrated += 1 return result async def migrate_business_feedbacks(self, result: MigrationResult) -> MigrationResult: """迁移消息反馈""" sqlite_reader = SQLiteReader() sqlite_fbs = sqlite_reader.read_all(SqliteMessageFeedback) result.records_total = len(sqlite_fbs) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(sqlite_fbs)} 条消息反馈") return result async with pg_manager.get_async_session_context() as session: for sqlite_fb in sqlite_fbs: existing = await session.execute(select(MessageFeedback).where(MessageFeedback.id == sqlite_fb.id)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue fb = MessageFeedback( id=sqlite_fb.id, message_id=sqlite_fb.message_id, user_id=sqlite_fb.user_id, rating=sqlite_fb.rating, reason=sqlite_fb.reason, created_at=_utc_dt(sqlite_fb.created_at), ) session.add(fb) result.records_migrated += 1 return result async def migrate_business_mcp_servers(self, result: MigrationResult) -> MigrationResult: """迁移 MCP 服务器""" sqlite_reader = SQLiteReader() sqlite_servers = sqlite_reader.read_all(SqliteMCPServer) result.records_total = len(sqlite_servers) if self.dry_run: for s in sqlite_servers: self.log(f"[DRY-RUN] 将创建 MCP 服务器: {s.name}") return result async with pg_manager.get_async_session_context() as session: for sqlite_server in sqlite_servers: existing = await session.execute(select(MCPServer).where(MCPServer.name == sqlite_server.name)) if existing.scalar_one_or_none(): result.records_skipped += 1 continue server = MCPServer( name=sqlite_server.name, description=sqlite_server.description, transport=sqlite_server.transport, url=sqlite_server.url, command=sqlite_server.command, args=sqlite_server.args, headers=sqlite_server.headers, timeout=sqlite_server.timeout, sse_read_timeout=sqlite_server.sse_read_timeout, tags=sqlite_server.tags, icon=sqlite_server.icon, enabled=sqlite_server.enabled, disabled_tools=sqlite_server.disabled_tools, created_by=sqlite_server.created_by, updated_by=sqlite_server.updated_by, created_at=_utc_dt(sqlite_server.created_at), updated_at=_utc_dt(sqlite_server.updated_at), ) session.add(server) result.records_migrated += 1 return result # ----- 知识库迁移阶段 ----- async def migrate_knowledge_bases(self, result: MigrationResult) -> MigrationResult: """迁移知识库""" base_dir = os.path.join(config.save_dir, "knowledge_base_data") global_meta_path = os.path.join(base_dir, "global_metadata.json") global_meta = _load_json(global_meta_path).get("databases", {}) kb_rows = [] kb_type_dirs = [ p for p in glob.glob(os.path.join(base_dir, "*_data")) if os.path.isdir(p) and os.path.basename(p) != "uploads" ] for kb_dir in kb_type_dirs: kb_type = os.path.basename(kb_dir)[: -len("_data")] meta_file = os.path.join(kb_dir, f"metadata_{kb_type}.json") meta = _load_json(meta_file) databases_meta = meta.get("databases", {}) for db_id, db_meta in databases_meta.items(): g = global_meta.get(db_id, {}) created_at = _utc_dt(g.get("created_at") or db_meta.get("created_at")) updated_at = _utc_dt(g.get("updated_at")) or created_at kb_rows.append({ "db_id": db_id, "name": g.get("name") or db_meta.get("name") or db_id, "description": g.get("description") or db_meta.get("description"), "kb_type": g.get("kb_type") or db_meta.get("kb_type") or kb_type, "embed_info": db_meta.get("embed_info") or g.get("embed_info"), "llm_info": db_meta.get("llm_info") or g.get("llm_info"), "query_params": db_meta.get("query_params") or g.get("query_params"), "additional_params": g.get("additional_params") or db_meta.get("metadata") or {}, "share_config": {"is_shared": True, "accessible_departments": []}, "mindmap": g.get("mindmap"), "sample_questions": g.get("sample_questions") or [], "created_at": created_at, "updated_at": updated_at, }) result.records_total = len(kb_rows) if self.dry_run: for kb in kb_rows: self.log(f"[DRY-RUN] 将创建知识库: {kb['name']} ({kb['db_id']})") return result from src.repositories.knowledge_base_repository import KnowledgeBaseRepository kb_repo = KnowledgeBaseRepository() for payload in kb_rows: db_id = payload["db_id"] existing = await kb_repo.get_by_id(db_id) if existing: result.records_skipped += 1 continue await kb_repo.create(payload) result.records_migrated += 1 return result async def migrate_knowledge_files(self, result: MigrationResult) -> MigrationResult: """迁移知识文件""" base_dir = os.path.join(config.save_dir, "knowledge_base_data") file_rows = [] kb_type_dirs = [ p for p in glob.glob(os.path.join(base_dir, "*_data")) if os.path.isdir(p) and os.path.basename(p) != "uploads" ] for kb_dir in kb_type_dirs: meta_file = os.path.join(kb_dir, f"metadata_{os.path.basename(kb_dir)[:-5]}.json") meta = _load_json(meta_file) files_meta = meta.get("files", {}) for file_id, fmeta in files_meta.items(): db_id = fmeta.get("database_id") if not db_id: continue file_rows.append({ "file_id": file_id, "db_id": db_id, "parent_id": fmeta.get("parent_id"), "filename": fmeta.get("filename") or "", "original_filename": fmeta.get("original_filename") or fmeta.get("file_name"), "file_type": fmeta.get("file_type") or fmeta.get("type"), "path": fmeta.get("path"), "minio_url": fmeta.get("minio_url"), "markdown_file": fmeta.get("markdown_file"), "status": fmeta.get("status"), "content_hash": fmeta.get("content_hash"), "file_size": fmeta.get("size") or fmeta.get("file_size"), "content_type": fmeta.get("content_type"), "processing_params": fmeta.get("processing_params"), "is_folder": bool(fmeta.get("is_folder", False)), "error_message": fmeta.get("error") or fmeta.get("error_message"), "created_by": str(fmeta.get("created_by")) if fmeta.get("created_by") else None, "updated_by": str(fmeta.get("updated_by")) if fmeta.get("updated_by") else None, "created_at": _utc_dt(fmeta.get("created_at")), "updated_at": _utc_dt(fmeta.get("updated_at")) or _utc_dt(fmeta.get("created_at")), }) result.records_total = len(file_rows) if self.dry_run: folders = [f for f in file_rows if f["is_folder"]] files = [f for f in file_rows if not f["is_folder"]] self.log(f"[DRY-RUN] 将创建 {len(folders)} 个文件夹和 {len(files)} 个文件") return result from src.repositories.knowledge_file_repository import KnowledgeFileRepository file_repo = KnowledgeFileRepository() # 先插入文件夹 folders = [(f["file_id"], f) for f in file_rows if f["is_folder"]] files = [(f["file_id"], f) for f in file_rows if not f["is_folder"]] for file_id, data in folders: data = data.copy() data.pop("file_id", None) # 移除重复的 file_id await file_repo.upsert(file_id=file_id, data=data) result.records_migrated += 1 for file_id, data in files: data = data.copy() data.pop("file_id", None) # 移除重复的 file_id await file_repo.upsert(file_id=file_id, data=data) result.records_migrated += 1 return result async def migrate_knowledge_evaluations(self, result: MigrationResult) -> MigrationResult: """迁移评估数据""" base_dir = os.path.join(config.save_dir, "knowledge_base_data") total_migrated = 0 kb_type_dirs = [ p for p in glob.glob(os.path.join(base_dir, "*_data")) if os.path.isdir(p) and os.path.basename(p) != "uploads" ] from src.repositories.evaluation_repository import EvaluationRepository from src.repositories.knowledge_base_repository import KnowledgeBaseRepository eval_repo = EvaluationRepository() kb_repo = KnowledgeBaseRepository() # 迁移评估基准 benchmark_rows = [] for kb_dir in kb_type_dirs: kb_type = os.path.basename(kb_dir)[: -len("_data")] meta_file = os.path.join(kb_dir, f"metadata_{kb_type}.json") meta = _load_json(meta_file) benchmarks_meta = meta.get("benchmarks", {}) for db_id, bmap in benchmarks_meta.items(): if not isinstance(bmap, dict): continue for benchmark_id, bmeta in bmap.items(): benchmark_rows.append({ "benchmark_id": benchmark_id, "db_id": db_id, "name": bmeta.get("name") or benchmark_id, "description": bmeta.get("description"), "question_count": int(bmeta.get("question_count") or 0), "has_gold_chunks": bool(bmeta.get("has_gold_chunks")), "has_gold_answers": bool(bmeta.get("has_gold_answers")), "data_file_path": bmeta.get("benchmark_file") or bmeta.get("data_file_path"), "created_by": str(bmeta.get("created_by")) if bmeta.get("created_by") else None, "created_at": _utc_dt(bmeta.get("created_at")), "updated_at": _utc_dt(bmeta.get("updated_at")) or _utc_dt(bmeta.get("created_at")), }) result.records_total += len(benchmark_rows) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(benchmark_rows)} 个评估基准") return result for payload in benchmark_rows: existing = await eval_repo.get_benchmark(payload["benchmark_id"]) if existing: result.records_skipped += 1 continue # 检查知识库是否存在 kb = await kb_repo.get_by_id(payload["db_id"]) if kb is None: self.log(f" 跳过评估基准 {payload['benchmark_id']}: 知识库 {payload['db_id']} 不存在") result.records_skipped += 1 continue await eval_repo.create_benchmark(payload) total_migrated += 1 # 迁移评估结果 result_rows = [] result_detail_rows = [] for kb_dir in kb_type_dirs: kb_type = os.path.basename(kb_dir)[: -len("_data")] meta_file = os.path.join(kb_dir, f"metadata_{kb_type}.json") meta = _load_json(meta_file) databases_meta = meta.get("databases", {}) for db_id in databases_meta.keys(): result_dir = os.path.join(kb_dir, db_id, "results") if not os.path.isdir(result_dir): continue for result_path in glob.glob(os.path.join(result_dir, "*.json")): try: data = _load_json(result_path) except Exception: continue task_id = data.get("task_id") or os.path.splitext(os.path.basename(result_path))[0] benchmark_id = data.get("benchmark_id") started_at = _utc_dt(data.get("started_at")) result_rows.append({ "task_id": task_id, "db_id": db_id, "benchmark_id": benchmark_id, "status": data.get("status") or "completed", "retrieval_config": data.get("retrieval_config") or {}, "metrics": data.get("metrics") or {}, "overall_score": data.get("overall_score"), "total_questions": int(data.get("total_questions") or 0), "completed_questions": int(data.get("completed_questions") or 0), "started_at": started_at, "completed_at": _utc_dt(data.get("completed_at")) or started_at, "created_by": str(data.get("created_by")) if data.get("created_by") else None, }) interim = data.get("interim_results") or data.get("results") or [] for idx, item in enumerate(interim): result_detail_rows.append({ "task_id": task_id, "query_index": idx, "query_text": item.get("query") or item.get("query_text") or "", "gold_chunk_ids": item.get("gold_chunk_ids"), "gold_answer": item.get("gold_answer"), "generated_answer": item.get("generated_answer"), "retrieved_chunks": item.get("retrieved_chunks"), "metrics": item.get("metrics") or {}, }) result.records_total += len(result_rows) + len(result_detail_rows) if self.dry_run: self.log(f"[DRY-RUN] 将创建 {len(result_rows)} 个评估结果和 {len(result_detail_rows)} 条详情") return result for payload in result_rows: existing = await eval_repo.get_result(payload["task_id"]) if existing: result.records_skipped += 1 continue # 检查知识库是否存在 kb = await kb_repo.get_by_id(payload["db_id"]) if kb is None: self.log(f" 跳过评估结果 {payload['task_id']}: 知识库 {payload['db_id']} 不存在") result.records_skipped += 1 continue await eval_repo.create_result(payload) total_migrated += 1 for detail in result_detail_rows: await eval_repo.upsert_result_detail( task_id=detail["task_id"], query_index=detail["query_index"], data={ "query_text": detail["query_text"], "gold_chunk_ids": detail["gold_chunk_ids"], "gold_answer": detail["gold_answer"], "generated_answer": detail["generated_answer"], "retrieved_chunks": detail["retrieved_chunks"], "metrics": detail["metrics"], }, ) total_migrated += 1 result.records_migrated = total_migrated return result async def migrate_knowledge_tasks(self, result: MigrationResult) -> MigrationResult: """迁移任务记录""" tasks_json_path = os.path.join(config.save_dir, "tasks", "tasks.json") task_items = _load_json(tasks_json_path).get("tasks", []) or [] result.records_total = len(task_items) if self.dry_run: self.log(f"[DRY-RUN] 将迁移 {len(task_items)} 个任务记录") return result from src.repositories.task_repository import TaskRepository task_repo = TaskRepository() for item in task_items: task_id = item.get("id") if not task_id: continue payload = item.get("payload") or {} await task_repo.upsert( task_id, { "name": item.get("name") or "Unnamed Task", "type": item.get("type") or "general", "status": item.get("status") or "pending", "progress": float(item.get("progress") or 0.0), "message": item.get("message") or "", "payload": payload, "result": item.get("result"), "error": item.get("error"), "cancel_requested": 1 if item.get("cancel_requested") else 0, "created_at": _utc_dt(item.get("created_at")), "updated_at": _utc_dt(item.get("updated_at")) or _utc_dt(item.get("created_at")), "started_at": _utc_dt(item.get("started_at")), "completed_at": _utc_dt(item.get("completed_at")), }, ) result.records_migrated += 1 return result # ----- 回滚函数 ----- async def rollback_business(self) -> None: """回滚业务数据""" self.log(_log_separator("回滚: 业务数据"), level="WARN") tables = [ MessageFeedback, OperationLog, ConversationStats, ToolCall, Message, Conversation, User, Department, MCPServer, ] for model in tables: async with pg_manager.get_async_session_context() as session: result = await session.execute(select(model)) records = result.scalars().all() for record in records: await session.delete(record) self.log(f" 已删除 {len(records)} 条 {model.__tablename__}") async def reset_sequences(self) -> None: """重置 PostgreSQL 序列值，防止主键冲突 迁移时直接使用了 SQLite 的原始 id 值，导致 PostgreSQL 的序列未同步。 此方法将序列值重置为当前最大 id + 1。 """ self.log(_log_separator("重置: PostgreSQL 序列"), level="WARN") tables_with_sequences = [ ("departments", "id"), ("users", "id"), ("conversations", "id"), ("messages", "id"), ("tool_calls", "id"), ("conversation_stats", "id"), ("operation_logs", "id"), ("message_feedbacks", "id"), ("mcp_servers", None), # name 是主键，不是 serial ("knowledge_bases", "id"), ("knowledge_files", "id"), ("evaluation_benchmarks", "id"), ("evaluation_results", "id"), ("evaluation_result_details", "id"), ] async with pg_manager.get_async_session_context() as session: for table_name, pk_column in tables_with_sequences: if pk_column is None: continue # 非自增主键，跳过 try: # 使用单条 SQL 获取 max_id 并重置序列 await session.execute( text(f""" SELECT setval( pg_get_serial_sequence('{table_name}', '{pk_column}'), COALESCE((SELECT MAX({pk_column}) FROM {table_name}), 0) + 1 ) """) ) self.log(f" {table_name}: 序列已重置") except Exception as e: self.log(f" {table_name}: 重置序列失败 - {e}", level="WARN") async def rollback_knowledge(self) -> None: """回滚知识库数据""" self.log(_log_separator("回滚: 知识库数据"), level="WARN") from src.repositories.evaluation_repository import EvaluationRepository from src.repositories.knowledge_base_repository import KnowledgeBaseRepository from src.repositories.knowledge_file_repository import KnowledgeFileRepository eval_repo = EvaluationRepository() kb_repo = KnowledgeBaseRepository() file_repo = KnowledgeFileRepository() # 回滚顺序：子表 -> 父表 await eval_repo.delete_all() self.log(" 已删除所有评估数据") rows = await kb_repo.get_all() for row in rows: await file_repo.delete_by_db_id(row.db_id) await kb_repo.delete(row.db_id) self.log(f" 已删除 {len(rows)} 个知识库及其文件") async def rollback_tasker(self) -> None: """回滚 Tasker 任务记录""" self.log(_log_separator("回滚: Tasker 任务记录"), level="WARN") from src.repositories.task_repository import TaskRepository task_repo = TaskRepository() await task_repo.delete_all() self.log(" 已删除所有任务记录") # ----- 验证函数 ----- async def verify_business(self) -> dict: """验证业务数据""" self.log(_log_separator("验证: 业务数据")) results = {} try: sqlite_reader = SQLiteReader() except FileNotFoundError: self.log("SQLite 数据库不存在，跳过验证", level="WARN") return {} sqlite_tables = { "departments": SqliteDepartment, "users": SqliteUser, "conversations": SqliteConversation, "messages": SqliteMessage, "tool_calls": SqliteToolCall, "conversation_stats": SqliteConversationStats, "operation_logs": SqliteOperationLog, "message_feedbacks": SqliteMessageFeedback, "mcp_servers": SqliteMCPServer, } pg_models = { "departments": Department, "users": User, "conversations": Conversation, "messages": Message, "tool_calls": ToolCall, "conversation_stats": ConversationStats, "operation_logs": OperationLog, "message_feedbacks": MessageFeedback, "mcp_servers": (MCPServer, "name"), # MCPServer 主键是 name } for table_name, sqlite_model in sqlite_tables.items(): sqlite_count = sqlite_reader.count_table(table_name) pg_model_info = pg_models[table_name] # 支持 (Model, pk_column) 元组形式 if isinstance(pg_model_info, tuple): pg_model, pk_column = pg_model_info else: pg_model, pk_column = pg_model_info, "id" async with pg_manager.get_async_session_context() as session: from sqlalchemy import func result = await session.execute(select(func.count(getattr(pg_model, pk_column)))) pg_count = result.scalar() or 0 match = sqlite_count == pg_count status = "✅" if match else "❌" results[table_name] = {"sqlite": sqlite_count, "pg": pg_count, "match": match} self.log(f" {status} {table_name}: SQLite={sqlite_count}, PG={pg_count}") return results async def verify_knowledge(self) -> dict: """验证知识库数据""" self.log(_log_separator("验证: 知识库数据")) results = {} base_dir = os.path.join(config.save_dir, "knowledge_base_data") # 统计 JSON 文件中的数据 json_kb_count = 0 json_file_count = 0 kb_type_dirs = [ p for p in glob.glob(os.path.join(base_dir, "*_data")) if os.path.isdir(p) and os.path.basename(p) != "uploads" ] for kb_dir in kb_type_dirs: kb_type = os.path.basename(kb_dir)[: -len("_data")] meta_file = os.path.join(kb_dir, f"metadata_{kb_type}.json") meta = _load_json(meta_file) json_kb_count += len(meta.get("databases", {})) json_file_count += len(meta.get("files", {})) from src.repositories.knowledge_base_repository import KnowledgeBaseRepository from src.repositories.knowledge_file_repository import KnowledgeFileRepository kb_repo = KnowledgeBaseRepository() file_repo = KnowledgeFileRepository() pg_kb_count = len(await kb_repo.get_all()) # 统计文件数量 all_files = [] rows = await kb_repo.get_all() for row in rows: files = await file_repo.list_by_db_id(row.db_id) all_files.extend(files) pg_file_count = len(all_files) results["knowledge_bases"] = {"json": json_kb_count, "pg": pg_kb_count, "match": json_kb_count == pg_kb_count} results["knowledge_files"] = {"json": json_file_count, "pg": pg_file_count, "match": json_file_count == pg_file_count} status_kb = "✅" if results["knowledge_bases"]["match"] else "❌" status_file = "✅" if results["knowledge_files"]["match"] else "❌" self.log(f" {status_kb} knowledge_bases: JSON={json_kb_count}, PG={pg_kb_count}") self.log(f" {status_file} knowledge_files: JSON={json_file_count}, PG={pg_file_count}") return results # ============================================================ # 阶段定义 # ============================================================ def get_stages() -> dict[str, MigrationStage]: """获取所有迁移阶段""" runner = MigrationRunner() return { # 业务数据阶段 (按外键依赖顺序) "business-departments": MigrationStage( name="business-departments", description="迁移部门数据 (departments)", migrate_fn=runner.migrate_business_departments, rollback_fn=None, # 依赖业务回滚整体处理 ), "business-users": MigrationStage( name="business-users", description="迁移用户数据 (users)，依赖 departments", migrate_fn=runner.migrate_business_users, depends_on=["business-departments"], ), "business-conversations": MigrationStage( name="business-conversations", description="迁移对话数据 (conversations)", migrate_fn=runner.migrate_business_conversations, depends_on=["business-users"], ), "business-messages": MigrationStage( name="business-messages", description="迁移消息数据 (messages)，依赖 conversations", migrate_fn=runner.migrate_business_messages, depends_on=["business-conversations"], ), "business-tool-calls": MigrationStage( name="business-tool-calls", description="迁移工具调用数据 (tool_calls)，依赖 messages", migrate_fn=runner.migrate_business_tool_calls, depends_on=["business-messages"], ), "business-stats": MigrationStage( name="business-stats", description="迁移对话统计数据 (conversation_stats)", migrate_fn=runner.migrate_business_stats, depends_on=["business-conversations"], ), "business-operation-logs": MigrationStage( name="business-operation-logs", description="迁移操作日志 (operation_logs)", migrate_fn=runner.migrate_business_operation_logs, depends_on=["business-users"], ), "business-feedbacks": MigrationStage( name="business-feedbacks", description="迁移消息反馈 (message_feedbacks)", migrate_fn=runner.migrate_business_feedbacks, depends_on=["business-messages"], ), "business-mcp-servers": MigrationStage( name="business-mcp-servers", description="迁移 MCP 服务器配置 (mcp_servers)", migrate_fn=runner.migrate_business_mcp_servers, ), # 知识库阶段 "knowledge-bases": MigrationStage( name="knowledge-bases", description="迁移知识库元数据 (knowledge_bases)", migrate_fn=runner.migrate_knowledge_bases, ), "knowledge-files": MigrationStage( name="knowledge-files", description="迁移知识文件元数据 (knowledge_files)，依赖 knowledge_bases", migrate_fn=runner.migrate_knowledge_files, depends_on=["knowledge-bases"], ), "knowledge-evaluations": MigrationStage( name="knowledge-evaluations", description="迁移评估数据 (benchmarks, results)", migrate_fn=runner.migrate_knowledge_evaluations, depends_on=["knowledge-bases"], ), # Tasker 阶段（独立于知识库） "tasker-tasks": MigrationStage( name="tasker-tasks", description="迁移 Tasker 任务记录 (tasks)", migrate_fn=runner.migrate_knowledge_tasks, ), } def get_stage_groups() -> dict[str, list[str]]: """获取阶段组（批量执行）""" return { "business": [ "business-departments", "business-users", "business-conversations", "business-messages", "business-tool-calls", "business-stats", "business-operation-logs", "business-feedbacks", "business-mcp-servers", ], "knowledge": [ "knowledge-bases", "knowledge-files", "knowledge-evaluations", ], "tasker": [ "tasker-tasks", ], "all": list(get_stages().keys()), } # ============================================================ # 主函数 # ============================================================ async def main() -> None: parser = argparse.ArgumentParser(description="统一数据迁移脚本") parser.add_argument("--dry-run", action="store_true", help="预览迁移，不执行") parser.add_argument("--execute", action="store_true", help="执行迁移") parser.add_argument("--verify", action="store_true", help="验证迁移结果") parser.add_argument("--rollback", action="store_true", help="回滚迁移") parser.add_argument("--stage", type=str, help="指定阶段或阶段组 (如: business, knowledge, business-users)") args = parser.parse_args() # 默认dry-run if not any([args.dry_run, args.execute, args.verify, args.rollback]): args.dry_run = True # 初始化 PostgreSQL pg_manager.initialize() await pg_manager.create_tables() runner = MigrationRunner(dry_run=args.dry_run) # 打印标题 if args.dry_run: mode = "预览模式" elif args.execute: mode = "执行模式" elif args.verify: mode = "验证模式" else: mode = "回滚模式" print("\n" + "=" * 60) print(f"🔧 数据迁移工具 | 模式: {mode}") print("=" * 60) if args.verify: # 验证模式 results_business = await runner.verify_business() results_knowledge = await runner.verify_knowledge() print("\n" + "=" * 60) print("📊 验证结果汇总") print("=" * 60) all_match = True for table_name, counts in {**results_business, **results_knowledge}.items(): if not counts.get("match", True): all_match = False print(f"全部匹配: {'✅ 是' if all_match else '❌ 否'}") return if args.rollback: # 回滚模式 if args.stage == "business" or args.stage is None: await runner.rollback_business() if args.stage == "knowledge" or args.stage is None: await runner.rollback_knowledge() if args.stage == "tasker" or args.stage is None: await runner.rollback_tasker() if args.stage == "business": print("\n✅ 已回滚业务数据") elif args.stage == "knowledge": print("\n✅ 已回滚知识库数据") elif args.stage == "tasker": print("\n✅ 已回滚 Tasker 任务记录") else: print("\n✅ 已回滚所有迁移数据") return # 迁移模式 stages = get_stages() stage_groups = get_stage_groups() # 确定要执行的阶段 if args.stage and args.stage in stage_groups: stage_names = stage_groups[args.stage] elif args.stage and args.stage in stages: stage_names = [args.stage] else: stage_names = stage_groups["all"] # 按依赖顺序排序 sorted_stages = [] resolved = set() while sorted_stages.__len__() < len(stage_names): progress = False for name in stage_names: if name in resolved: continue stage = stages[name] if all(dep in resolved for dep in stage.depends_on): sorted_stages.append(name) resolved.add(name) progress = True if not progress: raise ValueError(f"无法解析依赖: {set(stage_names) - resolved}") print(f"\n📋 将执行 {len(sorted_stages)} 个阶段:") for name in sorted_stages: print(f" - {name}") print() # 执行迁移 total_start = datetime.now() for stage_name in sorted_stages: stage = stages[stage_name] await runner.run_stage(stage) # 重置 PostgreSQL 序列，防止后续插入时主键冲突 if not args.dry_run: await runner.reset_sequences() total_duration = (datetime.now() - total_start).total_seconds() # 输出汇总 print("\n" + "=" * 60) print("📊 迁移汇总") print("=" * 60) total_migrated = sum(r.records_migrated for r in runner.results) total_skipped = sum(r.records_skipped for r in runner.results) failed = [r for r in runner.results if not r.success] print(f"总耗时: {total_duration:.1f}s") print(f"迁移记录: {total_migrated}") print(f"跳过记录: {total_skipped}") print(f"失败阶段: {len(failed)}") if failed: print("\n失败详情:") for r in failed: print(f" ❌ {r.stage_name}: {r.error}") if not args.dry_run: print("\n💡 建议运行 --verify 验证数据完整性") if __name__ == "__main__": asyncio.run(main())