Expand migrations CLI to support create/run/rollback/status

Author: abrookins

.codespellrc

@@ -0,0 +1,3 @@
+[codespell]
+skip = .git,poetry.lock,*.pyc,__pycache__
+ignore-words-list = redis,migrator,datetime,timestamp,asyncio,redisearch,pydantic,ulid,hnsw

.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.2.6
+    hooks:
+      - id: codespell
+        args: [--write-changes]
+        exclude: ^(poetry\.lock|\.git/|docs/.*\.md)$

aredis_om/cli/main.py

@@ -1,5 +1,5 @@
 """
-Redis-OM CLI - Main entry point for the async 'om' command.
+Redis OMCLI - Main entry point for the async 'om' command.
 """
 
 import click
@@ -11,7 +11,7 @@
 @click.group()
 @click.version_option()
 def om():
-    """Redis-OM Python CLI - Object mapping and migrations for Redis."""
+    """Redis OM Python CLI - Object mapping and migrations for Redis."""
     pass
 
 

aredis_om/model/cli/migrate.py

@@ -1,18 +1,204 @@
+import asyncio
+import os
+
 import click
 
-from aredis_om.model.migrations.migrator import Migrator
+from aredis_om.model.migrations.schema_migrator import SchemaMigrator
+from aredis_om.settings import get_root_migrations_dir
+
+
+def run_async(coro):
+    """Run an async coroutine in an isolated event loop to avoid interfering with pytest loops."""
+    import concurrent.futures
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future = executor.submit(asyncio.run, coro)
+        return future.result()
+
+
+@click.group()
+def migrate():
+    """Manage schema migrations for Redis OM models."""
+    pass
+
+
+@migrate.command()
+@click.option("--migrations-dir", help="Directory containing schema migration files")
+def status(migrations_dir: str | None):
+    """Show current schema migration status from files."""
+
+    async def _status():
+        dir_path = migrations_dir or os.path.join(
+            get_root_migrations_dir(), "schema-migrations"
+        )
+        migrator = SchemaMigrator(migrations_dir=dir_path)
+        status_info = await migrator.status()
+
+        click.echo("Schema Migration Status:")
+        click.echo(f"  Total migrations: {status_info['total_migrations']}")
+        click.echo(f"  Applied: {status_info['applied_count']}")
+        click.echo(f"  Pending: {status_info['pending_count']}")
+
+        if status_info["pending_migrations"]:
+            click.echo("\nPending migrations:")
+            for migration_id in status_info["pending_migrations"]:
+                click.echo(f"- {migration_id}")
+
+        if status_info["applied_migrations"]:
+            click.echo("\nApplied migrations:")
+            for migration_id in status_info["applied_migrations"]:
+                click.echo(f"- {migration_id}")
+
+    run_async(_status())
+
+
+@migrate.command()
+@click.option("--migrations-dir", help="Directory containing schema migration files")
+@click.option(
+    "--dry-run", is_flag=True, help="Show what would be done without applying changes"
+)
+@click.option("--verbose", "-v", is_flag=True, help="Enable verbose output")
+@click.option("--limit", type=int, help="Limit number of migrations to run")
+@click.option(
+    "--yes",
+    "-y",
+    is_flag=True,
+    help="Skip confirmation prompt to create directory or run",
+)
+def run(
+    migrations_dir: str | None,
+    dry_run: bool,
+    verbose: bool,
+    limit: int | None,
+    yes: bool,
+):
+    """Run pending schema migrations from files."""
+
+    async def _run():
+        dir_path = migrations_dir or os.path.join(
+            get_root_migrations_dir(), "schema-migrations"
+        )
+
+        if not os.path.exists(dir_path):
+            if yes or click.confirm(
+                f"Create schema migrations directory at '{dir_path}'?"
+            ):
+                os.makedirs(dir_path, exist_ok=True)
+            else:
+                click.echo("Aborted.")
+                return
+
+        migrator = SchemaMigrator(migrations_dir=dir_path)
+
+        # Show list for confirmation
+        if not dry_run and not yes:
+            status_info = await migrator.status()
+            if status_info["pending_migrations"]:
+                listing = "\n".join(
+                    f"- {m}"
+                    for m in status_info["pending_migrations"][
+                        : (limit or len(status_info["pending_migrations"]))
+                    ]
+                )
+                if not click.confirm(
+                    f"Run {min(limit or len(status_info['pending_migrations']), len(status_info['pending_migrations']))} migration(s)?\n{listing}"
+                ):
+                    click.echo("Aborted.")
+                    return
+
+        count = await migrator.run(dry_run=dry_run, limit=limit, verbose=verbose)
+        if verbose and not dry_run:
+            click.echo(f"Successfully applied {count} migration(s).")
+
+    run_async(_run())
+
+
+@migrate.command()
+@click.argument("name")
+@click.option("--migrations-dir", help="Directory to create migration in")
+@click.option(
+    "--yes", "-y", is_flag=True, help="Skip confirmation prompt to create directory"
+)
+def create(name: str, migrations_dir: str | None, yes: bool):
+    """Create a new schema migration snapshot file from current pending operations."""
+
+    async def _create():
+        dir_path = migrations_dir or os.path.join(
+            get_root_migrations_dir(), "schema-migrations"
+        )
+
+        if not os.path.exists(dir_path):
+            if yes or click.confirm(
+                f"Create schema migrations directory at '{dir_path}'?"
+            ):
+                os.makedirs(dir_path, exist_ok=True)
+            else:
+                click.echo("Aborted.")
+                return
+
+        migrator = SchemaMigrator(migrations_dir=dir_path)
+        filepath = await migrator.create_migration_file(name)
+        if filepath:
+            click.echo(f"Created migration: {filepath}")
+        else:
+            click.echo("No pending schema changes detected. Nothing to snapshot.")
+
+    run_async(_create())
+
+
+@migrate.command()
+@click.argument("migration_id")
+@click.option("--migrations-dir", help="Directory containing schema migration files")
+@click.option(
+    "--dry-run", is_flag=True, help="Show what would be done without applying changes"
+)
+@click.option("--verbose", "-v", is_flag=True, help="Enable verbose output")
+@click.option(
+    "--yes",
+    "-y",
+    is_flag=True,
+    help="Skip confirmation prompt to create directory or run",
+)
+def rollback(
+    migration_id: str,
+    migrations_dir: str | None,
+    dry_run: bool,
+    verbose: bool,
+    yes: bool,
+):
+    """Rollback a specific schema migration by ID."""
+
+    async def _rollback():
+        dir_path = migrations_dir or os.path.join(
+            get_root_migrations_dir(), "schema-migrations"
+        )
+
+        if not os.path.exists(dir_path):
+            if yes or click.confirm(
+                f"Create schema migrations directory at '{dir_path}'?"
+            ):
+                os.makedirs(dir_path, exist_ok=True)
+            else:
+                click.echo("Aborted.")
+                return
 
+        migrator = SchemaMigrator(migrations_dir=dir_path)
 
-@click.command()
-@click.option("--module", default="aredis_om")
-def migrate(module: str):
-    migrator = Migrator(module)
-    migrator.detect_migrations()
+        if not yes and not dry_run:
+            if not click.confirm(f"Rollback migration '{migration_id}'?"):
+                click.echo("Aborted.")
+                return
 
-    if migrator.migrations:
-        print("Pending migrations:")
-        for migration in migrator.migrations:
-            print(migration)
+        success = await migrator.rollback(
+            migration_id, dry_run=dry_run, verbose=verbose
+        )
+        if success:
+            if verbose:
+                click.echo(f"Successfully rolled back migration: {migration_id}")
+        else:
+            click.echo(
+                f"Migration '{migration_id}' does not support rollback or is not applied.",
+                err=True,
+            )
 
-        if input("Run migrations? (y/n) ") == "y":
-            migrator.run()
+    run_async(_rollback())