Little refactor in new migration