petpicturequeue/migrations/migrate.py

#!/usr/bin/env python3
"""
Migration runner for Pet Picture Queue
Handles database schema migrations in order
"""

import os
import sys
import importlib.util
import sqlite3
from datetime import datetime


def get_database_path():
    """Get the correct database path"""
    if os.path.exists("pet_pictures.db"):
        return "pet_pictures.db"
    elif os.path.exists("app/pet_pictures.db"):
        return "app/pet_pictures.db"
    else:
        return "pet_pictures.db"


def create_migration_table(cursor):
    """Create migrations tracking table"""
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS migrations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            migration_name TEXT NOT NULL UNIQUE,
            applied_at TIMESTAMP NOT NULL
        )
    """)


def get_applied_migrations(cursor):
    """Get list of applied migrations"""
    try:
        cursor.execute("SELECT migration_name FROM migrations ORDER BY id")
        return [row[0] for row in cursor.fetchall()]
    except sqlite3.OperationalError:
        # Migrations table doesn't exist yet
        return []


def get_migration_files():
    """Get list of migration files in order"""
    migrations_dir = os.path.dirname(__file__)
    migration_files = []
    
    for file in os.listdir(migrations_dir):
        if file.endswith('.py') and file != 'migrate.py' and file != '__init__.py':
            migration_files.append(file)
    
    # Sort by filename (assumes numbered naming like 001_xxx.py)
    migration_files.sort()
    return migration_files


def load_migration_module(migration_file):
    """Load a migration module dynamically"""
    migrations_dir = os.path.dirname(__file__)
    file_path = os.path.join(migrations_dir, migration_file)
    
    spec = importlib.util.spec_from_file_location("migration", file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    
    return module


def run_migrations():
    """Run pending migrations"""
    db_path = get_database_path()
    print(f"Using database: {db_path}")
    
    try:
        db = sqlite3.connect(db_path)
        cursor = db.cursor()
        
        # Create migrations table
        create_migration_table(cursor)
        db.commit()
        
        # Get current state
        applied_migrations = get_applied_migrations(cursor)
        available_migrations = get_migration_files()
        
        print(f"Applied migrations: {len(applied_migrations)}")
        print(f"Available migrations: {len(available_migrations)}")
        
        # Find pending migrations
        pending_migrations = []
        for migration_file in available_migrations:
            migration_name = migration_file.replace('.py', '')
            if migration_name not in applied_migrations:
                pending_migrations.append(migration_file)
        
        if not pending_migrations:
            print("✓ No pending migrations")
            return 0
        
        print(f"\nRunning {len(pending_migrations)} pending migrations:")
        
        # Run each pending migration
        for migration_file in pending_migrations:
            migration_name = migration_file.replace('.py', '')
            print(f"\n📦 Running migration: {migration_name}")
            
            try:
                # Load and run migration
                module = load_migration_module(migration_file)
                
                if hasattr(module, 'migrate_up'):
                    success = module.migrate_up(cursor)
                    if success is not False:  # Allow None or True
                        # Mark as applied
                        cursor.execute(
                            "INSERT INTO migrations (migration_name, applied_at) VALUES (?, ?)",
                            (migration_name, datetime.now())
                        )
                        db.commit()
                        print(f"✓ Migration {migration_name} applied successfully")
                    else:
                        print(f"• Migration {migration_name} skipped (already applied)")
                else:
                    print(f"❌ Migration {migration_name} has no migrate_up function")
                    
            except Exception as e:
                print(f"❌ Migration {migration_name} failed: {e}")
                db.rollback()
                return 1
        
        print(f"\n🎉 All migrations completed successfully!")
        
        # Show final schema
        cursor.execute("PRAGMA table_info(pet_pictures)")
        columns = cursor.fetchall()
        print(f"\nFinal pet_pictures schema:")
        for col in columns:
            nullable = "NOT NULL" if col[3] else "NULL"
            default = f"DEFAULT {col[4]}" if col[4] else ""
            print(f"  {col[1]} {col[2]} {nullable} {default}")
        
        return 0
        
    except Exception as e:
        print(f"❌ Migration runner failed: {e}")
        return 1
    finally:
        if 'db' in locals():
            db.close()


def main():
    """Main entry point"""
    print("=" * 60)
    print("Pet Picture Queue - Database Migration Runner")
    print(f"Started at: {datetime.now().isoformat()}")
    print("=" * 60)
    
    result = run_migrations()
    
    print("=" * 60)
    if result == 0:
        print("✅ Migration process completed successfully!")
    else:
        print("❌ Migration process failed!")
    print("=" * 60)
    
    return result


if __name__ == "__main__":
    sys.exit(main())