Fix upgrade bugs