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 test_execute_op_postgres(self, mock_cursor):
execute_op(DEFAULT_POSTGRES_OP, "some_id")
# correctly routes to postgres
mock_cursor.assert_called_once()
def test_execute_op_clickhouse(self, mock_sync_execute):
execute_op(DEFAULT_CH_OP, "some_id")
# correctly routes to ch
mock_sync_execute.assert_called_once_with(
"/* some_id */ SELECT 1", settings={"max_execution_time": 10})