/
badger.go
151 lines (127 loc) · 3.26 KB
/
badger.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
package store
import (
"errors"
"github.com/dgraph-io/badger/v3"
)
var _ KVStore = &BadgerKV{}
var _ Batch = &BadgerBatch{}
var (
// ErrKeyNotFound is returned if key is not found in KVStore.
ErrKeyNotFound = errors.New("key not found")
)
// BadgerKV is a implementation of KVStore using Badger v3.
type BadgerKV struct {
db *badger.DB
}
// Get returns value for given key, or error.
func (b *BadgerKV) Get(key []byte) ([]byte, error) {
txn := b.db.NewTransaction(false)
defer txn.Discard()
item, err := txn.Get(key)
if errors.Is(err, badger.ErrKeyNotFound) {
return nil, ErrKeyNotFound
}
if err != nil {
return nil, err
}
return item.ValueCopy(nil)
}
// Set saves key-value mapping in store.
func (b *BadgerKV) Set(key []byte, value []byte) error {
txn := b.db.NewTransaction(true)
err := txn.Set(key, value)
if err != nil {
txn.Discard()
return err
}
return txn.Commit()
}
// Delete removes key and corresponding value from store.
func (b *BadgerKV) Delete(key []byte) error {
txn := b.db.NewTransaction(true)
err := txn.Delete(key)
if err != nil {
txn.Discard()
return err
}
return txn.Commit()
}
// NewBatch creates new batch.
// Note: badger batches should be short lived as they use extra resources.
func (b *BadgerKV) NewBatch() Batch {
return &BadgerBatch{
txn: b.db.NewTransaction(true),
}
}
// BadgerBatch encapsulates badger transaction
type BadgerBatch struct {
txn *badger.Txn
}
// Set accumulates key-value entries in a transaction
func (bb *BadgerBatch) Set(key, value []byte) error {
if err := bb.txn.Set(key, value); err != nil {
return err
}
return nil
}
// Delete removes the key and associated value from store
func (bb *BadgerBatch) Delete(key []byte) error {
return bb.txn.Delete(key)
}
// Commit commits a transaction
func (bb *BadgerBatch) Commit() error {
return bb.txn.Commit()
}
// Discard cancels a transaction
func (bb *BadgerBatch) Discard() {
bb.txn.Discard()
}
var _ Iterator = &BadgerIterator{}
// PrefixIterator returns instance of prefix Iterator for BadgerKV.
func (b *BadgerKV) PrefixIterator(prefix []byte) Iterator {
txn := b.db.NewTransaction(false)
iter := txn.NewIterator(badger.DefaultIteratorOptions)
iter.Seek(prefix)
return &BadgerIterator{
txn: txn,
iter: iter,
prefix: prefix,
lastError: nil,
}
}
// BadgerIterator encapsulates prefix iterator for badger kv store.
type BadgerIterator struct {
txn *badger.Txn
iter *badger.Iterator
prefix []byte
lastError error
}
// Valid returns true if iterator is inside its prefix, false otherwise.
func (i *BadgerIterator) Valid() bool {
return i.iter.ValidForPrefix(i.prefix)
}
// Next progresses iterator to the next key-value pair.
func (i *BadgerIterator) Next() {
i.iter.Next()
}
// Key returns key pointed by iterator.
func (i *BadgerIterator) Key() []byte {
return i.iter.Item().KeyCopy(nil)
}
// Value returns value pointer by iterator.
func (i *BadgerIterator) Value() []byte {
val, err := i.iter.Item().ValueCopy(nil)
if err != nil {
i.lastError = err
}
return val
}
// Error returns last error that occurred during iteration.
func (i *BadgerIterator) Error() error {
return i.lastError
}
// Discard has to be called to free iterator resources.
func (i *BadgerIterator) Discard() {
i.iter.Close()
i.txn.Discard()
}