91 lines
1.9 KiB
Go
91 lines
1.9 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"
|
||
|
|
)
|
||
|
|
|
||
|
|
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))
|
||
|
|
})
|
||
|
|
}
|