package models import ( "github.com/dgraph-io/badger/v3" "go.uber.org/zap" ) type Persister struct { log *zap.SugaredLogger db *badger.DB } func NewPersister(log *zap.SugaredLogger, db *badger.DB) *Persister { aso := &Persister{log: log, db: db} return aso } func (p *Persister) Store(model model) error { err := p.db.Update(model.Save) return err } func (p *Persister) Count(model model) (int, error) { opts := badger.DefaultIteratorOptions opts.PrefetchValues = false var count int err := p.db.View(func(txn *badger.Txn) error { it := txn.NewIterator(opts) defer it.Close() prefix := []byte(model.Keybase()) for it.Seek(prefix); it.ValidForPrefix([]byte(prefix)); it.Next() { count++ } return nil }) return count, err } func (p *Persister) Collect(model model) ([]string, error) { var result []string err := p.db.View(func(txn *badger.Txn) error { opts := badger.DefaultIteratorOptions opts.PrefetchValues = false // TODO Maybe we want true here it := txn.NewIterator(opts) defer it.Close() prefix := []byte(model.Keybase()) for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() { item := it.Item() item.Value(func(v []byte) error { result = append(result, string(v)) return nil }) } return nil }) return result, err }