citu_vanna/data_pipeline/api/vector_restore_manager.py

"""
Vector表恢复管理器

提供pgvector表备份文件扫描和数据恢复功能，与VectorTableManager形成完整的备份恢复解决方案
"""

import os
import re
import time
import glob
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, List, Optional
import psycopg2
import logging


class VectorRestoreManager:
    """Vector表恢复管理器 - 仿照VectorTableManager设计"""
    
    def __init__(self, base_output_dir: str = None):
        """
        初始化恢复管理器，复用现有配置机制
        
        Args:
            base_output_dir: 基础输出目录，默认从data_pipeline.config获取
        """
        if base_output_dir is None:
            # 从配置文件获取默认目录
            from data_pipeline.config import SCHEMA_TOOLS_CONFIG
            base_output_dir = SCHEMA_TOOLS_CONFIG.get("output_directory", "./data_pipeline/training_data/")
        
        self.base_output_dir = Path(base_output_dir)
        
        # 从data_pipeline.config获取配置
        from data_pipeline.config import SCHEMA_TOOLS_CONFIG
        self.config = SCHEMA_TOOLS_CONFIG.get("vector_table_management", {})
        
        # 初始化日志
        self.logger = logging.getLogger("VectorRestoreManager")
        
        # 支持的表名
        self.supported_tables = self.config.get("supported_tables", [
            "langchain_pg_collection",
            "langchain_pg_embedding"
        ])
    
    def scan_backup_files(self, global_only: bool = False, task_id: str = None) -> Dict[str, Any]:
        """
        扫描可用的备份文件
        
        Args:
            global_only: 仅查询全局备份目录（training_data/vector_bak/）
            task_id: 指定task_id，仅查询该任务下的备份文件
            
        Returns:
            包含备份文件信息的字典
        """
        scan_start_time = datetime.now()
        backup_locations = []
        
        try:
            # 确定扫描范围
            if task_id:
                # 仅扫描指定任务
                directories_to_scan = [self.base_output_dir / task_id / "vector_bak"]
            elif global_only:
                # 仅扫描全局目录
                directories_to_scan = [self.base_output_dir / "vector_bak"]
            else:
                # 扫描所有目录
                directories_to_scan = self._get_all_vector_bak_directories()
            
            # 扫描每个目录
            for backup_dir in directories_to_scan:
                if not backup_dir.exists():
                    continue
                    
                # 查找有效的备份集
                backup_sets = self._find_backup_sets(backup_dir)
                if not backup_sets:
                    continue
                
                # 构建备份位置信息
                location_info = self._build_location_info(backup_dir, backup_sets)
                if location_info:
                    backup_locations.append(location_info)
            
            # 构建汇总信息
            summary = self._build_summary(backup_locations, scan_start_time)
            
            return {
                "backup_locations": backup_locations,
                "summary": summary
            }
            
        except Exception as e:
            self.logger.error(f"扫描备份文件失败: {e}")
            raise
    
    def restore_from_backup(self, backup_path: str, timestamp: str, 
                          tables: List[str] = None, db_connection: str = None,
                          truncate_before_restore: bool = False) -> Dict[str, Any]:
        """
        从备份文件恢复数据
        
        Args:
            backup_path: 备份文件所在的目录路径（相对路径）
            timestamp: 备份文件的时间戳
            tables: 要恢复的表名列表，None表示恢复所有表
            db_connection: PostgreSQL连接字符串，None则从config获取
            truncate_before_restore: 恢复前是否清空目标表
            
        Returns:
            恢复操作的详细结果
        """
        start_time = time.time()
        
        # 设置默认表列表
        if tables is None:
            tables = self.supported_tables.copy()
        
        # 验证表名
        invalid_tables = [t for t in tables if t not in self.supported_tables]
        if invalid_tables:
            raise ValueError(f"不支持的表名: {invalid_tables}")
        
        # 解析备份路径
        backup_dir = Path(backup_path)
        if not backup_dir.is_absolute():
            # 相对路径，相对于项目根目录
            project_root = Path(__file__).parent.parent.parent
            backup_dir = project_root / backup_path
        
        if not backup_dir.exists():
            raise FileNotFoundError(f"备份目录不存在: {backup_path}")
        
        # 验证备份文件存在
        missing_files = []
        backup_files = {}
        for table_name in tables:
            csv_file = backup_dir / f"{table_name}_{timestamp}.csv"
            if not csv_file.exists():
                missing_files.append(csv_file.name)
            else:
                backup_files[table_name] = csv_file
        
        if missing_files:
            raise FileNotFoundError(f"备份文件不存在: {', '.join(missing_files)}")
        
        # 初始化结果
        result = {
            "restore_performed": True,
            "truncate_performed": truncate_before_restore,
            "backup_info": {
                "backup_path": backup_path,
                "timestamp": timestamp,
                "backup_date": self._parse_timestamp_to_date(timestamp)
            },
            "truncate_results": {},
            "restore_results": {},
            "errors": [],
            "duration": 0
        }
        
        # 临时修改数据库连接配置
        original_config = None
        if db_connection:
            from data_pipeline.config import SCHEMA_TOOLS_CONFIG
            original_config = SCHEMA_TOOLS_CONFIG.get("default_db_connection")
            SCHEMA_TOOLS_CONFIG["default_db_connection"] = db_connection
        
        try:
            # 执行清空操作（如果需要）
            if truncate_before_restore:
                self.logger.info("🗑️ 开始清空目标表...")
                for table_name in tables:
                    truncate_result = self._truncate_table(table_name)
                    result["truncate_results"][table_name] = truncate_result
                    if not truncate_result.get("success", False):
                        result["errors"].append(f"{table_name}表清空失败")
            
            # 执行恢复操作
            self.logger.info("📥 开始恢复表数据...")
            for table_name in tables:
                csv_file = backup_files[table_name]
                restore_result = self._restore_table_from_csv(table_name, csv_file)
                result["restore_results"][table_name] = restore_result
                if not restore_result.get("success", False):
                    result["errors"].append(f"{table_name}表恢复失败")
            
            # 计算总耗时
            result["duration"] = time.time() - start_time
            
            # 记录最终状态
            if result["errors"]:
                self.logger.warning(f"⚠️ Vector表恢复完成，但有错误: {'; '.join(result['errors'])}")
            else:
                self.logger.info(f"✅ Vector表恢复完成，耗时: {result['duration']:.2f}秒")
            
            return result
            
        finally:
            # 恢复原始配置
            if original_config is not None:
                SCHEMA_TOOLS_CONFIG["default_db_connection"] = original_config
    
    def get_connection(self):
        """获取数据库连接 - 完全复用VectorTableManager的连接逻辑"""
        try:
            # 方法1：如果SCHEMA_TOOLS_CONFIG中有连接字符串，直接使用
            from data_pipeline.config import SCHEMA_TOOLS_CONFIG
            connection_string = SCHEMA_TOOLS_CONFIG.get("default_db_connection")
            if connection_string:
                conn = psycopg2.connect(connection_string)
            else:
                # 方法2：从app_config获取pgvector数据库配置
                import app_config
                pgvector_config = app_config.PGVECTOR_CONFIG
                conn = psycopg2.connect(
                    host=pgvector_config.get('host'),
                    port=pgvector_config.get('port'),
                    database=pgvector_config.get('dbname'),
                    user=pgvector_config.get('user'),
                    password=pgvector_config.get('password')
                )
            
            # 设置自动提交
            conn.autocommit = True
            return conn
            
        except Exception as e:
            self.logger.error(f"pgvector数据库连接失败: {e}")
            raise
    
    def _get_all_vector_bak_directories(self) -> List[Path]:
        """获取所有vector_bak目录"""
        directories = []
        
        # 全局备份目录
        global_backup_dir = self.base_output_dir / "vector_bak"
        if global_backup_dir.exists():
            directories.append(global_backup_dir)
        
        # 任务备份目录 (task_* 和 manual_*)
        for pattern in ["task_*", "manual_*"]:
            for task_dir in self.base_output_dir.glob(pattern):
                if task_dir.is_dir():
                    vector_bak_dir = task_dir / "vector_bak"
                    if vector_bak_dir.exists():
                        directories.append(vector_bak_dir)
        
        return directories
    
    def _find_backup_sets(self, backup_dir: Path) -> List[str]:
        """查找备份目录中的有效备份集"""
        # 查找所有CSV文件
        collection_files = list(backup_dir.glob("langchain_pg_collection_*.csv"))
        embedding_files = list(backup_dir.glob("langchain_pg_embedding_*.csv"))
        
        # 提取时间戳
        collection_timestamps = set()
        embedding_timestamps = set()
        
        for file in collection_files:
            timestamp = self._extract_timestamp_from_filename(file.name)
            if timestamp:
                collection_timestamps.add(timestamp)
        
        for file in embedding_files:
            timestamp = self._extract_timestamp_from_filename(file.name)
            if timestamp:
                embedding_timestamps.add(timestamp)
        
        # 找到同时存在两个文件的时间戳
        valid_timestamps = collection_timestamps & embedding_timestamps
        
        # 按时间戳降序排列（最新的在前）
        return sorted(valid_timestamps, reverse=True)
    
    def _extract_timestamp_from_filename(self, filename: str) -> Optional[str]:
        """从文件名中提取时间戳"""
        # 匹配格式：langchain_pg_collection_20250722_010318.csv
        pattern = r'langchain_pg_(?:collection|embedding)_(\d{8}_\d{6})\.csv'
        match = re.search(pattern, filename)
        return match.group(1) if match else None
    
    def _build_location_info(self, backup_dir: Path, backup_sets: List[str]) -> Optional[Dict[str, Any]]:
        """构建备份位置信息"""
        if not backup_sets:
            return None
        
        # 确定位置类型和相关信息
        relative_path = self._get_relative_path(backup_dir)
        location_type, task_id = self._determine_location_type(backup_dir)
        
        # 构建备份信息列表
        backups = []
        for timestamp in backup_sets:
            backup_info = self._build_backup_info(backup_dir, timestamp)
            if backup_info:
                backups.append(backup_info)
        
        location_info = {
            "type": location_type,
            "relative_path": relative_path,
            "backups": backups
        }
        
        if task_id:
            location_info["task_id"] = task_id
        
        return location_info
    
    def _get_relative_path(self, backup_dir: Path) -> str:
        """获取相对路径（Unix风格）"""
        try:
            # 计算相对于项目根目录的路径
            project_root = Path(__file__).parent.parent.parent
            relative_path = backup_dir.relative_to(project_root)
            # 转换为Unix风格路径
            return "./" + str(relative_path).replace("\\", "/")
        except ValueError:
            # 如果无法计算相对路径，直接转换
            return str(backup_dir).replace("\\", "/")
    
    def _determine_location_type(self, backup_dir: Path) -> tuple:
        """确定位置类型和task_id"""
        backup_dir_str = str(backup_dir)
        
        if "/vector_bak" in backup_dir_str.replace("\\", "/"):
            parent = backup_dir.parent.name
            if parent.startswith(("task_", "manual_")):
                return "task", parent
            else:
                return "global", None
        
        return "unknown", None
    
    def _build_backup_info(self, backup_dir: Path, timestamp: str) -> Optional[Dict[str, Any]]:
        """构建单个备份信息"""
        try:
            collection_file = backup_dir / f"langchain_pg_collection_{timestamp}.csv"
            embedding_file = backup_dir / f"langchain_pg_embedding_{timestamp}.csv"
            log_file = backup_dir / "vector_backup_log.txt"
            
            # 检查文件存在性
            if not (collection_file.exists() and embedding_file.exists()):
                return None
            
            # 获取文件大小
            collection_size = self._format_file_size(collection_file.stat().st_size)
            embedding_size = self._format_file_size(embedding_file.stat().st_size)
            
            # 解析备份日期
            backup_date = self._parse_timestamp_to_date(timestamp)
            
            return {
                "timestamp": timestamp,
                "collection_file": collection_file.name,
                "embedding_file": embedding_file.name,
                "collection_size": collection_size,
                "embedding_size": embedding_size,
                "backup_date": backup_date,
                "has_log": log_file.exists(),
                "log_file": log_file.name if log_file.exists() else None
            }
            
        except Exception as e:
            self.logger.warning(f"构建备份信息失败: {e}")
            return None
    
    def _parse_timestamp_to_date(self, timestamp: str) -> str:
        """将时间戳转换为可读日期格式"""
        try:
            # 解析格式：20250722_010318
            dt = datetime.strptime(timestamp, "%Y%m%d_%H%M%S")
            return dt.strftime("%Y-%m-%d %H:%M:%S")
        except Exception:
            return timestamp
    
    def _build_summary(self, backup_locations: List[Dict], scan_start_time: datetime) -> Dict[str, Any]:
        """构建汇总信息"""
        total_backup_sets = sum(len(loc["backups"]) for loc in backup_locations)
        global_backups = sum(len(loc["backups"]) for loc in backup_locations if loc["type"] == "global")
        task_backups = total_backup_sets - global_backups
        
        return {
            "total_locations": len(backup_locations),
            "total_backup_sets": total_backup_sets,
            "global_backups": global_backups,
            "task_backups": task_backups,
            "scan_time": scan_start_time.isoformat()
        }
    
    def _restore_table_from_csv(self, table_name: str, csv_file: Path) -> Dict[str, Any]:
        """从CSV文件恢复单个表 - 使用COPY FROM STDIN"""
        try:
            start_time = time.time()
            
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    # 检查是否是embedding表，需要特殊处理JSON格式
                    if table_name == "langchain_pg_embedding":
                        self._restore_embedding_table_with_json_fix(cursor, csv_file)
                    else:
                        # 其他表直接使用COPY FROM STDIN
                        with open(csv_file, 'r', encoding='utf-8') as f:
                            # 使用CSV HEADER选项自动跳过表头，无需手动next(f)
                            cursor.copy_expert(
                                f"COPY {table_name} FROM STDIN WITH (FORMAT CSV, HEADER)",
                                f
                            )
                    
                    # 验证导入结果
                    cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                    rows_restored = cursor.fetchone()[0]
            
            duration = time.time() - start_time
            file_size = csv_file.stat().st_size
            
            return {
                "success": True,
                "source_file": csv_file.name,
                "rows_restored": rows_restored,
                "file_size": self._format_file_size(file_size),
                "duration": duration
            }
            
        except Exception as e:
            return {
                "success": False,
                "source_file": csv_file.name,
                "error": str(e)
            }
    
    def _truncate_table(self, table_name: str) -> Dict[str, Any]:
        """清空指定表"""
        try:
            start_time = time.time()
            
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    # 获取清空前的行数
                    cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                    rows_before = cursor.fetchone()[0]
                    
                    # 执行TRUNCATE
                    cursor.execute(f"TRUNCATE TABLE {table_name}")
                    
                    # 验证清空结果
                    cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                    rows_after = cursor.fetchone()[0]
            
            duration = time.time() - start_time
            
            if rows_after == 0:
                return {
                    "success": True,
                    "rows_before": rows_before,
                    "rows_after": rows_after,
                    "duration": duration
                }
            else:
                raise Exception(f"清空失败，表中仍有 {rows_after} 行数据")
                
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def _format_file_size(self, size_bytes: int) -> str:
        """格式化文件大小显示"""
        if size_bytes == 0:
            return "0 B"
        
        size_names = ["B", "KB", "MB", "GB"]
        i = 0
        size = float(size_bytes)
        
        while size >= 1024.0 and i < len(size_names) - 1:
            size /= 1024.0
            i += 1
        
        return f"{size:.1f} {size_names[i]}" 
    
    def _restore_embedding_table_with_json_fix(self, cursor, csv_file: Path):
        """恢复embedding表，修复cmetadata列的JSON格式问题"""
        import csv
        import json
        import ast
        import io
        
        # 读取CSV并修复JSON格式
        corrected_data = io.StringIO()
        
        with open(csv_file, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            writer = csv.writer(corrected_data)
            
            # 处理表头
            header = next(reader)
            writer.writerow(header)
            
            # 找到cmetadata列的索引
            try:
                cmetadata_index = header.index('cmetadata')
            except ValueError:
                # 如果没有cmetadata列，直接使用原始CSV
                corrected_data.seek(0)
                corrected_data.truncate(0)
                f.seek(0)
                corrected_data.write(f.read())
                corrected_data.seek(0)
                cursor.copy_expert(
                    "COPY langchain_pg_embedding FROM STDIN WITH (FORMAT CSV, HEADER)",
                    corrected_data
                )
                return
            
            # 处理数据行
            for row in reader:
                if len(row) > cmetadata_index and row[cmetadata_index]:
                    try:
                        # 尝试将Python字典格式转换为JSON格式
                        # 如果已经是JSON格式，json.loads会成功
                        if row[cmetadata_index].startswith('{') and row[cmetadata_index].endswith('}'):
                            try:
                                # 先尝试作为JSON解析
                                json.loads(row[cmetadata_index])
                                # 已经是有效JSON，不需要转换
                            except json.JSONDecodeError:
                                # 不是有效JSON，尝试作为Python字典解析并转换
                                try:
                                    python_dict = ast.literal_eval(row[cmetadata_index])
                                    row[cmetadata_index] = json.dumps(python_dict, ensure_ascii=False)
                                except (ValueError, SyntaxError):
                                    # 如果都失败了，记录错误但继续处理
                                    self.logger.warning(f"无法解析cmetadata: {row[cmetadata_index]}")
                    except Exception as e:
                        self.logger.warning(f"处理cmetadata时出错: {e}")
                
                writer.writerow(row)
        
        # 使用修复后的数据进行导入
        corrected_data.seek(0)
        cursor.copy_expert(
            "COPY langchain_pg_embedding FROM STDIN WITH (FORMAT CSV, HEADER)",
            corrected_data
        )