orion/scripts/squash_migrations.py

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

This script squashes all existing migrations into fresh per-module initial migrations.
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 from all version_locations to a timestamped backup
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 files
2. Test with: make db-reset (drops schema, runs all migrations, seeds data)
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"

# All migration version directories (core + modules)
MODULE_MIGRATION_DIRS = [
    project_root / "alembic" / "versions",
    project_root / "app" / "modules" / "billing" / "migrations" / "versions",
    project_root / "app" / "modules" / "cart" / "migrations" / "versions",
    project_root / "app" / "modules" / "catalog" / "migrations" / "versions",
    project_root / "app" / "modules" / "cms" / "migrations" / "versions",
    project_root / "app" / "modules" / "customers" / "migrations" / "versions",
    project_root / "app" / "modules" / "dev_tools" / "migrations" / "versions",
    project_root / "app" / "modules" / "inventory" / "migrations" / "versions",
    project_root / "app" / "modules" / "loyalty" / "migrations" / "versions",
    project_root / "app" / "modules" / "marketplace" / "migrations" / "versions",
    project_root / "app" / "modules" / "messaging" / "migrations" / "versions",
    project_root / "app" / "modules" / "orders" / "migrations" / "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 from all version locations."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_dir = project_root / "alembic" / f"versions_backup_{timestamp}"
    backup_dir.mkdir(parents=True, exist_ok=True)

    total_backed_up = 0

    for versions_dir in MODULE_MIGRATION_DIRS:
        if not versions_dir.exists():
            continue

        migration_files = [f for f in versions_dir.glob("*.py") if f.name != "__init__.py"]
        if not migration_files:
            continue

        # Create a subdirectory in backup matching the source path
        rel_path = versions_dir.relative_to(project_root)
        target_dir = backup_dir / str(rel_path).replace("/", "_")
        target_dir.mkdir(parents=True, exist_ok=True)

        for f in migration_files:
            shutil.copy2(f, target_dir / f.name)
            total_backed_up += 1

        # Remove migration files from source (keep __init__.py)
        for f in migration_files:
            f.unlink()

    if total_backed_up > 0:
        print(f"Backed up {total_backed_up} migrations to {backup_dir.name}/")
        return backup_dir
    print("No migration files found to backup")
    backup_dir.rmdir()
    return None


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")

    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 files")
    print("2. On a fresh database, run: make db-reset")
    print("3. Verify all tables are created correctly")
    print("4. If satisfied, delete the backup directory")
    print("")
    if backup_dir:
        print("To restore from backup:")
        print(f"  Check {backup_dir}/ for backed up migration files")


if __name__ == "__main__":
    main()