More migration code fixes