Fix rollback in migration