update code to fit the new data model