def attempt_migration_rollback(migration_instance: AsyncMigration):
"""
Cycle through the operations in reverse order starting from the last completed op and run
the specified rollback statements.
"""
migration_instance.refresh_from_db()
ops = get_async_migration_definition(migration_instance.name).operations
# if the migration was completed the index is set 1 after, normally we should try rollback for current op
current_index = min(migration_instance.current_operation_index,
len(ops) - 1)
for op_index in range(current_index, -1, -1):
try:
op = ops[op_index]
execute_op(op, str(UUIDT()), rollback=True)
except Exception as e:
error = f"At operation {op_index} rollback failed with error:{str(e)}"
process_error(
migration_instance=migration_instance,
error=error,
rollback=False,
alert=True,
current_operation_index=op_index,
)
return
update_async_migration(migration_instance=migration_instance,
status=MigrationStatus.RolledBack,
progress=0)
def run_async_migration_next_op(
migration_name: str,
migration_instance: Optional[AsyncMigration] = None):
"""
Runs the next operation specified by the currently running migration
We run the next operation of the migration which needs attention
Returns (run_next, success)
Terminology:
- migration_instance: The migration object as stored in the DB
- migration_definition: The actual migration class outlining the operations (e.g. async_migrations/examples/example.py)
"""
if not migration_instance:
try:
migration_instance = AsyncMigration.objects.get(
name=migration_name, status=MigrationStatus.Running)
except AsyncMigration.DoesNotExist:
return (False, False)
else:
migration_instance.refresh_from_db()
assert migration_instance is not None
migration_definition = get_async_migration_definition(migration_name)
if migration_instance.current_operation_index > len(
migration_definition.operations) - 1:
complete_migration(migration_instance)
return (False, True)
error = None
current_query_id = str(UUIDT())
try:
op = migration_definition.operations[
migration_instance.current_operation_index]
execute_op(op, current_query_id)
update_async_migration(
migration_instance=migration_instance,
current_query_id=current_query_id,
current_operation_index=migration_instance.current_operation_index
+ 1,
)
except Exception as e:
error = f"Exception was thrown while running operation {migration_instance.current_operation_index} : {str(e)}"
process_error(migration_instance, error, alert=True)
if error:
return (False, False)
update_migration_progress(migration_instance)
return (True, False)
def update_migration_progress(migration_instance: AsyncMigration):
"""
We don't want to interrupt a migration if the progress check fails, hence try without handling exceptions
Progress is a nice-to-have bit of feedback about how the migration is doing, but not essential
"""
migration_instance.refresh_from_db()
try:
progress = get_async_migration_definition(migration_instance.name).progress(migration_instance) # type: ignore
update_async_migration(migration_instance=migration_instance, progress=progress)
except:
pass
def test_run_async_migration_next_op(self):
sm = AsyncMigration.objects.get(name="test")
update_async_migration(sm, status=MigrationStatus.Running)
run_async_migration_next_op("test", sm)
sm.refresh_from_db()
self.assertEqual(sm.current_operation_index, 1)
self.assertEqual(sm.progress, int(100 * 1 / 7))
run_async_migration_next_op("test", sm)
sm.refresh_from_db()
self.assertEqual(sm.current_operation_index, 2)
self.assertEqual(sm.progress, int(100 * 2 / 7))
run_async_migration_next_op("test", sm)
with connection.cursor() as cursor:
cursor.execute("SELECT * FROM test_async_migration")
res = cursor.fetchone()
self.assertEqual(res, ("a", "b"))
for i in range(5):
run_async_migration_next_op("test", sm)
sm.refresh_from_db()
self.assertEqual(sm.current_operation_index, 7)
self.assertEqual(sm.progress, 100)
self.assertEqual(sm.status, MigrationStatus.CompletedSuccessfully)
with connection.cursor() as cursor:
cursor.execute("SELECT * FROM test_async_migration")
res = cursor.fetchone()
self.assertEqual(res, ("a", "c"))