migrate.py

"""
Database migration script for QPyth Neon database.

This script initializes the database schema and can be run:
- Automatically on PR preview branch creation
- Manually for local development
- As part of CI/CD pipeline

Usage:
    python migrate.py              # Run all migrations
    python migrate.py --check      # Check current migration status
    python migrate.py --rollback   # Rollback last migration
"""

import argparse
import os
import sys
from datetime import datetime

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


def get_database_url() -> str:
    """Get database URL from environment."""
    database_url = os.getenv("DATABASE_URL")
    if not database_url:
        raise ValueError(
            "DATABASE_URL environment variable not set. "
            "Please set it or use a .env file."
        )
    return database_url


def run_migration(db) -> None:
    """
    Run database migrations.

    This creates all required tables if they don't exist.
    """
    print(f"[{datetime.now().isoformat()}] Starting database migration...")

    # Import database module
    from quantumpytho.modules.database import init_schema

    # Run schema initialization
    init_schema(db)

    print(f"[{datetime.now().isoformat()}] Migration completed successfully!")


def check_migration_status(db) -> None:
    """Check current migration status."""
    print(f"[{datetime.now().isoformat()}] Checking migration status...")

    # Check if tables exist
    tables = ["quantum_jobs", "vqe_results", "user_sessions", "api_logs"]

    for table in tables:
        result = db.fetch_one(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = %s)",
            (table,),
        )
        exists = result.get("exists", False) if result else False
        status = "✅" if exists else "❌"
        print(f"  {status} {table}")

    print(f"[{datetime.now().isoformat()}] Status check complete!")


def rollback_migration(db) -> None:
    """Rollback last migration (use with caution!)."""
    print(f"[{datetime.now().isoformat()}] WARNING: Rolling back migrations...")
    print("  This will DROP all tables!")

    confirm = input("  Are you sure? Type 'YES' to confirm: ")
    if confirm != "YES":
        print("  Rollback cancelled.")
        return

    # Drop all tables
    tables = ["api_logs", "user_sessions", "vqe_results", "quantum_jobs"]

    for table in tables:
        try:
            db.execute(f"DROP TABLE IF EXISTS {table} CASCADE")
            print(f"  ✅ Dropped table: {table}")
        except Exception as e:
            print(f"  ❌ Failed to drop {table}: {e}")

    print(f"[{datetime.now().isoformat()}] Rollback complete!")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="QPyth Database Migrations")
    parser.add_argument(
        "--check", action="store_true", help="Check current migration status"
    )
    parser.add_argument(
        "--rollback",
        action="store_true",
        help="Rollback last migration (DROPS all tables)",
    )
    args = parser.parse_args()

    # Get database URL
    try:
        database_url = get_database_url()
        print(f"Using database: {database_url.split('@')[1].split('/')[0]}")
    except ValueError as e:
        print(f"❌ Error: {e}")
        print("\nSet DATABASE_URL environment variable:")
        print("  export DATABASE_URL='postgresql://...'")
        print("\nOr create a .env file with:")
        print("  DATABASE_URL=postgresql://...")
        sys.exit(1)

    # Import and connect
    try:
        from quantumpytho.modules.database import DatabaseConfig, NeonDatabase

        config = DatabaseConfig(database_url=database_url)
        db = NeonDatabase(config)

        print("✅ Connected to database")

        # Run appropriate command
        if args.rollback:
            rollback_migration(db)
        elif args.check:
            check_migration_status(db)
        else:
            run_migration(db)

        # Close connection
        db.close()

    except Exception as e:
        print(f"❌ Migration failed: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()