ezshare/store/bolt.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
}