steem-block-db/db.go

143 lines
2.9 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"
// SteemDataStore is the object with which the rest of this tool interacts
type SteemDataStore struct {
kv KVStorer
}
func NewSteemDataStore(dir string) *SteemDataStore {
self := new(SteemDataStore)
self.kv = NewRedisKVStore()
return self
}
func (self *SteemDataStore) StoreBlockOps(blockNum int, blockOps []byte) {
panic("not implemented")
}
// 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)
Put(*string, *string) error
Close()
}
type RedisKVStore struct {
rc *redis.Client
}
func NewRedisKVStore() *RedisKVStore {
rkvs := new(RedisKVStore)
rkvs.Open("localhost:6379") //FIXME(sneak) use viper
return rkvs
}
func (self *RedisKVStore) Get(keyname *string) (*string, error) {
val, err := self.rc.Get(*keyname).Result()
if err != nil {
panic("unable to fetch key from redis")
//FIXME(sneak) we should probably distinguish between key not
//existing and fetch error
}
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(nil)
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
}