import os import json import sys # 新增系统模块导入 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from typing import Dict, Any, List, Optional, Type from dataclasses import dataclass, field from datetime import datetime from pymongo import MongoClient from pymongo.errors import ConnectionFailure from peewee import Model, Field, IntegrityError from src.common.database.database import db from src.common.database.database_model import ( ChatStreams, LLMUsage, Emoji, Messages, Images, ImageDescriptions, OnlineTime, PersonInfo, Knowledges, ThinkingLog, GraphNodes, GraphEdges ) from src.common.logger_manager import get_logger logger = get_logger("mongodb_to_sqlite") @dataclass class MigrationConfig: """迁移配置类""" mongo_collection: str target_model: Type[Model] field_mapping: Dict[str, str] batch_size: int = 500 enable_validation: bool = True skip_duplicates: bool = True unique_fields: List[str] = field(default_factory=list) # 用于重复检查的字段 @dataclass class MigrationStats: """迁移统计信息""" total_documents: int = 0 processed_count: int = 0 success_count: int = 0 error_count: int = 0 skipped_count: int = 0 duplicate_count: int = 0 errors: List[Dict[str, Any]] = field(default_factory=list) def add_error(self, doc_id: Any, error: str, doc_data: Optional[Dict] = None): """添加错误记录""" self.errors.append({ 'doc_id': str(doc_id), 'error': error, 'timestamp': datetime.now().isoformat(), 'doc_data': doc_data }) self.error_count += 1 class MongoToSQLiteMigrator: """MongoDB到SQLite数据迁移器 - 使用Peewee ORM""" def __init__(self, mongo_uri: Optional[str] = None, database_name: Optional[str] = None): self.database_name = database_name or os.getenv("DATABASE_NAME", "MegBot") self.mongo_uri = mongo_uri or self._build_mongo_uri() self.mongo_client: Optional[MongoClient] = None self.mongo_db = None # 迁移配置 self.migration_configs = self._initialize_migration_configs() def _build_mongo_uri(self) -> str: """构建MongoDB连接URI""" if mongo_uri := os.getenv("MONGODB_URI"): return mongo_uri user = os.getenv('MONGODB_USER') password = os.getenv('MONGODB_PASS') host = os.getenv('MONGODB_HOST', 'localhost') port = os.getenv('MONGODB_PORT', '27017') auth_source = os.getenv('MONGODB_AUTH_SOURCE', 'admin') if user and password: return f"mongodb://{user}:{password}@{host}:{port}/{self.database_name}?authSource={auth_source}" else: return f"mongodb://{host}:{port}/{self.database_name}" def _initialize_migration_configs(self) -> List[MigrationConfig]: """初始化迁移配置""" return [ # 表情包迁移配置 MigrationConfig( mongo_collection="emoji", target_model=Emoji, field_mapping={ "full_path": "full_path", "format": "format", "hash": "emoji_hash", "description": "description", "emotion": "emotion", "usage_count": "usage_count", "last_used_time": "last_used_time", "last_used_time": "record_time" # 这个纯粹是为了应付整体映射格式，实际上直接用当前时间戳填了record_time }, unique_fields=["full_path", "emoji_hash"] ), # 聊天流迁移配置 MigrationConfig( mongo_collection="chat_streams", target_model=ChatStreams, field_mapping={ "stream_id": "stream_id", "create_time": "create_time", "group_info.platform": "group_platform",# 由于Mongodb处理私聊时会让group_info值为null，而新的数据库不允许为null，所以私聊聊天流是没法迁移的，等更新吧。 "group_info.group_id": "group_id", # 同上 "group_info.group_name": "group_name", # 同上 "last_active_time": "last_active_time", "platform": "platform", "user_info.platform": "user_platform", "user_info.user_id": "user_id", "user_info.user_nickname": "user_nickname", "user_info.user_cardname": "user_cardname" }, unique_fields=["stream_id"] ), # LLM使用记录迁移配置 MigrationConfig( mongo_collection="llm_usage", target_model=LLMUsage, field_mapping={ "model_name": "model_name", "user_id": "user_id", "request_type": "request_type", "endpoint": "endpoint", "prompt_tokens": "prompt_tokens", "completion_tokens": "completion_tokens", "total_tokens": "total_tokens", "cost": "cost", "status": "status", "timestamp": "timestamp" }, unique_fields=["user_id", "timestamp"] # 组合唯一性 ), # 消息迁移配置 MigrationConfig( mongo_collection="messages", target_model=Messages, field_mapping={ "message_id": "message_id", "time": "time", "chat_id": "chat_id", "chat_info.stream_id": "chat_info_stream_id", "chat_info.platform": "chat_info_platform", "chat_info.user_info.platform": "chat_info_user_platform", "chat_info.user_info.user_id": "chat_info_user_id", "chat_info.user_info.user_nickname": "chat_info_user_nickname", "chat_info.user_info.user_cardname": "chat_info_user_cardname", "chat_info.group_info.platform": "chat_info_group_platform", "chat_info.group_info.group_id": "chat_info_group_id", "chat_info.group_info.group_name": "chat_info_group_name", "chat_info.create_time": "chat_info_create_time", "chat_info.last_active_time": "chat_info_last_active_time", "user_info.platform": "user_platform", "user_info.user_id": "user_id", "user_info.user_nickname": "user_nickname", "user_info.user_cardname": "user_cardname", "processed_plain_text": "processed_plain_text", "detailed_plain_text": "detailed_plain_text", "memorized_times": "memorized_times" }, unique_fields=["message_id"] ), # 图片迁移配置 MigrationConfig( mongo_collection="images", target_model=Images, field_mapping={ "hash": "emoji_hash", "description": "description", "path": "path", "timestamp": "timestamp", "type": "type" }, unique_fields=["path"] ), # 图片描述迁移配置 MigrationConfig( mongo_collection="image_descriptions", target_model=ImageDescriptions, field_mapping={ "type": "type", "hash": "image_description_hash", "description": "description", "timestamp": "timestamp" }, unique_fields=["image_description_hash", "type"] ), # 在线时长迁移配置 MigrationConfig( mongo_collection="online_time", target_model=OnlineTime, field_mapping={ "timestamp": "timestamp", "duration": "duration", "start_timestamp": "start_timestamp", "end_timestamp": "end_timestamp" }, unique_fields=["start_timestamp", "end_timestamp"] ), # 个人信息迁移配置 MigrationConfig( mongo_collection="person_info", target_model=PersonInfo, field_mapping={ "person_id": "person_id", "person_name": "person_name", "name_reason": "name_reason", "platform": "platform", "user_id": "user_id", "nickname": "nickname", "relationship_value": "relationship_value", "konw_time": "know_time", "msg_interval": "msg_interval", "msg_interval_list": "msg_interval_list" }, unique_fields=["person_id"] ), # 知识库迁移配置 MigrationConfig( mongo_collection="knowledges", target_model=Knowledges, field_mapping={ "content": "content", "embedding": "embedding" }, unique_fields=["content"] # 假设内容唯一 ), # 思考日志迁移配置 MigrationConfig( mongo_collection="thinking_log", target_model=ThinkingLog, field_mapping={ "chat_id": "chat_id", "trigger_text": "trigger_text", "response_text": "response_text", "trigger_info": "trigger_info_json", "response_info": "response_info_json", "timing_results": "timing_results_json", "chat_history": "chat_history_json", "chat_history_in_thinking": "chat_history_in_thinking_json", "chat_history_after_response": "chat_history_after_response_json", "heartflow_data": "heartflow_data_json", "reasoning_data": "reasoning_data_json", }, unique_fields=["chat_id", "created_at"] ), # 图节点迁移配置 MigrationConfig( mongo_collection="graph_data.nodes", target_model=GraphNodes, field_mapping={ "concept": "concept", "memory_items": "memory_items", "hash": "hash", "created_time": "created_time", "last_modified": "last_modified" }, unique_fields=["concept"] ), # 图边迁移配置 MigrationConfig( mongo_collection="graph_data.edges", target_model=GraphEdges, field_mapping={ "source": "source", "target": "target", "strength": "strength", "hash": "hash", "created_time": "created_time", "last_modified": "last_modified" }, unique_fields=["source", "target"] # 组合唯一性 ) ] def connect_mongodb(self) -> bool: """连接到MongoDB""" try: self.mongo_client = MongoClient( self.mongo_uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=10000, maxPoolSize=10 ) # 测试连接 self.mongo_client.admin.command('ping') self.mongo_db = self.mongo_client[self.database_name] logger.info(f"成功连接到MongoDB: {self.database_name}") return True except ConnectionFailure as e: logger.error(f"MongoDB连接失败: {e}") return False except Exception as e: logger.error(f"MongoDB连接异常: {e}") return False def disconnect_mongodb(self): """断开MongoDB连接""" if self.mongo_client: self.mongo_client.close() logger.info("MongoDB连接已关闭") def _get_nested_value(self, document: Dict[str, Any], field_path: str) -> Any: """获取嵌套字段的值""" if "." not in field_path: return document.get(field_path) parts = field_path.split(".") value = document for part in parts: if isinstance(value, dict): value = value.get(part) else: return None if value is None: break return value def _convert_field_value(self, value: Any, target_field: Field) -> Any: """根据目标字段类型转换值""" if value is None: return None field_type = target_field.__class__.__name__ try: if target_field.name == "record_time" and field_type == "DateTimeField": return datetime.now() if target_field.name == "record_time" and field_type == "DateTimeField": return self._convert_record_time(value) if field_type in ["CharField", "TextField"]: if isinstance(value, (list, dict)): return json.dumps(value, ensure_ascii=False) return str(value) if value is not None else "" elif field_type == "IntegerField": if isinstance(value, str): # 处理字符串数字 clean_value = value.strip() if clean_value.replace('.', '').replace('-', '').isdigit(): return int(float(clean_value)) return 0 return int(value) if value is not None else 0 elif field_type in ["FloatField", "DoubleField"]: return float(value) if value is not None else 0.0 elif field_type == "BooleanField": if isinstance(value, str): return value.lower() in ('true', '1', 'yes', 'on') return bool(value) elif field_type == "DateTimeField": if isinstance(value, (int, float)): return datetime.fromtimestamp(value) elif isinstance(value, str): try: # 尝试解析ISO格式日期 return datetime.fromisoformat(value.replace('Z', '+00:00')) except ValueError: try: # 尝试解析时间戳字符串 return datetime.fromtimestamp(float(value)) except ValueError: return datetime.now() return datetime.now() return value except (ValueError, TypeError) as e: logger.warning(f"字段值转换失败 ({field_type}): {value} -> {e}") return self._get_default_value_for_field(target_field) def _get_default_value_for_field(self, field: Field) -> Any: """获取字段的默认值""" field_type = field.__class__.__name__ if hasattr(field, 'default') and field.default is not None: return field.default if field.null: return None # 根据字段类型返回默认值 if field_type in ["CharField", "TextField"]: return "" elif field_type == "IntegerField": return 0 elif field_type in ["FloatField", "DoubleField"]: return 0.0 elif field_type == "BooleanField": return False elif field_type == "DateTimeField": return datetime.now() return None def _check_duplicate_by_unique_fields(self, model: Type[Model], data: Dict[str, Any], unique_fields: List[str]) -> bool: """根据唯一字段检查重复""" if not unique_fields: return False try: query = model.select() for field_name in unique_fields: if field_name in data and data[field_name] is not None: field_obj = getattr(model, field_name) query = query.where(field_obj == data[field_name]) return query.exists() except Exception as e: logger.debug(f"重复检查失败: {e}") return False def _create_model_instance(self, model: Type[Model], data: Dict[str, Any]) -> Optional[Model]: """使用ORM创建模型实例""" try: # 过滤掉不存在的字段 valid_data = {} for field_name, value in data.items(): if hasattr(model, field_name): valid_data[field_name] = value else: logger.debug(f"跳过未知字段: {field_name}") # 创建实例 instance = model.create(**valid_data) return instance except IntegrityError as e: # 处理唯一约束冲突等完整性错误 logger.debug(f"完整性约束冲突: {e}") return None except Exception as e: logger.error(f"创建模型实例失败: {e}") return None def migrate_collection(self, config: MigrationConfig) -> MigrationStats: """迁移单个集合 - 使用ORM方式""" stats = MigrationStats() logger.info(f"开始迁移: {config.mongo_collection} -> {config.target_model._meta.table_name}") try: # 获取MongoDB集合 mongo_collection = self.mongo_db[config.mongo_collection] stats.total_documents = mongo_collection.count_documents({}) if stats.total_documents == 0: logger.warning(f"集合 {config.mongo_collection} 为空，跳过迁移") return stats logger.info(f"待迁移文档数量: {stats.total_documents}") # 逐个处理文档 batch_count = 0 for mongo_doc in mongo_collection.find().batch_size(config.batch_size): try: stats.processed_count += 1 doc_id = mongo_doc.get('_id', 'unknown') # 构建目标数据 target_data = {} for mongo_field, sqlite_field in config.field_mapping.items(): value = self._get_nested_value(mongo_doc, mongo_field) # 获取目标字段对象并转换类型 if hasattr(config.target_model, sqlite_field): field_obj = getattr(config.target_model, sqlite_field) converted_value = self._convert_field_value(value, field_obj) target_data[sqlite_field] = converted_value # 重复检查 if config.skip_duplicates and self._check_duplicate_by_unique_fields( config.target_model, target_data, config.unique_fields ): stats.duplicate_count += 1 stats.skipped_count += 1 logger.debug(f"跳过重复记录: {doc_id}") continue # 使用ORM创建实例 with db.atomic(): # 每个实例的事务保护 instance = self._create_model_instance(config.target_model, target_data) if instance: stats.success_count += 1 else: stats.add_error(doc_id, "ORM创建实例失败", target_data) except Exception as e: doc_id = mongo_doc.get('_id', 'unknown') stats.add_error(doc_id, f"处理文档异常: {e}", mongo_doc) logger.error(f"处理文档失败 (ID: {doc_id}): {e}") # 进度报告 batch_count += 1 if batch_count % config.batch_size == 0: progress = (stats.processed_count / stats.total_documents) * 100 logger.info( f"迁移进度: {stats.processed_count}/{stats.total_documents} " f"({progress:.1f}%) - 成功: {stats.success_count}, " f"错误: {stats.error_count}, 跳过: {stats.skipped_count}" ) logger.info( f"迁移完成: {config.mongo_collection} -> {config.target_model._meta.table_name}\n" f"总计: {stats.total_documents}, 成功: {stats.success_count}, " f"错误: {stats.error_count}, 跳过: {stats.skipped_count}, 重复: {stats.duplicate_count}" ) except Exception as e: logger.error(f"迁移集合 {config.mongo_collection} 时发生异常: {e}") stats.add_error("collection_error", str(e)) return stats def migrate_all(self) -> Dict[str, MigrationStats]: """执行所有迁移任务""" logger.info("开始执行数据库迁移...") if not self.connect_mongodb(): logger.error("无法连接到MongoDB，迁移终止") return {} all_stats = {} try: for config in self.migration_configs: logger.info(f"\n开始处理集合: {config.mongo_collection}") stats = self.migrate_collection(config) all_stats[config.mongo_collection] = stats # 错误率检查 if stats.processed_count > 0: error_rate = stats.error_count / stats.processed_count if error_rate > 0.1: # 错误率超过10% logger.warning( f"集合 {config.mongo_collection} 错误率较高: {error_rate:.1%} " f"({stats.error_count}/{stats.processed_count})" ) finally: self.disconnect_mongodb() self._print_migration_summary(all_stats) return all_stats def _print_migration_summary(self, all_stats: Dict[str, MigrationStats]): """打印迁移汇总信息""" logger.info("\n" + "="*60) logger.info("数据迁移汇总报告") logger.info("="*60) total_processed = sum(stats.processed_count for stats in all_stats.values()) total_success = sum(stats.success_count for stats in all_stats.values()) total_errors = sum(stats.error_count for stats in all_stats.values()) total_skipped = sum(stats.skipped_count for stats in all_stats.values()) total_duplicates = sum(stats.duplicate_count for stats in all_stats.values()) # 表头 logger.info(f"{'集合名称':<20} | {'处理':<6} | {'成功':<6} | {'错误':<6} | {'跳过':<6} | {'重复':<6} | {'成功率':<8}") logger.info("-" * 75) for collection_name, stats in all_stats.items(): success_rate = (stats.success_count / stats.processed_count * 100) if stats.processed_count > 0 else 0 logger.info( f"{collection_name:<20} | " f"{stats.processed_count:<6} | " f"{stats.success_count:<6} | " f"{stats.error_count:<6} | " f"{stats.skipped_count:<6} | " f"{stats.duplicate_count:<6} | " f"{success_rate:<7.1f}%" ) logger.info("-" * 75) total_success_rate = (total_success / total_processed * 100) if total_processed > 0 else 0 logger.info( f"{'总计':<20} | " f"{total_processed:<6} | " f"{total_success:<6} | " f"{total_errors:<6} | " f"{total_skipped:<6} | " f"{total_duplicates:<6} | " f"{total_success_rate:<7.1f}%" ) if total_errors > 0: logger.warning(f"\n⚠️ 存在 {total_errors} 个错误，请检查日志详情") if total_duplicates > 0: logger.info(f"ℹ️ 跳过了 {total_duplicates} 个重复记录") logger.info("="*60) def add_migration_config(self, config: MigrationConfig): """添加新的迁移配置""" self.migration_configs.append(config) def migrate_single_collection(self, collection_name: str) -> Optional[MigrationStats]: """迁移单个指定的集合""" config = next((c for c in self.migration_configs if c.mongo_collection == collection_name), None) if not config: logger.error(f"未找到集合 {collection_name} 的迁移配置") return None if not self.connect_mongodb(): logger.error("无法连接到MongoDB") return None try: stats = self.migrate_collection(config) self._print_migration_summary({collection_name: stats}) return stats finally: self.disconnect_mongodb() def export_error_report(self, all_stats: Dict[str, MigrationStats], filepath: str): """导出错误报告""" error_report = { 'timestamp': datetime.now().isoformat(), 'summary': { collection: { 'total': stats.total_documents, 'processed': stats.processed_count, 'success': stats.success_count, 'errors': stats.error_count, 'skipped': stats.skipped_count, 'duplicates': stats.duplicate_count } for collection, stats in all_stats.items() }, 'errors': { collection: stats.errors for collection, stats in all_stats.items() if stats.errors } } try: with open(filepath, 'w', encoding='utf-8') as f: json.dump(error_report, f, ensure_ascii=False, indent=2) logger.info(f"错误报告已导出到: {filepath}") except Exception as e: logger.error(f"导出错误报告失败: {e}") def main(): """主程序入口""" migrator = MongoToSQLiteMigrator() # 执行迁移 migration_results = migrator.migrate_all() # 导出错误报告（如果有错误） if any(stats.error_count > 0 for stats in migration_results.values()): error_report_path = f"migration_errors_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" migrator.export_error_report(migration_results, error_report_path) logger.info("数据迁移完成！") if __name__ == "__main__": main()