1package leveldb
  2
  3import (
  4	"encoding/json"
  5
  6	log "github.com/sirupsen/logrus"
  7	"github.com/syndtr/goleveldb/leveldb"
  8	"github.com/syndtr/goleveldb/leveldb/errors"
  9	"github.com/syndtr/goleveldb/leveldb/opt"
 10	"github.com/syndtr/goleveldb/leveldb/util"
 11
 12	"github.com/iden3/go-merkletree/db"
 13)
 14
 15// LevelDbStorage implements the db.Storage interface
 16type LevelDbStorage struct {
 17	ldb    *leveldb.DB
 18	prefix []byte
 19}
 20
 21// LevelDbStorageTx implements the db.Tx interface
 22type LevelDbStorageTx struct {
 23	*LevelDbStorage
 24	cache db.KvMap
 25}
 26
 27// NewLevelDbStorage returns a new LevelDbStorage
 28func NewLevelDbStorage(path string, errorIfMissing bool) (*LevelDbStorage, error) {
 29	o := &opt.Options{
 30		ErrorIfMissing: errorIfMissing,
 31	}
 32	ldb, err := leveldb.OpenFile(path, o)
 33	if err != nil {
 34		return nil, err
 35	}
 36	return &LevelDbStorage{ldb, []byte{}}, nil
 37}
 38
 39type storageInfo struct {
 40	KeyCount   int
 41	ClaimCount int
 42}
 43
 44// Info implements the method Info of the interface db.Storage
 45func (l *LevelDbStorage) Info() string {
 46	snapshot, err := l.ldb.GetSnapshot()
 47	if err != nil {
 48		return err.Error()
 49	}
 50
 51	keycount := 0
 52	claimcount := 0
 53	iter := snapshot.NewIterator(nil, nil)
 54	for iter.Next() {
 55		if iter.Value()[0] == byte(1) {
 56			claimcount++
 57		}
 58
 59		keycount++
 60	}
 61	iter.Release()
 62	if err := iter.Error(); err != nil {
 63		return err.Error()
 64	}
 65	json, _ := json.MarshalIndent(
 66		storageInfo{
 67			KeyCount:   keycount,
 68			ClaimCount: claimcount,
 69		},
 70		"", "  ",
 71	)
 72	return string(json)
 73}
 74
 75// WithPrefix implements the method WithPrefix of the interface db.Storage
 76func (l *LevelDbStorage) WithPrefix(prefix []byte) db.Storage {
 77	return &LevelDbStorage{l.ldb, db.Concat(l.prefix, prefix)}
 78}
 79
 80// NewTx implements the method NewTx of the interface db.Storage
 81func (l *LevelDbStorage) NewTx() (db.Tx, error) {
 82	return &LevelDbStorageTx{l, make(db.KvMap)}, nil
 83}
 84
 85// Get retreives a value from a key in the db.Storage
 86func (l *LevelDbStorage) Get(key []byte) ([]byte, error) {
 87	v, err := l.ldb.Get(db.Concat(l.prefix, key[:]), nil)
 88	if err == errors.ErrNotFound {
 89		return nil, db.ErrNotFound
 90	}
 91	return v, err
 92}
 93
 94// Iterate implements the method Iterate of the interface db.Storage
 95func (l *LevelDbStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
 96	// FIXME: Use the prefix!
 97	snapshot, err := l.ldb.GetSnapshot()
 98	if err != nil {
 99		return err
100	}
101	iter := snapshot.NewIterator(util.BytesPrefix(l.prefix), nil)
102	defer iter.Release()
103	for iter.Next() {
104		localKey := iter.Key()[len(l.prefix):]
105		if cont, err := f(localKey, iter.Value()); err != nil {
106			return err
107		} else if !cont {
108			break
109		}
110	}
111	iter.Release()
112	return iter.Error()
113}
114
115// Get retreives a value from a key in the interface db.Tx
116func (tx *LevelDbStorageTx) Get(key []byte) ([]byte, error) {
117	var err error
118
119	fullkey := db.Concat(tx.prefix, key)
120
121	if value, ok := tx.cache.Get(fullkey); ok {
122		return value, nil
123	}
124
125	value, err := tx.ldb.Get(fullkey, nil)
126	if err == errors.ErrNotFound {
127		return nil, db.ErrNotFound
128	}
129
130	return value, err
131}
132
133// Put saves a key:value into the db.Storage
134func (tx *LevelDbStorageTx) Put(k, v []byte) {
135	tx.cache.Put(db.Concat(tx.prefix, k[:]), v)
136}
137
138// Add implements the method Add of the interface db.Tx
139func (tx *LevelDbStorageTx) Add(atx db.Tx) {
140	ldbtx := atx.(*LevelDbStorageTx)
141	for _, v := range ldbtx.cache {
142		tx.cache.Put(v.K, v.V)
143	}
144}
145
146// Commit implements the method Commit of the interface db.Tx
147func (tx *LevelDbStorageTx) Commit() error {
148	var batch leveldb.Batch
149	for _, v := range tx.cache {
150		batch.Put(v.K, v.V)
151	}
152
153	tx.cache = nil
154	return tx.ldb.Write(&batch, nil)
155}
156
157// Close implements the method Close of the interface db.Tx
158func (tx *LevelDbStorageTx) Close() {
159	tx.cache = nil
160}
161
162// Close implements the method Close of the interface db.Storage
163func (l *LevelDbStorage) Close() {
164	if err := l.ldb.Close(); err != nil {
165		panic(err)
166	}
167	log.Info("Database closed")
168}
169
170// LevelDB is an extra method that returns the *leveldb.DB
171func (l *LevelDbStorage) LevelDB() *leveldb.DB {
172	return l.ldb
173}
174
175// List implements the method List of the interface db.Storage
176func (l *LevelDbStorage) List(limit int) ([]db.KV, error) {
177	ret := []db.KV{}
178	err := l.Iterate(func(key []byte, value []byte) (bool, error) {
179		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
180		if len(ret) == limit {
181			return false, nil
182		}
183		return true, nil
184	})
185	return ret, err
186}