Updated migration to match final state