fiatjaf
113 lines
2.1 KiB
Go
113 lines
2.1 KiB
Go
package lmdb
|
|
|
|
import (
|
|
"os"
|
|
|
|
"github.com/PowerDNS/lmdb-go/lmdb"
|
|
"github.com/nbd-wtf/go-nostr/sdk/kvstore"
|
|
)
|
|
|
|
var _ kvstore.KVStore = (*Store)(nil)
|
|
|
|
type Store struct {
|
|
env *lmdb.Env
|
|
dbi lmdb.DBI
|
|
}
|
|
|
|
func NewStore(path string) (*Store, error) {
|
|
// create directory if it doesn't exist
|
|
if err := os.MkdirAll(path, 0755); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// initialize environment
|
|
env, err := lmdb.NewEnv()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// set max DBs and map size
|
|
env.SetMaxDBs(1)
|
|
env.SetMapSize(1 << 30) // 1GB
|
|
|
|
// open the environment
|
|
if err := env.Open(path, lmdb.NoTLS|lmdb.WriteMap, 0644); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
store := &Store{env: env}
|
|
|
|
// open the database
|
|
if err := env.Update(func(txn *lmdb.Txn) error {
|
|
dbi, err := txn.OpenDBI("store", lmdb.Create)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
store.dbi = dbi
|
|
return nil
|
|
}); err != nil {
|
|
env.Close()
|
|
return nil, err
|
|
}
|
|
|
|
return store, nil
|
|
}
|
|
|
|
func (s *Store) Get(key []byte) ([]byte, error) {
|
|
var value []byte
|
|
err := s.env.View(func(txn *lmdb.Txn) error {
|
|
v, err := txn.Get(s.dbi, key)
|
|
if lmdb.IsNotFound(err) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// make a copy since v is only valid during the transaction
|
|
value = make([]byte, len(v))
|
|
copy(value, v)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func (s *Store) Set(key []byte, value []byte) error {
|
|
return s.env.Update(func(txn *lmdb.Txn) error {
|
|
return txn.Put(s.dbi, key, value, 0)
|
|
})
|
|
}
|
|
|
|
func (s *Store) Delete(key []byte) error {
|
|
return s.env.Update(func(txn *lmdb.Txn) error {
|
|
return txn.Del(s.dbi, key, nil)
|
|
})
|
|
}
|
|
|
|
func (s *Store) Close() error {
|
|
s.env.Close()
|
|
return nil
|
|
}
|
|
|
|
func (s *Store) Scan(prefix []byte, fn func(key []byte, value []byte) bool) error {
|
|
return s.env.View(func(txn *lmdb.Txn) error {
|
|
cursor, err := txn.OpenCursor(s.dbi)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer cursor.Close()
|
|
|
|
for k, v, err := cursor.Get(prefix, nil, lmdb.SetRange); err == nil; k, v, err = cursor.Get(nil, nil, lmdb.Next) {
|
|
if !bytes.HasPrefix(k, prefix) {
|
|
break
|
|
}
|
|
if !fn(k, v) {
|
|
break
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
}
|