Add enhancements

snuba/cli/migrations.py

@@ -37,7 +37,7 @@ def list() -> None:
     setup_logging()
     check_clickhouse_connections(CLUSTERS)
     runner = Runner()
-    for group, group_migrations in runner.show_all():
+    for group, group_migrations in runner.show_all(include_nonexistent=True):
         readiness_state = get_group_readiness_state(group)
         click.echo(f"{group.value} (readiness_state: {readiness_state.value})")
         for migration_id, status, blocking, existing in group_migrations:

snuba/migrations/runner.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from datetime import datetime
 from functools import partial
 from typing import List, Mapping, MutableMapping, NamedTuple, Optional, Sequence, Tuple
@@ -134,7 +135,7 @@ def force_overwrite_status(
         )
 
     def show_all(
-        self, groups: Optional[Sequence[str]] = None
+        self, groups: Optional[Sequence[str]] = None, include_nonexistent: bool = False
     ) -> List[Tuple[MigrationGroup, List[MigrationDetails]]]:
         """
         Returns the list of migrations and their statuses for each group.
@@ -149,9 +150,9 @@ def show_all(
             migration_groups = get_active_migration_groups()
 
         migration_status = self._get_migration_status(migration_groups)
-        clickhouse_group_migrations: MutableMapping[MigrationGroup, List[str]] = {}
+        clickhouse_group_migrations = defaultdict(set)
         for group, migration_id in migration_status.keys():
-            clickhouse_group_migrations.setdefault(group, []).append(migration_id)
+            clickhouse_group_migrations[group].add(migration_id)
 
         def get_status(migration_key: MigrationKey) -> Status:
             return migration_status.get(migration_key, Status.NOT_STARTED)
@@ -173,16 +174,17 @@ def get_status(migration_key: MigrationKey) -> Status:
                     )
                 )
 
-            non_existing_migrations = set(
-                clickhouse_group_migrations.get(group, [])
-            ).difference(set(migration_ids))
-            for migration_id in non_existing_migrations:
-                migration_key = MigrationKey(group, migration_id)
-                group_migrations.append(
-                    MigrationDetails(
-                        migration_id, get_status(migration_key), False, False
+            if include_nonexistent:
+                non_existing_migrations = clickhouse_group_migrations.get(
+                    group, []
+                ).difference(set(migration_ids))
+                for migration_id in non_existing_migrations:
+                    migration_key = MigrationKey(group, migration_id)
+                    group_migrations.append(
+                        MigrationDetails(
+                            migration_id, get_status(migration_key), False, False
+                        )
                     )
-                )
 
             migrations.append((group, group_migrations))
 

tests/migrations/test_runner.py

@@ -110,6 +110,26 @@ def test_show_all_for_groups() -> None:
     assert all([migration.status == Status.COMPLETED for migration in migrations])
 
 
+@pytest.mark.clickhouse_db
+def test_show_all_nonexistent_migration() -> None:
+    runner = Runner()
+    assert all(
+        [
+            migration.status == Status.NOT_STARTED
+            for (_, group_migrations) in runner.show_all()
+            for migration in group_migrations
+        ]
+    )
+    runner.run_all(force=True)
+    assert all(
+        [
+            migration.status == Status.COMPLETED
+            for (_, group_migrations) in runner.show_all()
+            for migration in group_migrations
+        ]
+    )
+
+
 @pytest.mark.clickhouse_db
 def test_run_migration() -> None:
     runner = Runner()