From e256e2cd6bfe4b82028d8c9f71cbfefc5451c2ad Mon Sep 17 00:00:00 2001 From: David Tsukernik Date: Thu, 12 Sep 2024 15:40:33 -0700 Subject: [PATCH 1/4] Update migrations list command to show migrations that no longer exist in the codebase --- snuba/cli/migrations.py | 8 ++++++-- snuba/migrations/runner.py | 23 +++++++++++++++++++++-- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/snuba/cli/migrations.py b/snuba/cli/migrations.py index 95d180cfaa..d0fdb2d484 100644 --- a/snuba/cli/migrations.py +++ b/snuba/cli/migrations.py @@ -40,7 +40,7 @@ def list() -> None: for group, group_migrations in runner.show_all(): readiness_state = get_group_readiness_state(group) click.echo(f"{group.value} (readiness_state: {readiness_state.value})") - for migration_id, status, blocking in group_migrations: + for migration_id, status, blocking, existing in group_migrations: symbol = { Status.COMPLETED: "X", Status.NOT_STARTED: " ", @@ -53,7 +53,11 @@ def list() -> None: if status != Status.COMPLETED and blocking: blocking_text = " (blocking)" - click.echo(f"[{symbol}] {migration_id}{in_progress_text}{blocking_text}") + existing_text = "" if existing else " (this migration no longer exists)" + + click.echo( + f"[{symbol}] {migration_id}{in_progress_text}{blocking_text}{existing_text}" + ) click.echo() diff --git a/snuba/migrations/runner.py b/snuba/migrations/runner.py index 445e5f4ee6..27089c2f1d 100644 --- a/snuba/migrations/runner.py +++ b/snuba/migrations/runner.py @@ -64,6 +64,7 @@ class MigrationDetails(NamedTuple): migration_id: str status: Status blocking: bool + exists: bool class Runner: @@ -148,6 +149,9 @@ def show_all( migration_groups = get_active_migration_groups() migration_status = self._get_migration_status(migration_groups) + clickhouse_group_migrations = {} + for group, migration_id in migration_status.keys(): + clickhouse_group_migrations.setdefault(group, []).append(migration_id) def get_status(migration_key: MigrationKey) -> Status: return migration_status.get(migration_key, Status.NOT_STARTED) @@ -156,12 +160,27 @@ def get_status(migration_key: MigrationKey) -> Status: group_migrations: List[MigrationDetails] = [] group_loader = get_group_loader(group) - for migration_id in group_loader.get_migrations(): + migration_ids = group_loader.get_migrations() + for migration_id in migration_ids: migration_key = MigrationKey(group, migration_id) migration = group_loader.load_migration(migration_id) group_migrations.append( MigrationDetails( - migration_id, get_status(migration_key), migration.blocking + migration_id, + get_status(migration_key), + migration.blocking, + True, + ) + ) + + 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 ) ) From 848cc938b7f9fa72553a8aaa80c644cc8d5240f4 Mon Sep 17 00:00:00 2001 From: David Tsukernik Date: Thu, 12 Sep 2024 15:58:06 -0700 Subject: [PATCH 2/4] Fix typing issues --- snuba/admin/clickhouse/migration_checks.py | 2 +- snuba/admin/views.py | 2 +- snuba/migrations/runner.py | 2 +- .../test_migration_checks.py | 17 ++++++++++------- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/snuba/admin/clickhouse/migration_checks.py b/snuba/admin/clickhouse/migration_checks.py index 0439d56956..8fdd9e495b 100644 --- a/snuba/admin/clickhouse/migration_checks.py +++ b/snuba/admin/clickhouse/migration_checks.py @@ -100,7 +100,7 @@ def __init__( ).get_migrations() migration_statuses = {} - for migration_id, status, _ in migrations: + for migration_id, status, _, _ in migrations: migration_statuses[migration_id] = { "migration_id": migration_id, "status": status, diff --git a/snuba/admin/views.py b/snuba/admin/views.py index 91136e1646..9942b55419 100644 --- a/snuba/admin/views.py +++ b/snuba/admin/views.py @@ -208,7 +208,7 @@ def migrations_groups_list(group: str) -> Response: "status": status.value, "blocking": blocking, } - for migration_id, status, blocking in runner_group_migrations + for migration_id, status, blocking, _ in runner_group_migrations ] ), 200, diff --git a/snuba/migrations/runner.py b/snuba/migrations/runner.py index 27089c2f1d..1c5b8c8db6 100644 --- a/snuba/migrations/runner.py +++ b/snuba/migrations/runner.py @@ -149,7 +149,7 @@ def show_all( migration_groups = get_active_migration_groups() migration_status = self._get_migration_status(migration_groups) - clickhouse_group_migrations = {} + clickhouse_group_migrations: MutableMapping[MigrationGroup, List[str]] = {} for group, migration_id in migration_status.keys(): clickhouse_group_migrations.setdefault(group, []).append(migration_id) diff --git a/tests/admin/clickhouse_migrations/test_migration_checks.py b/tests/admin/clickhouse_migrations/test_migration_checks.py index 7ad0790759..198ef35dca 100644 --- a/tests/admin/clickhouse_migrations/test_migration_checks.py +++ b/tests/admin/clickhouse_migrations/test_migration_checks.py @@ -29,9 +29,9 @@ def group_loader() -> GroupLoader: RUN_MIGRATIONS: Sequence[MigrationDetails] = [ - MigrationDetails("0001", Status.COMPLETED, True), - MigrationDetails("0002", Status.NOT_STARTED, True), - MigrationDetails("0003", Status.NOT_STARTED, True), + MigrationDetails("0001", Status.COMPLETED, True, True), + MigrationDetails("0002", Status.NOT_STARTED, True, True), + MigrationDetails("0003", Status.NOT_STARTED, True, True), ] @@ -62,9 +62,9 @@ def test_status_checker_run( REVERSE_MIGRATIONS: Sequence[MigrationDetails] = [ - MigrationDetails("0001", Status.COMPLETED, True), - MigrationDetails("0002", Status.IN_PROGRESS, True), - MigrationDetails("0003", Status.NOT_STARTED, True), + MigrationDetails("0001", Status.COMPLETED, True, True), + MigrationDetails("0002", Status.IN_PROGRESS, True, True), + MigrationDetails("0003", Status.NOT_STARTED, True, True), ] @@ -155,7 +155,10 @@ def test_run_migration_checks_and_policies( mock_policy = Mock() checker = mock_checker() mock_runner.show_all.return_value = [ - (MigrationGroup("events"), [MigrationDetails("0001", Status.COMPLETED, True)]) + ( + MigrationGroup("events"), + [MigrationDetails("0001", Status.COMPLETED, True, True)], + ) ] mock_policy.can_run.return_value = policy_result[0] From 7ef42f8810607e491452da370549783d435ec893 Mon Sep 17 00:00:00 2001 From: David Tsukernik Date: Fri, 13 Sep 2024 11:07:12 -0700 Subject: [PATCH 3/4] Add enhancements --- snuba/cli/migrations.py | 2 +- snuba/migrations/runner.py | 26 ++++++++++++++------------ tests/migrations/test_runner.py | 20 ++++++++++++++++++++ 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/snuba/cli/migrations.py b/snuba/cli/migrations.py index d0fdb2d484..bacb266dbe 100644 --- a/snuba/cli/migrations.py +++ b/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: diff --git a/snuba/migrations/runner.py b/snuba/migrations/runner.py index 1c5b8c8db6..ca16bac91b 100644 --- a/snuba/migrations/runner.py +++ b/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)) diff --git a/tests/migrations/test_runner.py b/tests/migrations/test_runner.py index 44d096f014..6b1360dcf5 100644 --- a/tests/migrations/test_runner.py +++ b/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() From a0f04c02e482ce68aa10ba50e5899aafda8c8429 Mon Sep 17 00:00:00 2001 From: David Tsukernik Date: Fri, 13 Sep 2024 11:24:09 -0700 Subject: [PATCH 4/4] Fix mypy error --- snuba/migrations/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/snuba/migrations/runner.py b/snuba/migrations/runner.py index ca16bac91b..82a970a3b6 100644 --- a/snuba/migrations/runner.py +++ b/snuba/migrations/runner.py @@ -176,7 +176,7 @@ def get_status(migration_key: MigrationKey) -> Status: if include_nonexistent: non_existing_migrations = clickhouse_group_migrations.get( - group, [] + group, set() ).difference(set(migration_ids)) for migration_id in non_existing_migrations: migration_key = MigrationKey(group, migration_id)