qbolt/bolt.go

package main

import (
	"encoding/json"
	"errors"
	"os"
	"time"

	bolt "go.etcd.io/bbolt"
)

func Bolt_Open(readOnly bool, path string) (*bolt.DB, error) {
	opts := *bolt.DefaultOptions
	opts.Timeout = 10 * time.Second
	opts.ReadOnly = readOnly

	return bolt.Open(path, os.FileMode(0644), &opts)
}

func walkBuckets(tx *bolt.Tx, browse []string) (*bolt.Bucket, error) {
	bucket := tx.Bucket([]byte(browse[0]))
	if bucket == nil {
		return nil, errors.New("Unknown bucket")
	}

	for i := 1; i < len(browse); i += 1 {
		bucket = bucket.Bucket([]byte(browse[i]))
		if bucket == nil {
			return nil, errors.New("Unknown bucket")
		}
	}

	return bucket, nil
}

func withBrowse_ReadOnly(db *bolt.DB, browse []string, fn func(tx *bolt.Tx, bucket *bolt.Bucket) error) error {
	if len(browse) == 0 {
		// not a bucket
		return errors.New("No bucket selected")
	}

	return db.View(func(tx *bolt.Tx) error {

		bucket, err := walkBuckets(tx, browse)
		if err != nil {
			return err
		}

		// Walked the bucket chain, now run the user callback
		return fn(tx, bucket)

	})
}

func Bolt_CreateBucket(db *bolt.DB, browse []string, newBucket string) error {

	return db.Update(func(tx *bolt.Tx) error {

		if len(browse) == 0 {
			// Top-level bucket
			_, err := tx.CreateBucket([]byte(newBucket))
			return err

		} else {
			// Deeper bucket
			bucket, err := walkBuckets(tx, browse)
			if err != nil {
				return err
			}

			// Walked the bucket chain, now create the new bucket
			_, err = bucket.CreateBucket([]byte(newBucket))
			return err
		}
	})
}

func Bolt_DeleteBucket(db *bolt.DB, browse []string, delBucket string) error {
	return db.Update(func(tx *bolt.Tx) error {

		if len(browse) == 0 {
			// Top-level bucket
			return tx.DeleteBucket([]byte(delBucket))

		} else {
			// Deeper bucket
			bucket, err := walkBuckets(tx, browse)
			if err != nil {
				return err
			}

			// Walked the bucket chain, now delete the selected bucket
			return bucket.DeleteBucket([]byte(delBucket))
		}
	})
}

func Bolt_SetItem(db *bolt.DB, browse []string, key, val string) error {
	if len(browse) == 0 {
		return errors.New("Can't create top-level items")
	}

	return db.Update(func(tx *bolt.Tx) error {

		bucket, err := walkBuckets(tx, browse)
		if err != nil {
			return err
		}

		return bucket.Put([]byte(key), []byte(val))
	})
}

func Bolt_DeleteItem(db *bolt.DB, browse []string, key string) error {
	if len(browse) == 0 {
		return errors.New("Can't create top-level items")
	}

	return db.Update(func(tx *bolt.Tx) error {

		bucket, err := walkBuckets(tx, browse)
		if err != nil {
			return err
		}

		return bucket.Delete([]byte(key))
	})
}

func Bolt_DBStats(db *bolt.DB) (string, error) {
	jBytes, err := json.MarshalIndent(db.Stats(), "", "  ")
	if err != nil {
		return "", err
	}

	return string(jBytes), nil
}

func Bolt_BucketStats(db *bolt.DB, browse []string) (string, error) {
	var stats bolt.BucketStats

	err := withBrowse_ReadOnly(db, browse, func(tx *bolt.Tx, bucket *bolt.Bucket) error {
		stats = bucket.Stats()
		return nil
	})
	if err != nil {
		return "", err
	}

	jBytes, err := json.MarshalIndent(stats, "", "  ")
	if err != nil {
		return "", err
	}

	return string(jBytes), err
}

func Bolt_ListBuckets(db *bolt.DB, browse []string, cb func(b string)) error {

	if len(browse) == 0 {
		// root mode
		return db.View(func(tx *bolt.Tx) error {
			return tx.ForEach(func(k []byte, _ *bolt.Bucket) error {
				cb(string(k))
				return nil
			})
		})

	}

	// Nested-mode
	return withBrowse_ReadOnly(db, browse, func(tx *bolt.Tx, bucket *bolt.Bucket) error {
		return bucket.ForEach(func(k, v []byte) error {
			// non-nil v means it's a data item
			if v == nil {
				cb(string(k))
			}
			return nil
		})
	})
}

type ListItemInfo struct {
	Name    string
	DataLen int64
}

func Bolt_ListItems(db *bolt.DB, browse []string, cb func(ListItemInfo) error) error {

	if len(browse) == 0 {
		return errors.New("No bucket specified")
	}

	// Nested-mode
	return withBrowse_ReadOnly(db, browse, func(tx *bolt.Tx, bucket *bolt.Bucket) error {
		return bucket.ForEach(func(k, v []byte) error {
			if v == nil {
				return nil // nil v means it's a bucket, skip
			}

			return cb(ListItemInfo{string(k), int64(len(v))})
		})
	})
}

func Bolt_GetItem(db *bolt.DB, browse []string, key string) (string, error) {
	var ret string

	err := withBrowse_ReadOnly(db, browse, func(tx *bolt.Tx, bucket *bolt.Bucket) error {
		d := bucket.Get([]byte(key))
		ret = string(d)
		return nil
	})
	return ret, err
}

func Bolt_Close(db *bolt.DB) error {
	return db.Close()
}