Fix upgrade stuff