LinkDesk/backend/rollback_shot_project_id.py

#!/usr/bin/env python3
"""
Rollback script for shot project_id migration.
This script removes the project_id column and related constraints/indexes.

WARNING: This will permanently remove the project_id column and all related data.
Only use this if you need to rollback the migration.

Usage:
    python rollback_shot_project_id.py
"""

import sqlite3
import sys
from pathlib import Path
import logging

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)


def get_database_path():
    """Get the database path, trying multiple possible locations."""
    possible_paths = [
        "vfx_project_management.db",
        "database.db",
        "../vfx_project_management.db"
    ]
    
    for path in possible_paths:
        if Path(path).exists():
            return path
    
    return "vfx_project_management.db"


def rollback_migration():
    """Remove project_id column and related constraints from shots table."""
    db_path = get_database_path()
    logger.info(f"Using database: {db_path}")
    
    # Confirm rollback
    print("⚠️  WARNING: This will permanently remove the project_id column from shots table!")
    print("⚠️  All project_id data will be lost!")
    response = input("Are you sure you want to continue? (type 'YES' to confirm): ")
    
    if response != 'YES':
        print("Rollback cancelled.")
        return True
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # Check if shots table exists
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='shots'")
        if not cursor.fetchone():
            logger.error("Shots table not found")
            return False
        
        # Check if project_id column exists
        cursor.execute("PRAGMA table_info(shots)")
        columns = [col[1] for col in cursor.fetchall()]
        
        if 'project_id' not in columns:
            logger.info("project_id column does not exist, nothing to rollback")
            return True
        
        logger.info("Starting rollback of shot project_id migration...")
        
        # Create new table without project_id column
        logger.info("Creating new shots table without project_id...")
        cursor.execute("""
            CREATE TABLE shots_rollback (
                id INTEGER PRIMARY KEY,
                episode_id INTEGER NOT NULL,
                name VARCHAR NOT NULL,
                description VARCHAR,
                frame_start INTEGER NOT NULL DEFAULT 1001,
                frame_end INTEGER NOT NULL DEFAULT 1001,
                status VARCHAR(11) NOT NULL DEFAULT 'not_started',
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                deleted_at TIMESTAMP,
                deleted_by INTEGER,
                FOREIGN KEY (episode_id) REFERENCES episodes(id),
                FOREIGN KEY (deleted_by) REFERENCES users(id)
            )
        """)
        
        # Copy data from old table to new table (excluding project_id)
        logger.info("Copying data to new table...")
        cursor.execute("""
            INSERT INTO shots_rollback 
            SELECT id, episode_id, name, description, frame_start, frame_end, 
                   status, created_at, updated_at, deleted_at, deleted_by
            FROM shots
        """)
        
        # Verify data copy
        cursor.execute("SELECT COUNT(*) FROM shots")
        original_count = cursor.fetchone()[0]
        cursor.execute("SELECT COUNT(*) FROM shots_rollback")
        new_count = cursor.fetchone()[0]
        
        if original_count != new_count:
            logger.error(f"Data copy failed: original={original_count}, new={new_count}")
            return False
        
        logger.info(f"Successfully copied {new_count} shots to new table")
        
        # Drop old table and rename new table
        logger.info("Replacing table...")
        cursor.execute("DROP TABLE shots")
        cursor.execute("ALTER TABLE shots_rollback RENAME TO shots")
        
        # Recreate basic indexes (without project_id related ones)
        logger.info("Recreating basic indexes...")
        
        cursor.execute("CREATE INDEX idx_shots_name ON shots(name)")
        cursor.execute("CREATE INDEX idx_shots_episode_id ON shots(episode_id)")
        cursor.execute("CREATE INDEX idx_shots_active ON shots(id) WHERE deleted_at IS NULL")
        
        # Commit changes
        conn.commit()
        
        logger.info("Rollback completed successfully!")
        logger.info(f"Removed project_id column from {new_count} shots")
        
        return True
        
    except sqlite3.Error as e:
        logger.error(f"Database error: {e}")
        if conn:
            conn.rollback()
        return False
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        if conn:
            conn.rollback()
        return False
    finally:
        if conn:
            conn.close()


def verify_rollback():
    """Verify that the rollback was successful."""
    db_path = get_database_path()
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        logger.info("Verifying rollback...")
        
        # Check that project_id column is gone
        cursor.execute("PRAGMA table_info(shots)")
        columns = [col[1] for col in cursor.fetchall()]
        
        if 'project_id' in columns:
            logger.error("❌ project_id column still exists")
            return False
        
        logger.info("✅ project_id column successfully removed")
        
        # Check that basic functionality still works
        cursor.execute("SELECT COUNT(*) FROM shots")
        shot_count = cursor.fetchone()[0]
        logger.info(f"✅ {shot_count} shots remain in table")
        
        # Check that episode relationships still work
        cursor.execute("""
            SELECT COUNT(*) FROM shots s
            JOIN episodes e ON s.episode_id = e.id
        """)
        valid_relationships = cursor.fetchone()[0]
        
        if valid_relationships != shot_count:
            logger.error(f"❌ Episode relationships broken: {valid_relationships}/{shot_count}")
            return False
        
        logger.info("✅ Episode relationships intact")
        
        return True
        
    except sqlite3.Error as e:
        logger.error(f"Verification error: {e}")
        return False
    finally:
        if conn:
            conn.close()


if __name__ == "__main__":
    print("VFX Project Management - Shot Project ID Rollback")
    print("=" * 60)
    
    success = rollback_migration()
    
    if success:
        print("\n" + "=" * 60)
        print("VERIFYING ROLLBACK")
        print("=" * 60)
        
        if verify_rollback():
            print("\n✅ SUCCESS: Rollback completed successfully!")
            sys.exit(0)
        else:
            print("\n❌ FAILED: Rollback verification failed!")
            sys.exit(1)
    else:
        print("\n❌ FAILED: Rollback failed!")
        sys.exit(1)