ezshare/store/bolt.go
2021-12-04 11:30:42 +01:00

114 lines
2.4 KiB
Go

package store
import (
"fmt"
"gitea.benny.dog/torjus/ezshare/pb"
"github.com/google/uuid"
bolt "go.etcd.io/bbolt"
"google.golang.org/protobuf/proto"
)
var _ FileStore = &BoltStore{}
type BoltStore struct {
db *bolt.DB
}
var bktKey = []byte("files")
func NewBoltStore(path string) (*BoltStore, error) {
s := &BoltStore{}
db, err := bolt.Open(path, 0660, &bolt.Options{})
if err != nil {
return nil, fmt.Errorf("unable to open store: %w", err)
}
s.db = db
err = db.Update(func(t *bolt.Tx) error {
if _, err := t.CreateBucketIfNotExists(bktKey); err != nil {
return err
}
return nil
})
if err != nil {
return nil, fmt.Errorf("unable to create bucket: %w", err)
}
return s, nil
}
func (s *BoltStore) Close() error {
return s.db.Close()
}
func (s *BoltStore) GetFile(id string) (*pb.File, error) {
var file pb.File
var data []byte
if err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
data = bkt.Get([]byte(id))
return nil
}); err != nil {
return nil, fmt.Errorf("error getting file: %w", err)
}
if data == nil {
return nil, ErrNoSuchFile
}
err := proto.Unmarshal(data, &file)
if err != nil {
return nil, fmt.Errorf("unable to unmarshal file: %w", err)
}
return &file, nil
}
func (s *BoltStore) StoreFile(file *pb.File) (string, error) {
file.FileId = uuid.Must(uuid.NewRandom()).String()
data, err := proto.Marshal(file)
if err != nil {
return "", fmt.Errorf("unable to marshal file: %w", err)
}
err = s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
return bkt.Put([]byte(file.FileId), data)
})
if err != nil {
return "", fmt.Errorf("unable to store file: %w", err)
}
return file.FileId, nil
}
func (s *BoltStore) DeleteFile(id string) error {
return s.db.Update(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
data := bkt.Get([]byte(id))
if data == nil {
return ErrNoSuchFile
}
return bkt.Delete([]byte(id))
})
}
func (s *BoltStore) ListFiles() ([]*pb.ListFilesResponse_ListFileInfo, error) {
var response []*pb.ListFilesResponse_ListFileInfo
err := s.db.View(func(t *bolt.Tx) error {
bkt := t.Bucket(bktKey)
return bkt.ForEach(func(k, v []byte) error {
var f pb.File
err := proto.Unmarshal(v, &f)
if err != nil {
return err
}
response = append(response, &pb.ListFilesResponse_ListFileInfo{FileId: f.FileId, Metadata: f.Metadata})
return nil
})
})
if err != nil {
return nil, err
}
return response, nil
}