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