1package pebble
  2
  3import (
  4	"bytes"
  5	"encoding/json"
  6
  7	"github.com/cockroachdb/pebble"
  8	"github.com/iden3/go-merkletree/db"
  9	log "github.com/sirupsen/logrus"
 10)
 11
 12// PebbleStorage implements the db.Storage interface
 13type PebbleStorage struct {
 14	pdb    *pebble.DB
 15	prefix []byte
 16}
 17
 18// PebbleStorageTx implements the db.Tx interface
 19type PebbleStorageTx struct {
 20	*PebbleStorage
 21	batch *pebble.Batch
 22}
 23
 24// NewPebbleStorage returns a new PebbleStorage
 25func NewPebbleStorage(path string, errorIfMissing bool) (*PebbleStorage, error) {
 26	o := &pebble.Options{
 27		ErrorIfNotExists: errorIfMissing,
 28	}
 29	rdb, err := pebble.Open(path, o)
 30	if err != nil {
 31		return nil, err
 32	}
 33	return &PebbleStorage{rdb, []byte{}}, nil
 34}
 35
 36type storageInfo struct {
 37	KeyCount   int
 38	ClaimCount int
 39}
 40
 41// Info implements the method Info of the interface db.Storage
 42func (p *PebbleStorage) Info() string {
 43	keycount := 0
 44	claimcount := 0
 45	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
 46		if value[0] == byte(1) {
 47			claimcount++
 48		}
 49
 50		keycount++
 51		return true, nil
 52	})
 53	if err != nil {
 54		return err.Error()
 55	}
 56	json, _ := json.MarshalIndent(
 57		storageInfo{
 58			KeyCount:   keycount,
 59			ClaimCount: claimcount,
 60		},
 61		"", "  ",
 62	)
 63	return string(json)
 64}
 65
 66// WithPrefix implements the method WithPrefix of the interface db.Storage
 67func (p *PebbleStorage) WithPrefix(prefix []byte) db.Storage {
 68	return &PebbleStorage{p.pdb, db.Concat(p.prefix, prefix)}
 69}
 70
 71// NewTx implements the method NewTx of the interface db.Storage
 72func (p *PebbleStorage) NewTx() (db.Tx, error) {
 73	return &PebbleStorageTx{p, p.pdb.NewIndexedBatch()}, nil
 74}
 75
 76// Get retreives a value from a key in the db.Storage
 77func (p *PebbleStorage) Get(key []byte) ([]byte, error) {
 78	v, closer, err := p.pdb.Get(db.Concat(p.prefix, key[:]))
 79	if err == pebble.ErrNotFound {
 80		return nil, db.ErrNotFound
 81	}
 82	closer.Close()
 83	return v, err
 84}
 85
 86// Iterate implements the method Iterate of the interface db.Storage
 87func (p *PebbleStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
 88	// NewIter already provides a point-in-time view of the current DB
 89	// state, but if is used for long term (is not the case), should use an
 90	// iterator over an snapshot:
 91	// snapshot := p.pdb.NewSnapshot()
 92	// defer snapshot.Close()
 93	// iter := snapshot.NewIter(nil)
 94	iter := p.pdb.NewIter(nil)
 95	defer iter.Close()
 96
 97	iter.First() // move the iterator to the first key/value pair
 98	if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) {
 99	} else {
100		localKey := iter.Key()[len(p.prefix):]
101		if _, err := f(localKey, iter.Value()); err != nil {
102			return err
103		}
104	}
105	for iter.Next() {
106		if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) {
107			continue
108		}
109		localKey := iter.Key()[len(p.prefix):]
110		if cont, err := f(localKey, iter.Value()); err != nil {
111			return err
112		} else if !cont {
113			break
114		}
115	}
116	return iter.Error()
117}
118
119// Get retreives a value from a key in the interface db.Tx
120func (tx *PebbleStorageTx) Get(key []byte) ([]byte, error) {
121	var err error
122
123	fullkey := db.Concat(tx.prefix, key)
124
125	v, closer, err := tx.batch.Get(fullkey)
126	if err == pebble.ErrNotFound {
127		return nil, db.ErrNotFound
128	}
129	closer.Close()
130
131	return v, err
132}
133
134// Put saves a key:value into the db.Storage
135func (tx *PebbleStorageTx) Put(k, v []byte) error {
136	return tx.batch.Set(db.Concat(tx.prefix, k[:]), v, nil)
137}
138
139// Add implements the method Add of the interface db.Tx
140func (tx *PebbleStorageTx) Add(atx db.Tx) error {
141	patx := atx.(*PebbleStorageTx)
142	return tx.batch.Apply(patx.batch, nil)
143}
144
145// Commit implements the method Commit of the interface db.Tx
146func (tx *PebbleStorageTx) Commit() error {
147	return tx.batch.Commit(nil)
148}
149
150// Close implements the method Close of the interface db.Tx
151func (tx *PebbleStorageTx) Close() {
152	_ = tx.batch.Close()
153}
154
155// Close implements the method Close of the interface db.Storage
156func (p *PebbleStorage) Close() {
157	if err := p.pdb.Close(); err != nil {
158		panic(err)
159	}
160	log.Info("Database closed")
161}
162
163// Pebble is an extra method that returns the *pebble.DB
164func (p *PebbleStorage) Pebble() *pebble.DB {
165	return p.pdb
166}
167
168// List implements the method List of the interface db.Storage
169func (p *PebbleStorage) List(limit int) ([]db.KV, error) {
170	ret := []db.KV{}
171	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
172		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
173		if len(ret) == limit {
174			return false, nil
175		}
176		return true, nil
177	})
178	return ret, err
179}