You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

139 lines
2.7KB

  1. package main
  2. import log "github.com/sirupsen/logrus"
  3. import "io/ioutil"
  4. import "github.com/dgraph-io/badger"
  5. import "github.com/spf13/afero"
  6. import "github.com/go-redis/redis"
  7. // KeyValueStorer is an interface for the backend kv store used by
  8. // SteemDataStore
  9. // it could be fs, badgerdb, whatever
  10. type KVStorer interface {
  11. Open(string)
  12. Get(*string) (*string, error)
  13. Exists(*string) (bool, error)
  14. Put(*string, *string) error
  15. Close()
  16. }
  17. type RedisKVStore struct {
  18. rc *redis.Client
  19. }
  20. func NewRedisKVStore(hostname string) *RedisKVStore {
  21. rkvs := new(RedisKVStore)
  22. rkvs.Open(hostname) //FIXME(sneak) use viper
  23. return rkvs
  24. }
  25. func (self *RedisKVStore) Exists(keyname *string) (bool, error) {
  26. val, err := self.rc.Exists(*keyname).Result()
  27. if err != nil {
  28. return false, err
  29. }
  30. if val == int64(1) {
  31. return true, nil
  32. } else {
  33. return false, nil
  34. }
  35. }
  36. func (self *RedisKVStore) Get(keyname *string) (*string, error) {
  37. val, err := self.rc.Get(*keyname).Result()
  38. if err != nil {
  39. return nil, err
  40. }
  41. return &val, nil
  42. }
  43. func (self *RedisKVStore) Put(keyname *string, value *string) error {
  44. err := self.rc.Set(*keyname, *value, 0).Err()
  45. if err != nil {
  46. panic("unable to write to redis")
  47. }
  48. return nil
  49. }
  50. func (self *RedisKVStore) Close() {
  51. self.rc.Close()
  52. }
  53. func (self *RedisKVStore) Open(hostname string) {
  54. self.rc = redis.NewClient(&redis.Options{
  55. Addr: hostname,
  56. Password: "", // no password set
  57. DB: 0, // use default DB
  58. })
  59. _, err := self.rc.Ping().Result()
  60. if err != nil {
  61. panic(err)
  62. }
  63. }
  64. type AferoFSKVStore struct {
  65. fs *afero.Fs
  66. }
  67. // BadgerKVStore is an object that conforms to KeyValueStorer for use
  68. // by SteemDataStore to persist Steem data
  69. type BadgerKVStore struct {
  70. db *badger.DB
  71. }
  72. func NewBadgerKVStore(dir string) *BadgerKVStore {
  73. kv := new(BadgerKVStore)
  74. kv.Open(dir)
  75. return kv
  76. }
  77. func (kv *BadgerKVStore) Open(dir string) {
  78. dir, err := ioutil.TempDir("", "badger")
  79. if err != nil {
  80. log.Fatal(err)
  81. }
  82. opts := badger.DefaultOptions
  83. opts.Dir = dir
  84. opts.ValueDir = dir
  85. kv.db, err = badger.Open(opts)
  86. if err != nil {
  87. log.Fatal(err)
  88. }
  89. }
  90. func (kv *BadgerKVStore) Close() {
  91. kv.db.Close()
  92. }
  93. func (kv *BadgerKVStore) Put(key *string, value *string) error {
  94. txn := kv.db.NewTransaction(true) // Read-write txn
  95. err := txn.Set([]byte(*key), []byte(*value))
  96. if err != nil {
  97. log.Fatal(err)
  98. }
  99. err = txn.Commit()
  100. if err != nil {
  101. log.Fatal(err)
  102. }
  103. return nil
  104. }
  105. func (kv *BadgerKVStore) Get(key *string) (*string, error) {
  106. txn := kv.db.NewTransaction(false)
  107. item, err := txn.Get([]byte(*key))
  108. if err != nil {
  109. return nil, err
  110. }
  111. val, err := item.ValueCopy(nil)
  112. if err != nil {
  113. return nil, err
  114. }
  115. s := string(val)
  116. return &s, nil
  117. }