mcp/database/migrate.py

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