Fix outstanding upgrade bugs