kleinTodo/common/bolt.go

package common

import (
	"fmt"
	bolt "go.etcd.io/bbolt"
	"sync"
)

type DataStore interface {
	SaveValueToBucket(bucket, key, value string) error
	CreateBucket(bucket string) error
	GetFromBucketByKey(bucket, key string) (string, error)
	GetAllBuckets() ([]string, error)
	GetAllFromBucket(bucket string) (map[string]string, error)
	GetAllKeysFromBucket(bucket string) ([]string, error)
	EmptyBucket(bucket string) error
	Close() error
}

type BoltStore struct {
	DB *bolt.DB
}

func NewBoltStore(path string) (*BoltStore, error) {
	db, err := bolt.Open(path, 0600, nil)
	if err != nil {
		return nil, fmt.Errorf("could not open db: %w", err)
	}
	return &BoltStore{DB: db}, nil
}

func (s *BoltStore) Close() error {
	return s.DB.Close()
}

var (
	dataStore *BoltStore
	once      sync.Once
	err       error
)

func GetTodoDataStore() (*BoltStore, error) {
	once.Do(func() {
		// We assign to the outer 'dataStore' and 'err' variables.
		dataStore, err = NewBoltStore("data/todo.db")
	})
	if err != nil {
		return nil, err
	}

	return dataStore, nil
}

// SaveValueToBucket Save data to bucket
func (s *BoltStore) SaveValueToBucket(bucket, key, value string) error {
	return s.DB.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(bucket))
		if err != nil {
			return err
		}
		return b.Put([]byte(key), []byte(value))
	})
}

// RemoveValueFromBucket remove value
func (s *BoltStore) RemoveValueFromBucket(bucket, key string) error {
	return s.DB.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(bucket))
		if err != nil {
			return err
		}
		return b.Delete([]byte(key))
	})
}

// CreateBucket Create bucket if not exists
func (s *BoltStore) CreateBucket(bucket string) error {
	return s.DB.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucket))
		return err
	})
}

// GetFromBucketByKey Returns value from bucket by key or empty string
func (s *BoltStore) GetFromBucketByKey(bucket, key string) string {
	var value string
	s.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if b == nil {
			return nil // Bucket doesn't exist, nothing to do
		}
		// Get returns nil if the key doesn't exist
		valBytes := b.Get([]byte(key))
		if valBytes != nil {
			value = string(valBytes)
		}
		return nil
	})
	return value
}

// GetAllBuckets Returns map of all buckets
func (s *BoltStore) GetAllBuckets() []string {
	var a = make([]string, 0)
	s.DB.View(func(tx *bolt.Tx) error {
		return tx.ForEach(func(name []byte, b *bolt.Bucket) error {
			a = append(a, string(name))
			return nil
		})
	})
	return a
}

// GetAllFromBucket Returns map of all from bucket
func (s *BoltStore) GetAllFromBucket(bucket string) map[string]string {
	results := make(map[string]string)
	s.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if b == nil {
			return nil
		}
		// Using a cursor is the recommended way to iterate
		c := b.Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			results[string(k)] = string(v)
		}
		return nil
	})
	return results
}

// GetAllKeysFromBucket Returns keys from bucket
func (s *BoltStore) GetAllKeysFromBucket(bucket string) []string {
	var keys []string
	s.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if b == nil {
			return nil
		}
		return b.ForEach(func(k, v []byte) error {
			keys = append(keys, string(k))
			return nil
		})
	})
	return keys
}

// EmptyBucket Returns value from bucket by key or empty string
func (s *BoltStore) EmptyBucket(bucket string) error {
	return s.DB.Update(func(tx *bolt.Tx) error {
		if err := tx.DeleteBucket([]byte(bucket)); err != nil {
			return err
		}
		if _, err := tx.CreateBucket([]byte(bucket)); err != nil {
			return err
		}
		return nil
	})
}