citu_vanna/data_pipeline/api/simple_db_manager.py

"""
Data Pipeline API 简化数据库管理器

复用现有的pgvector数据库连接机制，提供Data Pipeline任务的数据库操作功能
"""

import json
from datetime import datetime
from typing import Dict, Any, List, Optional, Tuple

import psycopg2
from psycopg2.extras import RealDictCursor, Json

from app_config import PGVECTOR_CONFIG
from core.logging import get_data_pipeline_logger


class SimpleTaskManager:
    """简化的任务管理器，复用现有pgvector连接"""
    
    def __init__(self):
        """初始化任务管理器"""
        self.logger = get_data_pipeline_logger("SimpleTaskManager")
        self._connection = None
    
    def _get_connection(self):
        """获取pgvector数据库连接"""
        if self._connection is None or self._connection.closed:
            try:
                self._connection = 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')
                )
                self._connection.autocommit = True
            except Exception as e:
                self.logger.error(f"pgvector数据库连接失败: {e}")
                raise
        return self._connection
    
    def close_connection(self):
        """关闭数据库连接"""
        if self._connection and not self._connection.closed:
            self._connection.close()
            self._connection = None
    
    def generate_task_id(self) -> str:
        """生成任务ID，格式: task_YYYYMMDD_HHMMSS"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"task_{timestamp}"
    
    def create_task(self, 
                   table_list_file: str,
                   business_context: str,
                   db_name: str = None,
                   **kwargs) -> str:
        """创建新任务"""
        task_id = self.generate_task_id()
        
        # 从 app_config 获取业务数据库连接信息
        from app_config import APP_DB_CONFIG
        
        # 构建业务数据库连接字符串（用于参数记录）
        business_db_connection = self._build_db_connection_string(APP_DB_CONFIG)
        
        # 使用传入的db_name或从APP_DB_CONFIG提取
        if not db_name:
            db_name = APP_DB_CONFIG.get('dbname', 'business_db')
        
        # 构建参数
        parameters = {
            "db_connection": business_db_connection,  # 业务数据库连接（用于schema_workflow执行）
            "table_list_file": table_list_file,
            "business_context": business_context,
            **kwargs
        }
        
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("""
                    INSERT INTO data_pipeline_tasks (
                        id, task_type, status, parameters, created_by, 
                        db_name, business_context, output_directory
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    task_id, 
                    'data_workflow', 
                    'pending', 
                    Json(parameters),
                    'api',
                    db_name,
                    business_context,
                    f"./data_pipeline/training_data/{task_id}"
                ))
                
            self.logger.info(f"任务创建成功: {task_id}")
            return task_id
            
        except Exception as e:
            self.logger.error(f"任务创建失败: {e}")
            raise
    
    def get_task(self, task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务信息"""
        try:
            conn = self._get_connection()
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("SELECT * FROM data_pipeline_tasks WHERE id = %s", (task_id,))
                result = cursor.fetchone()
                return dict(result) if result else None
        except Exception as e:
            self.logger.error(f"获取任务信息失败: {e}")
            raise
    
    def update_task_status(self, task_id: str, status: str, error_message: Optional[str] = None):
        """更新任务状态"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                update_fields = ["status = %s"]
                values = [status]
                
                if status == 'in_progress' and not self._get_task_started_at(task_id):
                    update_fields.append("started_at = CURRENT_TIMESTAMP")
                
                if status in ['completed', 'failed']:
                    update_fields.append("completed_at = CURRENT_TIMESTAMP")
                
                if error_message:
                    update_fields.append("error_message = %s")
                    values.append(error_message)
                
                values.append(task_id)
                
                cursor.execute(f"""
                    UPDATE data_pipeline_tasks 
                    SET {', '.join(update_fields)}
                    WHERE id = %s
                """, values)
                
                self.logger.info(f"任务状态更新: {task_id} -> {status}")
        except Exception as e:
            self.logger.error(f"任务状态更新失败: {e}")
            raise
    
    def update_step_status(self, task_id: str, step_name: str, step_status: str):
        """更新步骤状态"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("""
                    UPDATE data_pipeline_tasks 
                    SET step_status = jsonb_set(step_status, %s, %s)
                    WHERE id = %s
                """, ([step_name], json.dumps(step_status), task_id))
                
                self.logger.debug(f"步骤状态更新: {task_id}.{step_name} -> {step_status}")
        except Exception as e:
            self.logger.error(f"步骤状态更新失败: {e}")
            raise
    
    def create_execution(self, task_id: str, execution_step: str) -> str:
        """创建执行记录"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        execution_id = f"{task_id}_step_{execution_step}_exec_{timestamp}"
        
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("""
                    INSERT INTO data_pipeline_task_executions (
                        task_id, execution_step, status, execution_id
                    ) VALUES (%s, %s, %s, %s)
                """, (task_id, execution_step, 'running', execution_id))
                
                self.logger.info(f"执行记录创建: {execution_id}")
                return execution_id
        except Exception as e:
            self.logger.error(f"执行记录创建失败: {e}")
            raise
    
    def complete_execution(self, execution_id: str, status: str, error_message: Optional[str] = None):
        """完成执行记录"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                # 计算执行时长
                cursor.execute("""
                    SELECT started_at FROM data_pipeline_task_executions 
                    WHERE execution_id = %s
                """, (execution_id,))
                result = cursor.fetchone()
                
                duration_seconds = None
                if result and result[0]:
                    duration_seconds = int((datetime.now() - result[0]).total_seconds())
                
                # 更新执行记录
                update_fields = ["status = %s", "completed_at = CURRENT_TIMESTAMP"]
                values = [status]
                
                if duration_seconds is not None:
                    update_fields.append("duration_seconds = %s")
                    values.append(duration_seconds)
                
                if error_message:
                    update_fields.append("error_message = %s")
                    values.append(error_message)
                
                values.append(execution_id)
                
                cursor.execute(f"""
                    UPDATE data_pipeline_task_executions 
                    SET {', '.join(update_fields)}
                    WHERE execution_id = %s
                """, values)
                
                self.logger.info(f"执行记录完成: {execution_id} -> {status}")
        except Exception as e:
            self.logger.error(f"执行记录完成失败: {e}")
            raise
    
    def record_log(self, task_id: str, log_level: str, message: str, 
                   execution_id: Optional[str] = None, step_name: Optional[str] = None):
        """记录日志到数据库"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("""
                    INSERT INTO data_pipeline_task_logs (
                        task_id, execution_id, log_level, message, step_name
                    ) VALUES (%s, %s, %s, %s, %s)
                """, (task_id, execution_id, log_level, message, step_name))
        except Exception as e:
            self.logger.error(f"日志记录失败: {e}")
    
    def get_task_logs(self, task_id: str, limit: int = 100) -> List[Dict[str, Any]]:
        """获取任务日志"""
        try:
            conn = self._get_connection()
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("""
                    SELECT * FROM data_pipeline_task_logs 
                    WHERE task_id = %s 
                    ORDER BY timestamp DESC 
                    LIMIT %s
                """, (task_id, limit))
                
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            self.logger.error(f"获取任务日志失败: {e}")
            raise
    
    def get_task_executions(self, task_id: str) -> List[Dict[str, Any]]:
        """获取任务执行记录"""
        try:
            conn = self._get_connection()
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute("""
                    SELECT * FROM data_pipeline_task_executions 
                    WHERE task_id = %s 
                    ORDER BY started_at DESC
                """, (task_id,))
                
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            self.logger.error(f"获取执行记录失败: {e}")
            raise
    
    def get_tasks_list(self, limit: int = 50, offset: int = 0, status_filter: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取任务列表"""
        try:
            conn = self._get_connection()
            with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                where_clause = ""
                params = []
                
                if status_filter:
                    where_clause = "WHERE status = %s"
                    params.append(status_filter)
                
                params.extend([limit, offset])
                
                cursor.execute(f"""
                    SELECT * FROM data_pipeline_tasks 
                    {where_clause}
                    ORDER BY created_at DESC 
                    LIMIT %s OFFSET %s
                """, params)
                
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            self.logger.error(f"获取任务列表失败: {e}")
            raise
    
    def _get_task_started_at(self, task_id: str) -> Optional[datetime]:
        """获取任务开始时间"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("SELECT started_at FROM data_pipeline_tasks WHERE id = %s", (task_id,))
                result = cursor.fetchone()
                return result[0] if result and result[0] else None
        except Exception:
            return None
    
    def _build_db_connection_string(self, db_config: dict) -> str:
        """构建数据库连接字符串"""
        try:
            host = db_config.get('host', 'localhost')
            port = db_config.get('port', 5432)
            dbname = db_config.get('dbname', 'database')
            user = db_config.get('user', 'postgres')
            password = db_config.get('password', '')
            
            return f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
        except Exception:
            return "postgresql://localhost:5432/database"
    
    def _extract_db_name(self, connection_string: str) -> str:
        """从连接字符串提取数据库名称"""
        try:
            if '/' in connection_string:
                db_name = connection_string.split('/')[-1]
                if '?' in db_name:
                    db_name = db_name.split('?')[0]
                return db_name if db_name else "database"
            else:
                return "database"
        except Exception:
            return "database"