kit/db/boltdb/boltdb.go

package boltdb

import (
	"bytes"
	"errors"
	"time"

	bolt "go.etcd.io/bbolt"
)

var (
	ErrBucketInvalid     = errors.New("bucket is invalid")
	ErrKeyInvalid        = errors.New("key is invalid")
	ErrValueInvalid      = errors.New("value is invalid")
	ErrBucketNotFound    = errors.New("bucket not found")
	ErrKeyNotFound       = errors.New("key not found")
	ErrPrefixInvalid     = errors.New("prefix is invalid")
	ErrRangeStartInvalid = errors.New("range start is invalid")
	ErrRangeEndInvalid   = errors.New("range end is invalid")
)

type boltDB struct {
	db *bolt.DB
}

func New(config ...Config) (db *boltDB) {
	cfg := configDefault(config...)
	db = &boltDB{}
	db.db, _ = bolt.Open(cfg.Name, 0600, &bolt.Options{
		Timeout:    time.Duration(cfg.Timeout) * time.Second,
		NoGrowSync: cfg.NoGrowSync,
		ReadOnly:   cfg.ReadOnly,
	})
	return
}

// 获取数据库地址
func (cc *boltDB) Path() (rv string) {
	return cc.db.Path()
}

// 关闭
func (cc *boltDB) Close() (rer error) {
	return cc.db.Close()
}

// 建表
func (cc *boltDB) CreateBucket(bucketName ...string) (rer error) {
	if bucketName == nil {
		return
	}
	rer = cc.db.Update(func(tx *bolt.Tx) error {
		var errs error
		for _, name := range bucketName {
			if name != "" {
				if _, err := tx.CreateBucketIfNotExists([]byte(name)); err != nil {
					errs = errors.Join(errs, err)
				}
			}
		}
		return errs
	})
	return
}

// 删表
func (cc *boltDB) DeleteBucket(bucketName ...string) (rer error) {
	if bucketName == nil {
		return
	}
	rer = cc.db.Update(func(tx *bolt.Tx) error {
		var errs error
		for _, name := range bucketName {
			if name != "" {
				if err := tx.DeleteBucket([]byte(name)); err != nil {
					errs = errors.Join(errs, err)
				}
			}
		}
		return errs
	})
	return
}

// 保存
func (cc *boltDB) Set(bucketName, key string, value []byte) (rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrKeyInvalid
	if key == "" {
		return
	}
	rer = ErrValueInvalid
	if value == nil {
		return
	}
	return cc.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		return b.Put([]byte(key), value)
	})
}

// 读取
func (cc *boltDB) Get(bucketName, key string) (rv []byte, rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrKeyInvalid
	if key == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		rv = b.Get([]byte(key))
		if rv == nil {
			return ErrKeyNotFound
		}
		return nil
	})
	return
}

// 存在
func (cc *boltDB) Has(bucketName, key string) (ok bool) {
	_, err := cc.Get(bucketName, key)
	return errors.Is(err, nil)
}

// 不存在
func (cc *boltDB) HasNot(bucketName, key string) (ok bool) {
	_, err := cc.Get(bucketName, key)
	return errors.Is(err, ErrKeyNotFound)
}

// 删除
func (cc *boltDB) Delete(bucketName, key string) (rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrKeyInvalid
	if key == "" {
		return
	}
	return cc.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		return b.Delete([]byte(key))
	})
}

// 自定义遍历
func (cc *boltDB) ForEach(bucketName string, fn func(k, v []byte) error) (rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		return b.ForEach(fn)
	})
	return
}

// 查询所有list
func (cc *boltDB) GetList(bucketName string) (rv [][]byte, rer error) {
	rer = cc.ForEach(bucketName, func(k, v []byte) error {
		rv = append(rv, v)
		return nil
	})
	return
}

// 查询所有map
func (cc *boltDB) GetMap(bucketName string) (rv map[string][]byte, rer error) {
	rv = make(map[string][]byte)
	rer = cc.ForEach(bucketName, func(k, v []byte) error {
		rv[string(k)] = v
		return nil
	})
	return
}

// 前缀扫描list
func (cc *boltDB) ScanPrefixList(bucketName, prefix string) (rv [][]byte, rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrPrefixInvalid
	if prefix == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(prefix)); k != nil && bytes.HasPrefix(k, []byte(prefix)); k, v = c.Next() {
			rv = append(rv, v)
		}
		return nil
	})
	return
}

// 前缀扫描map
func (cc *boltDB) ScanPrefixMap(bucketName, prefix string) (rv map[string][]byte, rer error) {
	rv = make(map[string][]byte)
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrPrefixInvalid
	if prefix == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(prefix)); k != nil && bytes.HasPrefix(k, []byte(prefix)); k, v = c.Next() {
			rv[string(k)] = v
		}
		return nil
	})
	return
}

// 后缀扫描list
func (cc *boltDB) ScanSuffixList(bucketName, prefix string) (rv [][]byte, rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrPrefixInvalid
	if prefix == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(prefix)); k != nil && bytes.HasSuffix(k, []byte(prefix)); k, v = c.Next() {
			rv = append(rv, v)
		}
		return nil
	})
	return
}

// 后缀扫描map
func (cc *boltDB) ScanSuffixMap(bucketName, prefix string) (rv map[string][]byte, rer error) {
	rv = make(map[string][]byte)
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrPrefixInvalid
	if prefix == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(prefix)); k != nil && bytes.HasSuffix(k, []byte(prefix)); k, v = c.Next() {
			rv[string(k)] = v
		}
		return nil
	})
	return
}

// 范围扫描list
func (cc *boltDB) ScanRangeList(bucketName, start, end string) (rv [][]byte, rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrRangeStartInvalid
	if start == "" {
		return
	}
	rer = ErrRangeEndInvalid
	if end == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(start)); k != nil && bytes.Compare(k, []byte(end)) <= 0; k, v = c.Next() {
			rv = append(rv, v)
		}
		return nil
	})
	return
}

// 范围扫描list
func (cc *boltDB) ScanRangeMap(bucketName, start, end string) (rv map[string][]byte, rer error) {
	rv = make(map[string][]byte)
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = ErrRangeStartInvalid
	if start == "" {
		return
	}
	rer = ErrRangeEndInvalid
	if end == "" {
		return
	}
	rer = cc.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		c := b.Cursor()
		for k, v := c.Seek([]byte(start)); k != nil && bytes.Compare(k, []byte(end)) <= 0; k, v = c.Next() {
			rv[string(k)] = v
		}
		return nil
	})
	return
}

// 批处理插入
func (cc *boltDB) BatchSet(bucketName string, value map[string][]byte) (rer error) {
	rer = ErrBucketInvalid
	if bucketName == "" {
		return
	}
	rer = cc.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		if b == nil {
			return ErrBucketNotFound
		}
		for k, v := range value {
			if k == "" {
				return ErrKeyInvalid
			}
			if value == nil {
				return ErrValueInvalid
			}
			if err := b.Put([]byte(k), v); err != nil {
				return err
			}
		}
		return nil
	})
	return
}