139 lines
2.7 KiB
Go
139 lines
2.7 KiB
Go
package main
|
|
|
|
import log "github.com/sirupsen/logrus"
|
|
import "io/ioutil"
|
|
import "github.com/dgraph-io/badger"
|
|
import "github.com/spf13/afero"
|
|
import "github.com/go-redis/redis"
|
|
|
|
// KeyValueStorer is an interface for the backend kv store used by
|
|
// SteemDataStore
|
|
// it could be fs, badgerdb, whatever
|
|
|
|
type KVStorer interface {
|
|
Open(string)
|
|
Get(*string) (*string, error)
|
|
Exists(*string) (bool, error)
|
|
Put(*string, *string) error
|
|
Close()
|
|
}
|
|
|
|
type RedisKVStore struct {
|
|
rc *redis.Client
|
|
}
|
|
|
|
func NewRedisKVStore(hostname string) *RedisKVStore {
|
|
rkvs := new(RedisKVStore)
|
|
rkvs.Open(hostname) //FIXME(sneak) use viper
|
|
return rkvs
|
|
}
|
|
|
|
func (self *RedisKVStore) Exists(keyname *string) (bool, error) {
|
|
val, err := self.rc.Exists(*keyname).Result()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if val == int64(1) {
|
|
return true, nil
|
|
} else {
|
|
return false, nil
|
|
}
|
|
}
|
|
|
|
func (self *RedisKVStore) Get(keyname *string) (*string, error) {
|
|
val, err := self.rc.Get(*keyname).Result()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &val, nil
|
|
}
|
|
|
|
func (self *RedisKVStore) Put(keyname *string, value *string) error {
|
|
err := self.rc.Set(*keyname, *value, 0).Err()
|
|
if err != nil {
|
|
panic("unable to write to redis")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (self *RedisKVStore) Close() {
|
|
self.rc.Close()
|
|
}
|
|
|
|
func (self *RedisKVStore) Open(hostname string) {
|
|
self.rc = redis.NewClient(&redis.Options{
|
|
Addr: hostname,
|
|
Password: "", // no password set
|
|
DB: 0, // use default DB
|
|
})
|
|
|
|
_, err := self.rc.Ping().Result()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
type AferoFSKVStore struct {
|
|
fs *afero.Fs
|
|
}
|
|
|
|
// BadgerKVStore is an object that conforms to KeyValueStorer for use
|
|
// by SteemDataStore to persist Steem data
|
|
|
|
type BadgerKVStore struct {
|
|
db *badger.DB
|
|
}
|
|
|
|
func NewBadgerKVStore(dir string) *BadgerKVStore {
|
|
kv := new(BadgerKVStore)
|
|
kv.Open(dir)
|
|
return kv
|
|
}
|
|
|
|
func (kv *BadgerKVStore) Open(dir string) {
|
|
dir, err := ioutil.TempDir("", "badger")
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
opts := badger.DefaultOptions
|
|
opts.Dir = dir
|
|
opts.ValueDir = dir
|
|
kv.db, err = badger.Open(opts)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func (kv *BadgerKVStore) Close() {
|
|
kv.db.Close()
|
|
}
|
|
|
|
func (kv *BadgerKVStore) Put(key *string, value *string) error {
|
|
txn := kv.db.NewTransaction(true) // Read-write txn
|
|
err := txn.Set([]byte(*key), []byte(*value))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
err = txn.Commit()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (kv *BadgerKVStore) Get(key *string) (*string, error) {
|
|
txn := kv.db.NewTransaction(false)
|
|
item, err := txn.Get([]byte(*key))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
val, err := item.ValueCopy(nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s := string(val)
|
|
return &s, nil
|
|
}
|