citu_vanna/data_pipeline/trainer/vector_table_manager.py

import asyncio
import time
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, List
import psycopg2
import logging


class VectorTableManager:
    """Vector表管理器，负责备份和清空操作"""
    
    def __init__(self, task_output_dir: str, task_id: str = None):
        """
        Args:
            task_output_dir: 任务输出目录（用于存放备份文件）
            task_id: 任务ID（用于日志记录）
        Note:
            数据库连接将从data_pipeline.config.SCHEMA_TOOLS_CONFIG自动获取
        """
        self.task_output_dir = task_output_dir
        self.task_id = task_id
        
        # 从data_pipeline.config获取配置
        from data_pipeline.config import SCHEMA_TOOLS_CONFIG
        self.config = SCHEMA_TOOLS_CONFIG.get("vector_table_management", {})
        
        # 初始化日志
        if task_id:
            from data_pipeline.dp_logging import get_logger
            self.logger = get_logger("VectorTableManager", task_id)
        else:
            import logging
            self.logger = logging.getLogger("VectorTableManager")
    
    def execute_vector_management(self, backup: bool, truncate: bool) -> Dict[str, Any]:
        """执行vector表管理操作的主流程"""
        
        start_time = time.time()
        
        # 1. 参数验证和自动启用逻辑
        if truncate and not backup:
            backup = True
            self.logger.info("🔄 启用truncate时自动启用backup")
        
        if not backup and not truncate:
            self.logger.info("⏭️ 未启用vector表管理，跳过操作")
            return {"backup_performed": False, "truncate_performed": False}
        
        # 2. 初始化结果统计
        result = {
            "backup_performed": backup,
            "truncate_performed": truncate,
            "tables_backed_up": {},
            "truncate_results": {},
            "errors": [],
            "backup_directory": None,
            "duration": 0
        }
        
        try:
            # 3. 创建备份目录
            backup_dir = Path(self.task_output_dir) / self.config.get("backup_directory", "vector_bak")
            if backup:
                backup_dir.mkdir(parents=True, exist_ok=True)
                result["backup_directory"] = str(backup_dir)
                self.logger.info(f"📁 备份目录: {backup_dir}")
            
            # 4. 执行备份操作
            if backup:
                self.logger.info("🗂️ 开始备份vector表...")
                backup_results = self.backup_vector_tables()
                result["tables_backed_up"] = backup_results
                
                # 检查备份是否全部成功
                backup_failed = any(not r.get("success", False) for r in backup_results.values())
                if backup_failed:
                    result["errors"].append("部分表备份失败")
                    if truncate:
                        self.logger.error("❌ 备份失败，取消清空操作")
                        result["truncate_performed"] = False
                        truncate = False
            
            # 5. 执行清空操作（仅在备份成功时）
            if truncate:
                self.logger.info("🗑️ 开始清空vector表...")
                truncate_results = self.truncate_vector_tables()
                result["truncate_results"] = truncate_results
                
                # 检查清空是否成功
                truncate_failed = any(not r.get("success", False) for r in truncate_results.values())
                if truncate_failed:
                    result["errors"].append("部分表清空失败")
            
            # 6. 生成备份日志文件
            if backup and backup_dir.exists():
                self._write_backup_log(backup_dir, result)
            
            # 7. 计算总耗时
            result["duration"] = time.time() - start_time
            
            # 8. 记录最终状态
            if result["errors"]:
                self.logger.warning(f"⚠️ Vector表管理完成，但有错误: {'; '.join(result['errors'])}")
            else:
                self.logger.info(f"✅ Vector表管理完成，耗时: {result['duration']:.2f}秒")
            
            return result
            
        except Exception as e:
            result["duration"] = time.time() - start_time
            result["errors"].append(f"执行失败: {str(e)}")
            self.logger.error(f"❌ Vector表管理失败: {e}")
            raise
    
    def backup_vector_tables(self) -> Dict[str, Any]:
        """备份vector表数据"""
        
        # 1. 创建备份目录
        backup_dir = Path(self.task_output_dir) / self.config.get("backup_directory", "vector_bak")
        backup_dir.mkdir(parents=True, exist_ok=True)
        
        # 2. 生成时间戳
        timestamp = datetime.now().strftime(self.config.get("timestamp_format", "%Y%m%d_%H%M%S"))
        
        # 3. 执行备份（每个表分别处理）
        results = {}
        supported_tables = self.config.get("supported_tables", ["langchain_pg_collection", "langchain_pg_embedding"])
        
        for table_name in supported_tables:
            try:
                # 3.1 定义文件路径（.tmp临时文件）
                temp_file = backup_dir / f"{table_name}_{timestamp}.csv.tmp"
                final_file = backup_dir / f"{table_name}_{timestamp}.csv"
                
                # 确保使用绝对路径（PostgreSQL COPY命令要求）
                temp_file_abs = temp_file.resolve()
                
                # 3.2 通过psycopg2使用流式客户端导出（支持大数据量）
                start_time = time.time()
                row_count = 0
                batch_size = 10000  # 每批处理1万条记录
                
                with self.get_connection() as conn:
                    # 临时关闭autocommit以支持流式处理
                    old_autocommit = conn.autocommit
                    conn.autocommit = False
                    
                    try:
                        with conn.cursor() as cursor:
                            # 设置游标为流式模式
                            cursor.itersize = batch_size
                            
                            # 执行编码设置
                            cursor.execute("SET client_encoding TO 'UTF8'")
                            
                            # 执行查询
                            cursor.execute(f"SELECT * FROM {table_name}")
                            
                            # 获取列名
                            colnames = [desc[0] for desc in cursor.description]
                            
                            # 使用流式方式写入CSV文件
                            import csv
                            with open(temp_file_abs, 'w', newline='', encoding='utf-8') as csvfile:
                                writer = csv.writer(csvfile)
                                
                                # 写入表头
                                writer.writerow(colnames)
                                
                                # 流式读取和写入数据
                                while True:
                                    rows = cursor.fetchmany(batch_size)
                                    if not rows:
                                        break
                                        
                                    # 批量写入当前批次的数据
                                    for row in rows:
                                        writer.writerow(row)
                                        row_count += 1
                                    
                                    # 记录进度（大数据量时有用）
                                    if row_count % (batch_size * 5) == 0:  # 每5万条记录记录一次
                                        self.logger.info(f"📊 {table_name} 已导出 {row_count} 行数据...")
                        
                        # 提交事务
                        conn.commit()
                        
                    finally:
                        # 恢复原来的autocommit设置
                        conn.autocommit = old_autocommit
                
                self.logger.info(f"📊 {table_name} 流式导出完成，总计 {row_count} 行")
                
                # 3.3 导出完成后，重命名文件 (.tmp -> .csv)
                if temp_file.exists():
                    temp_file.rename(final_file)
                    
                    # 3.4 获取文件信息
                    file_stat = final_file.stat()
                    duration = time.time() - start_time
                    
                    results[table_name] = {
                        "success": True,
                        "row_count": row_count,
                        "file_size": self._format_file_size(file_stat.st_size),
                        "backup_file": final_file.name,
                        "duration": duration
                    }
                    
                    self.logger.info(f"✅ {table_name} 备份成功: {row_count}行 -> {final_file.name}")
                else:
                    raise Exception(f"临时文件 {temp_file} 未生成")
                    
            except Exception as e:
                results[table_name] = {
                    "success": False,
                    "error": str(e)
                }
                self.logger.error(f"❌ {table_name} 备份失败: {e}")
                
                # 清理可能的临时文件
                if temp_file.exists():
                    temp_file.unlink()
        
        return results
    
    def truncate_vector_tables(self) -> Dict[str, Any]:
        """清空vector表数据（只清空langchain_pg_embedding）"""
        
        results = {}
        
        # 只清空配置中指定的表（通常只有langchain_pg_embedding）
        truncate_tables = self.config.get("truncate_tables", ["langchain_pg_embedding"])
        
        for table_name in truncate_tables:
            try:
                # 记录清空前的行数（用于统计）
                count_sql = f"SELECT COUNT(*) FROM {table_name}"
                
                start_time = time.time()
                with self.get_connection() as conn:
                    with conn.cursor() as cursor:
                        # 1. 获取清空前的行数
                        cursor.execute(count_sql)
                        rows_before = cursor.fetchone()[0]
                        
                        # 2. 执行TRUNCATE
                        cursor.execute(f"TRUNCATE TABLE {table_name}")
                        
                        # 3. 验证清空结果
                        cursor.execute(count_sql)
                        rows_after = cursor.fetchone()[0]
                
                duration = time.time() - start_time
                
                if rows_after == 0:
                    results[table_name] = {
                        "success": True,
                        "rows_before": rows_before,
                        "rows_after": rows_after,
                        "duration": duration
                    }
                    self.logger.info(f"✅ {table_name} 清空成功: {rows_before}行 -> 0行")
                else:
                    raise Exception(f"清空失败，表中仍有 {rows_after} 行数据")
                    
            except Exception as e:
                results[table_name] = {
                    "success": False,
                    "error": str(e)
                }
                self.logger.error(f"❌ {table_name} 清空失败: {e}")
        
        return results
    
    def get_connection(self):
        """获取pgvector数据库连接（从data_pipeline.config获取配置）"""
        import psycopg2
        
        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 _write_backup_log(self, backup_dir: Path, result: Dict[str, Any]):
        """写入详细的备份日志"""
        log_file = backup_dir / "vector_backup_log.txt"
        
        try:
            with open(log_file, 'w', encoding='utf-8') as f:
                f.write("=== Vector Table Backup Log ===\n")
                f.write(f"Backup Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"Task ID: {self.task_id or 'Unknown'}\n")
                f.write(f"Duration: {result.get('duration', 0):.2f}s\n\n")
                
                # 备份状态
                f.write("Tables Backup Status:\n")
                for table_name, info in result.get("tables_backed_up", {}).items():
                    if info.get("success", False):
                        f.write(f"✓ {table_name}: {info['row_count']} rows -> {info['backup_file']} ({info['file_size']})\n")
                    else:
                        f.write(f"✗ {table_name}: FAILED - {info.get('error', 'Unknown error')}\n")
                
                # 清空状态
                if result.get("truncate_performed", False):
                    f.write("\nTruncate Status:\n")
                    for table_name, info in result.get("truncate_results", {}).items():
                        if info.get("success", False):
                            f.write(f"✓ {table_name}: TRUNCATED ({info['rows_before']} rows removed)\n")
                        else:
                            f.write(f"✗ {table_name}: FAILED - {info.get('error', 'Unknown error')}\n")
                else:
                    f.write("\nTruncate Status:\n- Not performed\n")
                
                # 错误汇总
                if result.get("errors"):
                    f.write(f"\nErrors: {'; '.join(result['errors'])}\n")
                    
        except Exception as e:
            self.logger.warning(f"写入备份日志失败: {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]}"