run_performance_migration.py

#!/usr/bin/env python3
"""
Script to run performance optimization database migration.
"""

import os
import sys

# Set environment for development
os.environ['FLASK_ENV'] = 'development'
# Use absolute path for SQLite database
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
db_path = os.path.join(current_dir, 'instance', 'chat_agent.db')
os.environ['DATABASE_URL'] = f'sqlite:///{db_path}'

from app import create_app
from chat_agent.models.base import db

def run_migration():
    """Run the performance indexes migration."""
    app = create_app()
    
    with app.app_context():
        print("Running performance indexes migration...")
        
        # Read the migration file
        try:
            with open('migrations/002_performance_indexes_sqlite.sql', 'r') as f:
                migration_sql = f.read()
        except FileNotFoundError:
            print("Migration file not found!")
            return False
        
        # Parse SQL statements more carefully
        import re
        
        # Remove comments and extract CREATE INDEX statements
        lines = migration_sql.split('\n')
        current_statement = []
        statements = []
        
        for line in lines:
            line = line.strip()
            
            # Skip comment lines
            if line.startswith('--') or not line:
                continue
            
            # Add line to current statement
            current_statement.append(line)
            
            # If line ends with semicolon, we have a complete statement
            if line.endswith(';'):
                statement = ' '.join(current_statement)
                statements.append(statement)
                current_statement = []
        
        # Add any remaining statement
        if current_statement:
            statement = ' '.join(current_statement)
            statements.append(statement)
        
        executed_count = 0
        skipped_count = 0
        
        print(f"Found {len(statements)} SQL statements to execute")
        
        for i, statement in enumerate(statements):
            print(f"Processing statement {i+1}: {statement[:50]}...")
                
            try:
                # Execute CREATE INDEX statements
                if 'CREATE INDEX' in statement:
                    print(f"Executing: {statement[:80]}...")
                    from sqlalchemy import text
                    db.session.execute(text(statement))
                    executed_count += 1
                else:
                    print(f"Skipping non-CREATE INDEX statement: {statement[:50]}...")
                    skipped_count += 1
                    
            except Exception as e:
                print(f"Error executing statement: {e}")
                print(f"Statement was: {statement[:100]}...")
                continue
        
        try:
            db.session.commit()
            print(f"Migration completed successfully!")
            print(f"Executed: {executed_count} statements")
            print(f"Skipped: {skipped_count} statements")
            return True
        except Exception as e:
            print(f"Error committing migration: {e}")
            db.session.rollback()
            return False

if __name__ == "__main__":
    success = run_migration()
    sys.exit(0 if success else 1)