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)