Added a few checks to upgrade logic.