Refactor DB connect