Add code for insert and update