1  import pathlib
  2  from typing import Optional
  3  
  4  from typer import Argument
  5  from typer import Option
  6  from typer import echo
  7  
  8  from evidently.cli.main import app
  9  
 10  
 11  def run_cli_migrations(
 12      database_url: str,
 13      revision: str = "head",
 14      downgrade: bool = False,
 15      autogenerate: bool = False,
 16      message: Optional[str] = None,
 17  ) -> None:
 18      """Run Alembic migrations for SQL storage.
 19  
 20      Args:
 21          database_url: Database connection URL (e.g., 'postgresql://user:pass@localhost/db')
 22          revision: Revision to upgrade/downgrade to (default: 'head')
 23          downgrade: If True, downgrade instead of upgrade
 24          autogenerate: If True, create a new migration from model changes
 25          message: Message for new migration (required if autogenerate=True)
 26      """
 27      from evidently.ui.service.storage.sql.utils import run_migrations
 28  
 29      if autogenerate:
 30          if not message:
 31              raise ValueError("--message is required when using --autogenerate")
 32          echo(f"Creating new migration: {message}")
 33          run_migrations(database_url, revision, downgrade, autogenerate, message)
 34          echo("Migration file created. Please review it before applying.")
 35      elif downgrade:
 36          echo(f"Downgrading to revision: {revision}")
 37          run_migrations(database_url, revision, downgrade, autogenerate, message)
 38          echo("Downgrade complete.")
 39      else:
 40          echo(f"Upgrading to revision: {revision}")
 41          run_migrations(database_url, revision, downgrade, autogenerate, message)
 42          echo("Upgrade complete.")
 43  
 44  
 45  @app.command("migrate")
 46  def migrate(
 47      database_url: str = Argument(..., help="Database connection URL (e.g., 'postgresql://user:pass@localhost/db')"),
 48      revision: str = Option("head", help="Revision to upgrade/downgrade to"),
 49      downgrade: bool = Option(False, "--downgrade", "-d", help="Downgrade instead of upgrade"),
 50      autogenerate: bool = Option(False, "--autogenerate", "-a", help="Create new migration from model changes"),
 51      message: Optional[str] = Option(
 52          None, "--message", "-m", help="Message for new migration (required with --autogenerate)"
 53      ),
 54  ):
 55      """Run database migrations for SQL storage.
 56  
 57      Examples:
 58          # Upgrade to latest migration
 59          evidently migrate postgresql://user:pass@localhost/evidently
 60  
 61          # Upgrade to specific revision
 62          evidently migrate postgresql://user:pass@localhost/evidently --revision abc123
 63  
 64          # Create new migration from model changes
 65          evidently migrate postgresql://user:pass@localhost/evidently --autogenerate -m "add new column"
 66  
 67          # Downgrade one revision
 68          evidently migrate postgresql://user:pass@localhost/evidently --downgrade --revision -1
 69      """
 70      run_cli_migrations(
 71          database_url=database_url, revision=revision, downgrade=downgrade, autogenerate=autogenerate, message=message
 72      )
 73  
 74  
 75  @app.command("migrate-status")
 76  def migrate_status(
 77      database_url: str = Argument(..., help="Database connection URL"),
 78  ):
 79      """Check current migration status."""
 80      try:
 81          from alembic import command
 82          from alembic.config import Config
 83      except ImportError as e:
 84          raise ImportError(
 85              "Alembic is required for migrations. Please install it with: pip install evidently[sql]"
 86          ) from e
 87  
 88      import evidently.ui.service.storage.sql.migrations
 89  
 90      migrations_dir = pathlib.Path(evidently.ui.service.storage.sql.migrations.__file__).parent
 91      alembic_ini = migrations_dir / "alembic.ini"
 92  
 93      config = Config(str(alembic_ini))
 94      config.set_main_option("sqlalchemy.url", database_url)
 95      config.set_section_option("alembic", "script_location", str(migrations_dir))
 96      # Add the evidently src directory to sys.path so imports work
 97      import evidently
 98  
 99      evidently_src = pathlib.Path(evidently.__file__).parent.parent
100      config.set_section_option("alembic", "prepend_sys_path", str(evidently_src))
101  
102      echo("Current database revision:")
103      command.current(config, verbose=True)
104  
105      echo("\nMigration history:")
106      command.history(config, verbose=True)