Maas-Model-platform/backend/scripts/apply_search_migrations.py

#!/usr/bin/env python3
"""
搜索功能数据库迁移脚本

功能：
- 应用搜索相关的数据库迁移
- 创建用户搜索偏好表
- 创建搜索使用记录表
- 扩展ai_message表支持搜索信息
- 添加相关索引

使用方法：
    python apply_search_migrations.py
    python apply_search_migrations.py --rollback  # 回滚迁移
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List

# 添加项目根目录到Python路径
sys.path.append(str(Path(__file__).parent.parent))

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from app.database import DATABASE_URL, engine


class SearchMigrationManager:
    """搜索功能迁移管理器"""
    
    def __init__(self, database_url: str = None):
        """
        初始化迁移管理器
        
        Args:
            database_url: 数据库连接URL，默认使用项目配置
        """
        self.database_url = database_url or DATABASE_URL
        self.engine = create_engine(self.database_url)
        
        # 定义迁移文件列表（按执行顺序）
        self.migration_files = [
            "030_create_user_search_preferences.sql",
            "031_create_search_usage_log.sql", 
            "032_extend_ai_message_search_support.sql"
        ]
        
        # 迁移文件目录
        self.migrations_dir = Path(__file__).parent.parent / "migrations"
    
    def read_migration_file(self, filename: str) -> str:
        """读取迁移文件内容"""
        file_path = self.migrations_dir / filename
        
        if not file_path.exists():
            raise FileNotFoundError(f"迁移文件不存在: {file_path}")
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        return content
    
    def extract_forward_migration(self, content: str) -> str:
        """提取正向迁移SQL"""
        lines = content.split('\n')
        forward_lines = []
        in_forward_section = False
        
        for line in lines:
            if '正向迁移:' in line or 'Forward Migration' in line:
                in_forward_section = True
                continue
            elif '回滚迁移:' in line or 'Rollback Migration' in line:
                break
            elif in_forward_section and not line.strip().startswith('--'):
                forward_lines.append(line)
        
        return '\n'.join(forward_lines).strip()
    
    def extract_rollback_migration(self, content: str) -> str:
        """提取回滚迁移SQL"""
        lines = content.split('\n')
        rollback_lines = []
        in_rollback_section = False
        
        for line in lines:
            if '回滚迁移:' in line or 'Rollback Migration' in line:
                in_rollback_section = True
                continue
            elif in_rollback_section and line.strip().startswith('-- '):
                # 移除注释符号
                sql_line = line.strip()[3:]  # 移除 '-- '
                if sql_line.strip():
                    rollback_lines.append(sql_line)
        
        return '\n'.join(rollback_lines).strip()
    
    def execute_sql(self, sql: str, description: str = ""):
        """执行SQL语句"""
        if not sql.strip():
            print(f"跳过空SQL: {description}")
            return
        
        try:
            with self.engine.connect() as conn:
                # 分割SQL语句（按分号分割）
                statements = [stmt.strip() for stmt in sql.split(';') if stmt.strip()]
                
                for stmt in statements:
                    if stmt.strip():
                        print(f"执行SQL: {stmt[:100]}...")
                        conn.execute(text(stmt))
                        conn.commit()
                
                print(f"✅ 成功执行: {description}")
                
        except Exception as e:
            print(f"❌ 执行失败: {description}")
            print(f"错误信息: {e}")
            raise
    
    def check_table_exists(self, table_name: str, schema: str = 'aigcspace') -> bool:
        """检查表是否存在"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text("""
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables 
                        WHERE table_schema = :schema 
                        AND table_name = :table_name
                    )
                """), {"schema": schema, "table_name": table_name})
                return result.scalar()
        except Exception as e:
            print(f"检查表存在性失败: {e}")
            return False
    
    def check_column_exists(self, table_name: str, column_name: str, schema: str = 'aigcspace') -> bool:
        """检查列是否存在"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text("""
                    SELECT EXISTS (
                        SELECT FROM information_schema.columns 
                        WHERE table_schema = :schema 
                        AND table_name = :table_name 
                        AND column_name = :column_name
                    )
                """), {"schema": schema, "table_name": table_name, "column_name": column_name})
                return result.scalar()
        except Exception as e:
            print(f"检查列存在性失败: {e}")
            return False
    
    def apply_migrations(self):
        """应用所有搜索相关迁移"""
        print("开始应用搜索功能数据库迁移...")
        print(f"数据库URL: {self.database_url}")
        print(f"迁移文件数量: {len(self.migration_files)}")
        print()
        
        for i, filename in enumerate(self.migration_files, 1):
            print(f"[{i}/{len(self.migration_files)}] 处理迁移文件: {filename}")
            
            try:
                # 读取迁移文件
                content = self.read_migration_file(filename)
                
                # 提取正向迁移SQL
                forward_sql = self.extract_forward_migration(content)
                
                if forward_sql:
                    # 执行迁移
                    self.execute_sql(forward_sql, f"迁移文件 {filename}")
                else:
                    print(f"⚠️  未找到正向迁移SQL: {filename}")
                
            except Exception as e:
                print(f"❌ 迁移失败: {filename}")
                print(f"错误信息: {e}")
                return False
            
            print()
        
        print("🎉 所有搜索功能迁移已成功应用！")
        
        # 验证迁移结果
        self.verify_migrations()
        
        return True
    
    def rollback_migrations(self):
        """回滚所有搜索相关迁移"""
        print("开始回滚搜索功能数据库迁移...")
        print(f"数据库URL: {self.database_url}")
        print()
        
        # 反向执行回滚
        for i, filename in enumerate(reversed(self.migration_files), 1):
            print(f"[{i}/{len(self.migration_files)}] 回滚迁移文件: {filename}")
            
            try:
                # 读取迁移文件
                content = self.read_migration_file(filename)
                
                # 提取回滚迁移SQL
                rollback_sql = self.extract_rollback_migration(content)
                
                if rollback_sql:
                    # 执行回滚
                    self.execute_sql(rollback_sql, f"回滚迁移文件 {filename}")
                else:
                    print(f"⚠️  未找到回滚迁移SQL: {filename}")
                
            except Exception as e:
                print(f"❌ 回滚失败: {filename}")
                print(f"错误信息: {e}")
                return False
            
            print()
        
        print("🎉 所有搜索功能迁移已成功回滚！")
        return True
    
    def verify_migrations(self):
        """验证迁移结果"""
        print("验证迁移结果...")
        
        # 检查表是否创建成功
        tables_to_check = [
            "user_search_preferences",
            "search_usage_log"
        ]
        
        for table in tables_to_check:
            exists = self.check_table_exists(table)
            status = "✅" if exists else "❌"
            print(f"{status} 表 {table}: {'存在' if exists else '不存在'}")
        
        # 检查ai_message表的新列
        columns_to_check = [
            ("ai_message", "search_info"),
            ("ai_message", "search_results")
        ]
        
        for table, column in columns_to_check:
            exists = self.check_column_exists(table, column)
            status = "✅" if exists else "❌"
            print(f"{status} 列 {table}.{column}: {'存在' if exists else '不存在'}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="搜索功能数据库迁移脚本")
    parser.add_argument(
        "--rollback",
        action="store_true",
        help="回滚迁移（默认为应用迁移）"
    )
    parser.add_argument(
        "--database-url",
        type=str,
        help="数据库连接URL（默认使用项目配置）"
    )
    
    args = parser.parse_args()
    
    try:
        # 初始化迁移管理器
        manager = SearchMigrationManager(args.database_url)
        
        # 执行迁移或回滚
        if args.rollback:
            success = manager.rollback_migrations()
        else:
            success = manager.apply_migrations()
        
        if success:
            print("\n✅ 操作完成！")
        else:
            print("\n❌ 操作失败！")
            sys.exit(1)
            
    except Exception as e:
        print(f"\n❌ 执行失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()