boltease/bucket_funcs.go

package boltease

import (
	"encoding/binary"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

func DoUpdate(b *DB, path []string, fn func(bkt *bolt.Bucket) error) error {
	var err error
	if !b.dbIsOpen {
		if err = b.OpenDB(); err != nil {
			return err
		}
		defer b.CloseDB()
	}

	err = b.MkBucketPath(path)
	if err != nil {
		return err
	}
	err = b.boltDB.Update(func(tx *bolt.Tx) error {
		bkt := tx.Bucket([]byte(path[0]))
		if bkt == nil {
			return fmt.Errorf("couldn't find bucket " + path[0])
		}
		for idx := 1; idx < len(path); idx++ {
			bkt, err = bkt.CreateBucketIfNotExists([]byte(path[idx]))
			if err != nil {
				return err
			}
		}
		// bkt should have the last bucket in the path
		return fn(bkt)
	})
	return err
}

func DoRead(b *DB, path []string, fn func(bkt *bolt.Bucket) error) error {
	var err error
	if !b.dbIsOpen {
		if err = b.OpenDB(); err != nil {
			return err
		}
		defer b.CloseDB()
	}
	err = b.boltDB.View(func(tx *bolt.Tx) error {
		bkt := tx.Bucket([]byte(path[0]))
		if bkt == nil {
			return fmt.Errorf("couldn't find bucket " + path[0])
		}
		for idx := 1; idx < len(path); idx++ {
			bkt = bkt.Bucket([]byte(path[idx]))
			if bkt == nil {
				return fmt.Errorf("couldn't find bucket " + strings.Join(path[:idx+1], "/"))
			}
		}
		// bkt should have the last bucket in the path
		return fn(bkt)
	})
	return err
}

// These bucket functions are meant to be run inside of a boltdb
// transaction (Update/View)
// GetBucket must be run inside of an Update transaction
// Otherwise, any Get* function can be in either.
// And any Set* function needs to be in an Update.
func GetBucket(bkt *bolt.Bucket, path ...string) (*bolt.Bucket, error) {
	if bkt == nil {
		return nil, fmt.Errorf("bucket must not be nil")
	}
	var ret *bolt.Bucket
	var err error
	ret, err = bkt.CreateBucketIfNotExists([]byte(path[0]))
	if len(path) > 1 {
		for idx := 1; idx < len(path); idx++ {
			ret, err = bkt.CreateBucketIfNotExists([]byte(path[idx]))
			if err != nil {
				return nil, err
			}
		}
	}
	return ret, err
}

func GetString(bkt *bolt.Bucket, key string) (string, error) {
	if bkt == nil {
		return "", fmt.Errorf("bucket must not be nil")
	}
	return string(bkt.Get([]byte(key))), nil
}

func SetString(bkt *bolt.Bucket, key, val string) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	return bkt.Put([]byte(key), []byte(val))
}

func GetInt(bkt *bolt.Bucket, key string) (int, error) {
	if bkt == nil {
		return 0, fmt.Errorf("bucket must not be nil")
	}
	r, err := GetString(bkt, key)
	if err != nil {
		return 0, err
	}
	return strconv.Atoi(r)
}

func SetInt(bkt *bolt.Bucket, key string, val int) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	return SetString(bkt, key, strconv.Itoa(val))
}

func GetBool(bkt *bolt.Bucket, key string) (bool, error) {
	r, err := GetString(bkt, key)
	if err != nil {
		return false, err
	}
	if r == "true" || r == "1" {
		return true, nil
	} else if r == "false" || r == "0" {
		return false, nil
	} else {
		return false, fmt.Errorf("cannot parse as a boolean")
	}
}

func SetBool(bkt *bolt.Bucket, key string, val bool) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	if val {
		return SetString(bkt, key, "true")
	}
	return SetString(bkt, key, "false")
}

func GetTimestamp(bkt *bolt.Bucket, key string) (time.Time, error) {
	r, err := GetString(bkt, key)
	if err != nil {
		return time.Unix(0, 0), err
	}
	return time.Parse(time.RFC3339, r)
}

func SetTimestamp(bkt *bolt.Bucket, key string, val time.Time) error {
	return SetString(bkt, key, val.Format(time.RFC3339))
}

func GetBucketList(bkt *bolt.Bucket) ([]string, error) {
	var ret []string
	if bkt == nil {
		return ret, fmt.Errorf("bucket must not be nil")
	}
	bkt.ForEach(func(k, v []byte) error {
		if v == nil {
			// Must be a bucket
			ret = append(ret, string(k))
		}
		return nil
	})
	return ret, nil
}

func GetKeyList(bkt *bolt.Bucket) ([]string, error) {
	var ret []string
	if bkt == nil {
		return ret, fmt.Errorf("bucket must not be nil")
	}
	bkt.ForEach(func(k, v []byte) error {
		if v != nil {
			ret = append(ret, string(k))
		}
		return nil
	})
	return ret, nil
}

func DeletePair(bkt *bolt.Bucket, key string) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	// Make sure that the key belongs to a pair
	if tst := bkt.Bucket([]byte(key)); tst != nil {
		return fmt.Errorf("key is a bucket")
	}
	return bkt.Delete([]byte(key))
}

func DeleteBucket(bkt *bolt.Bucket, key string) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	// Make sure that the key belongs to a pair
	if tst := bkt.Bucket([]byte(key)); tst == nil {
		return fmt.Errorf("key is not a bucket")
	}
	return bkt.DeleteBucket([]byte(key))
}

func GetStringList(bkt *bolt.Bucket) ([]string, error) {
	var ret []string
	if bkt == nil {
		return ret, fmt.Errorf("bucket must not be nil")
	}
	bkt.ForEach(func(k, v []byte) error {
		if v != nil {
			// Must be a key
			ret = append(ret, string(v))
		}
		return nil
	})
	return ret, nil
}

func SetStringList(bkt *bolt.Bucket, val []string) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	genseq := func(id uint64) []byte {
		bId := make([]byte, 8)
		binary.BigEndian.PutUint64(bId, id)
		return bId
	}
	var err error
	// We need to get the sequence number, if it's greater than the length of
	// val, we're going to have to delete values
	seq := bkt.Sequence()
	bkt.SetSequence(0)
	for _, v := range val {
		id, _ := bkt.NextSequence()
		if err = bkt.Put(genseq(id), []byte(v)); err != nil {
			return err
		}
	}
	currSeq := bkt.Sequence()
	for ; seq < currSeq; seq++ {
		err = bkt.Delete(genseq(seq))
		if err != nil {
			return err
		}
	}
	return nil
}

func AddToStringList(bkt *bolt.Bucket, key string, values ...string) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	genseq := func(id uint64) []byte {
		bId := make([]byte, 8)
		binary.BigEndian.PutUint64(bId, id)
		return bId
	}
	var err error
	for _, v := range values {
		id, _ := bkt.NextSequence()
		if err = bkt.Put(genseq(id), []byte(v)); err != nil {
			return err
		}
	}
	return nil
}

func GetKeyValueMap(bkt *bolt.Bucket) (map[string][]byte, error) {
	ret := make(map[string][]byte)
	if bkt == nil {
		return ret, fmt.Errorf("bucket must not be nil")
	}
	bkt.ForEach(func(k, v []byte) error {
		if v != nil {
			ret[string(k)] = v
		}
		return nil
	})
	return ret, nil
}

func SetKeyValueMap(bkt *bolt.Bucket, kvMap map[string][]byte) error {
	if bkt == nil {
		return fmt.Errorf("bucket must not be nil")
	}
	var err error
	for k, v := range kvMap {
		err = bkt.Put([]byte(k), v)
		if err != nil {
			return err
		}
	}
	return nil
}

func GetKeyValueStringMap(bkt *bolt.Bucket) (map[string]string, error) {
	ret := make(map[string]string)
	if bkt == nil {
		return ret, fmt.Errorf("bucket must not be nil")
	}
	btMap, err := GetKeyValueMap(bkt)
	if err != nil {
		return ret, err
	}
	for k, v := range btMap {
		ret[k] = string(v)
	}
	return ret, nil
}

func SetKeyValueStringMap(bkt *bolt.Bucket, kvMap map[string]string) error {
	chain := make(map[string][]byte)
	for k, v := range kvMap {
		chain[k] = []byte(v)
	}
	return SetKeyValueMap(bkt, chain)
}