Add back some upgrade code missed on rebase