steem-block-db/db.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"
import "fmt"
import "strconv"

const appPrefix = "steem-block-fetcher"

type SteemDataStorer interface {
	SetCurrentBlockHeight(BlockNumber) error
	CurrentBlockHeight() BlockNumber
	HaveOpsForBlock(BlockNumber) bool
	StoreBlockOps(BlockNumber, *[]byte) error
}

// 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) SetCurrentBlockHeight(blockNum BlockNumber) error {
	panic("unimplemented")
	return nil
}

func (self *SteemDataStore) StoreBlockOps(blockNum BlockNumber, blockOps *[]byte) error {
	panic("unimplemented")
}

func (self *SteemDataStore) HaveOpsForBlock(blockNum BlockNumber) bool {
	panic("unimplemented")
}

func (self *SteemDataStore) CurrentBlockHeight() BlockNumber {
	keyname := fmt.Sprintf("%s:global:CurrentBlockHeight", appPrefix)
	val, err := self.kv.Get(&keyname)
	//FIXME(sneak) make this init the key to 0 if not found
	if err != nil {
		panic("unable to query redis")
	}
	intval, err := strconv.ParseUint(*val, 10, 64)
	return BlockNumber(intval)
}

// 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 {
		log.Panicf("unable to fetch key '%s' from redis", *keyname)
		//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
}