Refactor DB code