switch do diskv

This commit is contained in:
Aaron Bieber 2020-02-04 17:22:19 -07:00
parent 94c9645820
commit 0c6fba31d5
5 changed files with 38 additions and 86 deletions

1
.gitignore vendored
View File

@ -1,4 +1,5 @@
.todo
db/*
*.db
mcchunkie
deploy.sh

3
go.mod
View File

@ -3,8 +3,9 @@ module git.sr.ht/~qbit/mcchunkie
go 1.13
require (
github.com/google/btree v1.0.0 // indirect
github.com/matrix-org/gomatrix v0.0.0-20200128155335-9e7906b6766d
go.etcd.io/bbolt v1.3.3
github.com/peterbourgon/diskv v2.0.1+incompatible
golang.org/x/crypto v0.0.0-20200128174031-69ecbb4d6d5d
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3
golang.org/x/sys v0.0.0-20200124204421-9fbb57f87de9

4
go.sum
View File

@ -1,5 +1,9 @@
github.com/google/btree v1.0.0 h1:0udJVsspx3VBr5FwtLhQQtuAsVc79tTq0ocGIPAU6qo=
github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
github.com/matrix-org/gomatrix v0.0.0-20200128155335-9e7906b6766d h1:Vf/EQgAfg8/CBUQv9te7UJreZ9iKKouB2gb8UIRM4jQ=
github.com/matrix-org/gomatrix v0.0.0-20200128155335-9e7906b6766d/go.mod h1:3fxX6gUjWyI/2Bt7J1OLhpCzOfO/bB3AiX0cJtEKud0=
github.com/peterbourgon/diskv v2.0.1+incompatible h1:UBdAOUP5p4RWqPBg048CAvpKN+vxiaj6gdUUzhl4XmI=
github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=
go.etcd.io/bbolt v1.3.3 h1:MUGmc65QhB3pIlaQ5bB4LwqSj6GIonVJXpZiaKNyaKk=
go.etcd.io/bbolt v1.3.3/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

18
main.go
View File

@ -40,7 +40,7 @@ func main() {
flag.StringVar(&server, "server", "", "matrix server")
flag.StringVar(&avatar, "avatar", "", "set the avatar of the bot to specified url")
flag.BoolVar(&setup, "s", false, "setup account")
flag.StringVar(&db, "db", "mcchunkie.db", "full path to database file")
flag.StringVar(&db, "db", "db", "full path to database directory")
flag.Parse()
@ -56,13 +56,13 @@ func main() {
}
if server == "" {
server, err = store.get("config", "server")
server, err = store.get("server")
if server == "" {
log.Fatalln("please specify a server")
}
} else {
store.set("config", "server", server)
store.set("server", server)
}
log.Printf("connecting to %s\n", server)
@ -93,16 +93,16 @@ func main() {
// No longer need tty now that we have our info
pledge("stdio unveil rpath wpath cpath flock dns inet")
store.set("account", "username", username)
store.set("account", "access_token", resp.AccessToken)
store.set("account", "user_id", resp.UserID)
store.set("username", username)
store.set("access_token", resp.AccessToken)
store.set("user_id", resp.UserID)
accessToken = resp.AccessToken
userID = resp.UserID
} else {
username, _ = store.get("account", "username")
accessToken, _ = store.get("account", "access_token")
userID, _ = store.get("account", "user_id")
username, _ = store.get("username")
accessToken, _ = store.get("access_token")
userID, _ = store.get("user_id")
}
cli.SetCredentials(userID, accessToken)

View File

@ -3,93 +3,39 @@ package main
import (
"bytes"
"encoding/gob"
"log"
"github.com/matrix-org/gomatrix"
bolt "go.etcd.io/bbolt"
"github.com/peterbourgon/diskv"
)
// MCStore implements a gomatrix.Storer and exposes a bbolt db to be used for
// MCStore implements a gomatrix.Storer and exposes a diskv db to be used for
// application storage (account info, config info etc).
type MCStore struct {
db *bolt.DB
db *diskv.Diskv
}
// NewStore creates a new MCStore instance populated with the proper buckets.
func NewStore(path string) (*MCStore, error) {
db, err := bolt.Open(path, 0666, nil)
if err != nil {
return nil, err
}
flatTransform := func(s string) []string { return []string{} }
db := diskv.New(diskv.Options{
BasePath: "db",
Transform: flatTransform,
CacheSizeMax: 1024 * 1024,
})
s := &MCStore{db: db}
err = s.db.Update(func(tx *bolt.Tx) error {
buckets := []string{"filter", "batch", "room", "account", "config", "errata"}
for _, b := range buckets {
if _, err := tx.CreateBucketIfNotExists([]byte(b)); err != nil {
return err
}
}
return nil
})
if err != nil {
s.db.Close()
return nil, err
}
return s, nil
}
func (s *MCStore) set(bucket, key, value string) error {
err := s.db.Update(func(tx *bolt.Tx) error {
bkt, err := tx.CreateBucketIfNotExists([]byte(bucket))
if err != nil {
return err
}
err = bkt.Put([]byte(key), []byte(value))
if err != nil {
return err
}
return nil
})
if err != nil {
log.Fatal(err)
}
return nil
func (s *MCStore) set(key string, value string) {
v := []byte(value)
s.db.Write(key, v)
}
func (s *MCStore) get(bucket, key string) (string, error) {
var result string
return result, s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket([]byte(bucket))
data := bkt.Get([]byte(key))
if data == nil {
return nil
}
result = string(data)
return nil
})
}
func (s *MCStore) getBytes(bucket, key string) ([]byte, error) {
var result []byte
return result, s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket([]byte(bucket))
data := bkt.Get([]byte(key))
if data == nil {
return nil
}
result = data
return nil
})
func (s *MCStore) get(key string) (string, error) {
b, err := s.db.Read(key)
return string(b), err
}
func (s *MCStore) encodeRoom(room *gomatrix.Room) ([]byte, error) {
@ -115,36 +61,36 @@ func (s *MCStore) decodeRoom(room []byte) (*gomatrix.Room, error) {
// SaveFilterID exposed for gomatrix
func (s *MCStore) SaveFilterID(userID, filterID string) {
_ = s.set("filter", userID, filterID)
s.set(userID, filterID)
}
// LoadFilterID exposed for gomatrix
func (s *MCStore) LoadFilterID(userID string) string {
filter, _ := s.get("filter", userID)
filter, _ := s.get(userID)
return string(filter)
}
// SaveNextBatch exposed for gomatrix
func (s *MCStore) SaveNextBatch(userID, nextBatchToken string) {
_ = s.set("batch", userID, nextBatchToken)
s.set(userID, nextBatchToken)
}
// LoadNextBatch exposed for gomatrix
func (s *MCStore) LoadNextBatch(userID string) string {
batch, _ := s.get("batch", userID)
batch, _ := s.get(userID)
return string(batch)
}
// SaveRoom exposed for gomatrix
func (s *MCStore) SaveRoom(room *gomatrix.Room) {
b, _ := s.encodeRoom(room)
_ = s.set("room", room.ID, string(b))
s.set(room.ID, string(b))
}
// LoadRoom exposed for gomatrix
func (s *MCStore) LoadRoom(roomID string) *gomatrix.Room {
b, _ := s.getBytes("room", roomID)
room, _ := s.decodeRoom(b)
b, _ := s.get(roomID)
room, _ := s.decodeRoom([]byte(b))
return room
}