Updated code to reflect change in db schema