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