Cleanup code for update/add