Improve and fix upgrade hooks code

Make it easier and safer to add an upgrade hook.
It also fixes first upgrade hook which was totally broken.