""" Database Migration Management Script Provides helper functions for managing database migrations with Alembic """ import os import sys import logging from pathlib import Path # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from alembic.config import Config from alembic import command logger = logging.getLogger(__name__) def get_alembic_config() -> Config: """Get Alembic configuration""" # Path to alembic.ini alembic_ini = Path(__file__).parent.parent.parent / "alembic.ini" if not alembic_ini.exists(): raise FileNotFoundError(f"alembic.ini not found at {alembic_ini}") config = Config(str(alembic_ini)) return config def create_migration(message: str): """Create a new migration""" config = get_alembic_config() command.revision(config, message=message, autogenerate=True) logger.info(f"Created migration: {message}") def upgrade_database(revision: str = "head"): """Upgrade database to a revision""" config = get_alembic_config() command.upgrade(config, revision) logger.info(f"Upgraded database to {revision}") def downgrade_database(revision: str): """Downgrade database to a revision""" config = get_alembic_config() command.downgrade(config, revision) logger.info(f"Downgraded database to {revision}") def show_current_revision(): """Show current database revision""" config = get_alembic_config() command.current(config) def show_migration_history(): """Show migration history""" config = get_alembic_config() command.history(config) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Database Migration Management") subparsers = parser.add_subparsers(dest="command", help="Command to run") # Create migration create_parser = subparsers.add_parser("create", help="Create a new migration") create_parser.add_argument("message", help="Migration message") # Upgrade database upgrade_parser = subparsers.add_parser("upgrade", help="Upgrade database") upgrade_parser.add_argument( "--revision", default="head", help="Revision to upgrade to (default: head)" ) # Downgrade database downgrade_parser = subparsers.add_parser("downgrade", help="Downgrade database") downgrade_parser.add_argument("revision", help="Revision to downgrade to") # Show current revision subparsers.add_parser("current", help="Show current database revision") # Show history subparsers.add_parser("history", help="Show migration history") args = parser.parse_args() logging.basicConfig(level=logging.INFO) if args.command == "create": create_migration(args.message) elif args.command == "upgrade": upgrade_database(args.revision) elif args.command == "downgrade": downgrade_database(args.revision) elif args.command == "current": show_current_revision() elif args.command == "history": show_migration_history() else: parser.print_help()