package models import ( "github.com/dgraph-io/badger/v3" "go.uber.org/zap" ) type Persister struct { log *zap.SugaredLogger db *badger.DB } type PersisterResult struct { Value []byte Key []byte } func NewPersister(log *zap.SugaredLogger, db *badger.DB) *Persister { logger := log.With("type", "db") aso := &Persister{log: logger, db: db} return aso } func (p *Persister) StoreWithCallback(m model, onSave func()) error { return p.storeModel(m, onSave) } func (p *Persister) Store(m model) error { return p.storeModel(m, func() {}) } func (p *Persister) storeModel(model model, onSave func()) error { log := p.log.With("model", model.Name()).With("DedupKey", model.DedupKey()).With("Key", model.Key()) log.Debug("Store()") err := p.db.View(func(txn *badger.Txn) error { var getErr error var updateErr error log.Debug("checking") _, getErr = txn.Get([]byte(model.DedupKey())) if getErr == badger.ErrKeyNotFound { log.Debug("not found") updateErr = p.db.Update(model.Save) // stores the outbox item if updateErr != nil { return updateErr } log.Debug("saving and calling storeCallback") updateErr = p.db.Update(model.SaveDedup) // stores the sha256 onSave() } return updateErr }) return err } func (p *Persister) Delete(model model) error { log := p.log.With("model", model.Name()).With("DedupKey", model.DedupKey()).With("Key", model.Key()) log.Debug("Delete()") err := p.db.Update(func(txn *badger.Txn) error { txn.Delete([]byte(model.Key())) return nil }) 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) ([][]byte, error) { var result [][]byte err := p.db.View(func(txn *badger.Txn) error { opts := badger.DefaultIteratorOptions opts.PrefetchValues = true 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, v) return nil }) } return nil }) return result, err } func (p *Persister) Find(model model) ([]byte, error) { var result []byte var item *badger.Item var notFound bool err := p.db.View(func(txn *badger.Txn) error { var getErr error item, getErr = txn.Get([]byte(model.Key())) if getErr == badger.ErrKeyNotFound { notFound = true return nil } if getErr != nil { return getErr } return nil }) if notFound { return nil, nil // on key not found we want to return nil, so that gin returns 404 and not 500 } if err != nil { return nil, err } err = item.Value(func(v []byte) error { result = v return nil }) if err != nil { return nil, err } return result, nil }