#!/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()