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})