steem-block-db/kvstore.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
}