orion/scripts/squash_migrations.py

#!/usr/bin/env python3
"""
Migration Squash Script

This script squashes all existing migrations into a single initial migration.
Run this after setting up PostgreSQL to simplify the migration history.

Prerequisites:
- PostgreSQL must be running: make docker-up
- DATABASE_URL environment variable must be set to PostgreSQL

Usage:
    python scripts/squash_migrations.py

What this script does:
1. Backs up existing migrations to alembic/versions_backup_YYYYMMDD/
2. Creates a fresh initial migration from current models
3. Stamps the database as being at the new migration

After running:
1. Review the new migration in alembic/versions/
2. Test with: make migrate-up (on a fresh database)
3. If satisfied, delete the backup directory
"""

import os
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

VERSIONS_DIR = project_root / "alembic" / "versions"


def check_prerequisites():
    """Verify PostgreSQL is configured."""
    database_url = os.getenv("DATABASE_URL", "")

    if not database_url.startswith("postgresql"):
        print("ERROR: DATABASE_URL must be a PostgreSQL URL")
        print(f"Current: {database_url[:50]}...")
        print("")
        print("Set DATABASE_URL or start PostgreSQL with: make docker-up")
        sys.exit(1)

    print(f"Database: {database_url.split('@')[0]}@...")
    return True


def backup_migrations():
    """Backup existing migrations."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_dir = project_root / "alembic" / f"versions_backup_{timestamp}"

    if not VERSIONS_DIR.exists():
        print("No existing migrations to backup")
        return None

    migration_files = list(VERSIONS_DIR.glob("*.py"))
    if not migration_files:
        print("No migration files found")
        return None

    print(f"Backing up {len(migration_files)} migrations to {backup_dir.name}/")
    shutil.copytree(VERSIONS_DIR, backup_dir)

    # Clear versions directory (keep __pycache__ if exists)
    for f in VERSIONS_DIR.glob("*.py"):
        f.unlink()

    return backup_dir


def create_fresh_migration():
    """Generate fresh initial migration from models."""
    print("Generating fresh initial migration...")

    result = subprocess.run(
        [
            sys.executable, "-m", "alembic", "revision",
            "--autogenerate", "-m", "initial_postgresql_schema"
        ],
        cwd=project_root,
        capture_output=True,
        text=True
    )

    if result.returncode != 0:
        print("ERROR: Failed to generate migration")
        print(result.stderr)
        sys.exit(1)

    print(result.stdout)

    # Find the new migration file
    new_migrations = list(VERSIONS_DIR.glob("*initial_postgresql_schema*.py"))
    if new_migrations:
        print(f"Created: {new_migrations[0].name}")
        return new_migrations[0]

    return None


def clean_migration_file(migration_path: Path):
    """Remove SQLite-specific patterns from migration."""
    if not migration_path:
        return

    content = migration_path.read_text()

    # Remove batch_alter_table references (not needed for PostgreSQL)
    if "batch_alter_table" in content:
        print("Note: Migration contains batch_alter_table - this is not needed for PostgreSQL")
        # We don't auto-remove as it might be intentional

    print(f"Review migration at: {migration_path}")


def stamp_database():
    """Stamp the database as being at the new migration."""
    print("Stamping database with new migration...")

    result = subprocess.run(
        [sys.executable, "-m", "alembic", "stamp", "head"],
        cwd=project_root,
        capture_output=True,
        text=True
    )

    if result.returncode != 0:
        print("WARNING: Could not stamp database (may need to run migrate-up first)")
        print(result.stderr)
    else:
        print("Database stamped at head")


def main():
    print("=" * 60)
    print("MIGRATION SQUASH SCRIPT")
    print("=" * 60)
    print("")

    # Check prerequisites
    check_prerequisites()
    print("")

    # Confirm with user
    response = input("This will backup and replace all migrations. Continue? [y/N] ")
    if response.lower() != 'y':
        print("Aborted")
        sys.exit(0)

    print("")

    # Backup existing migrations
    backup_dir = backup_migrations()
    print("")

    # Create fresh migration
    new_migration = create_fresh_migration()
    print("")

    # Clean up the migration file
    clean_migration_file(new_migration)
    print("")

    # Summary
    print("=" * 60)
    print("SQUASH COMPLETE")
    print("=" * 60)
    print("")
    if backup_dir:
        print(f"Backup location: {backup_dir}")
    print("")
    print("Next steps:")
    print("1. Review the new migration file")
    print("2. On a fresh database, run: make migrate-up")
    print("3. Verify all tables are created correctly")
    print("4. If satisfied, delete the backup directory")
    print("")
    print("To restore from backup:")
    print(f"  rm -rf alembic/versions/*.py")
    print(f"  cp -r {backup_dir}/* alembic/versions/")


if __name__ == "__main__":
    main()