Refactor db logic