More upgrade refactors