def run_async_migration(migration_name: str, fresh_start: bool = True) -> None: if fresh_start: start_async_migration(migration_name) return # Resumable operations run_async_migration_operations(migration_name)
def handle(self, *args, **options): setup_async_migrations(ignore_posthog_version=True) necessary_migrations = get_necessary_migrations() if options["plan"] or options["check"]: print() if len(necessary_migrations) == 0: print("Async migrations up to date!") return print("List of async migrations to be applied:") for migration in necessary_migrations: print(f"- {migration.name}") print() if options["check"]: exit(1) return for migration in necessary_migrations: logger.info(f"Applying async migration {migration.name}") started_successfully = start_async_migration(migration.name, ignore_posthog_version=True) migration.refresh_from_db() if not started_successfully or migration.status != MigrationStatus.CompletedSuccessfully: last_error = AsyncMigrationError.objects.filter(async_migration=migration).last() last_error_msg = f", last error: {last_error.description}" if last_error else "" logger.info(f"Unable to complete async migration {migration.name}{last_error_msg}.") raise ImproperlyConfigured( f"Migrate job failed because necessary async migration {migration.name} could not complete." ) logger.info(f"✅ Migration {migration.name} successful")
def test_rollback_migration(self): self.migration.sec.reset_count() migration_successful = start_async_migration("test") self.assertEqual(migration_successful, True) sm = AsyncMigration.objects.get(name="test") attempt_migration_rollback(sm) sm.refresh_from_db() exception = None try: with connection.cursor() as cursor: cursor.execute("SELECT * FROM test_async_migration") except Exception as e: exception = e self.assertIn('relation "test_async_migration" does not exist', str(exception)) self.assertEqual(sm.status, MigrationStatus.RolledBack) self.assertEqual(sm.progress, 0) self.assertEqual(self.migration.sec.side_effect_count, 3) self.assertEqual(self.migration.sec.side_effect_rollback_count, 3)
def test_run_migration_in_full(self): self.migration.sec.reset_count() migration_successful = start_async_migration("test") sm = AsyncMigration.objects.get(name="test") with connection.cursor() as cursor: cursor.execute("SELECT * FROM test_async_migration") res = cursor.fetchone() self.assertEqual(res, ("a", "c")) self.assertTrue(migration_successful) self.assertEqual(sm.name, "test") self.assertEqual(sm.description, self.TEST_MIGRATION_DESCRIPTION) self.assertEqual(sm.status, MigrationStatus.CompletedSuccessfully) self.assertEqual(sm.progress, 100) errors = AsyncMigrationError.objects.filter(async_migration=sm) self.assertEqual(errors.count(), 0) self.assertTrue(UUIDT.is_valid_uuid(sm.current_query_id)) self.assertEqual(sm.current_operation_index, 7) self.assertEqual(sm.posthog_min_version, "1.0.0") self.assertEqual(sm.posthog_max_version, "100000.0.0") self.assertEqual(sm.finished_at.day, datetime.today().day) self.assertEqual(self.migration.sec.side_effect_count, 3) self.assertEqual(self.migration.sec.side_effect_rollback_count, 0)
def test_run_migration_in_full(self): from ee.clickhouse.client import sync_execute migration_successful = start_async_migration(MIGRATION_NAME) sm = AsyncMigration.objects.get(name=MIGRATION_NAME) self.assertTrue(migration_successful) self.assertEqual(sm.status, MigrationStatus.CompletedSuccessfully) self.assertEqual(sm.progress, 100) self.assertEqual(sm.current_operation_index, 9) errors = AsyncMigrationError.objects.filter(async_migration=sm) self.assertEqual(errors.count(), 0) create_table_res = sync_execute( f"SHOW CREATE TABLE {CLICKHOUSE_DATABASE}.events") events_count_res = sync_execute( f"SELECT COUNT(*) FROM {CLICKHOUSE_DATABASE}.events") backup_events_count_res = sync_execute( f"SELECT COUNT(*) FROM {CLICKHOUSE_DATABASE}.events_backup_0002_events_sample_by" ) self.assertTrue( "ORDER BY (team_id, toDate(timestamp), event, cityHash64(distinct_id), cityHash64(uuid))" in create_table_res[0][0]) self.assertEqual(events_count_res[0][0], 5) self.assertEqual(backup_events_count_res[0][0], 5)
def test_migration(self): from posthog.client import sync_execute p1, p2, p3, p4, p5, p6 = [UUID(int=i) for i in range(6)] self.create_distinct_id(team_id=1, distinct_id="a", person_id=str(p1), sign=1) self.create_distinct_id(team_id=2, distinct_id="a", person_id=str(p2), sign=1) # Merged user self.create_distinct_id(team_id=2, distinct_id="b", person_id=str(p3), sign=1) self.create_distinct_id(team_id=2, distinct_id="b", person_id=str(p3), sign=-1) self.create_distinct_id(team_id=2, distinct_id="b", person_id=str(p4), sign=1) # Deleted user self.create_distinct_id(team_id=2, distinct_id="c", person_id=str(p5), sign=1) self.create_distinct_id(team_id=2, distinct_id="c", person_id=str(p5), sign=-1) self.create_distinct_id(team_id=3, distinct_id="d", person_id=str(p6), sign=1) setup_async_migrations() migration_successful = start_async_migration(MIGRATION_NAME) self.assertTrue(migration_successful) rows = sync_execute( "SELECT team_id, distinct_id, person_id, version FROM person_distinct_id2 ORDER BY team_id, distinct_id" ) self.assertEqual(rows, [(1, "a", p1, 0), (2, "a", p2, 0), (2, "b", p4, 0), (3, "d", p6, 0)])
def test_migration(self): # :TRICKY: Relies on tables being migrated as unreplicated before. _create_event(team=self.team, distinct_id="test", event="$pageview") _create_event(team=self.team, distinct_id="test2", event="$pageview") settings.CLICKHOUSE_REPLICATION = True setup_async_migrations() migration_successful = start_async_migration(MIGRATION_NAME) self.assertTrue(migration_successful) self.verify_table_engines_correct() self.assertEqual(self.get_event_table_row_count(), 2)
def test_rollback_migration_failure(self): migration_name = "test_with_rollback_exception" create_async_migration(name=migration_name) self.migration.sec.reset_count() migration_successful = start_async_migration(migration_name) self.assertEqual(migration_successful, True) sm = AsyncMigration.objects.get(name=migration_name) attempt_migration_rollback(sm) sm.refresh_from_db() self.assertEqual(sm.status, MigrationStatus.Errored) self.assertEqual(sm.current_operation_index, 1)
def test_rollback(self): # :TRICKY: Relies on tables being migrated as unreplicated before. _create_event(team=self.team, distinct_id="test", event="$pageview") _create_event(team=self.team, distinct_id="test2", event="$pageview") settings.CLICKHOUSE_REPLICATION = True setup_async_migrations() migration = get_async_migration_definition(MIGRATION_NAME) self.assertEqual(len(migration.operations), 53) migration.operations[30].sql = "THIS WILL FAIL!" # type: ignore migration_successful = start_async_migration(MIGRATION_NAME) self.assertFalse(migration_successful) self.assertEqual(AsyncMigration.objects.get(name=MIGRATION_NAME).status, MigrationStatus.RolledBack) self.verify_table_engines_correct(expected_engine_types=("ReplacingMergeTree", "CollapsingMergeTree", "Kafka"))
def test_migration(self): # :TRICKY: Relies on tables being migrated as unreplicated before. _create_event(team=self.team, distinct_id="test", event="$pageview") _create_event(team=self.team, distinct_id="test2", event="$pageview") settings.CLICKHOUSE_REPLICATION = True setup_async_migrations() migration_successful = start_async_migration(MIGRATION_NAME) self.assertTrue(migration_successful) self.verify_table_engines_correct( expected_engine_types=( "ReplicatedReplacingMergeTree", "ReplicatedCollapsingMergeTree", "Distributed", "Kafka", ) ) self.assertEqual(self.get_event_table_row_count(), 2)