bolt.go

package main
import (
    "bytes"
    "encoding/json"
    "errors"
    "reflect"
    "github.com/boltdb/bolt"
    "github.com/timshannon/bolthold"
)
// BoltDB WebUI 工具:
//    cd ./data
//    boltdbweb  -d=db.bolt -p=9000 -s="G:/Work/GoPath/src/github.com/evnix/boltdbweb/"
const (
    boltDBFileName = "data/db.bolt"
)
var (
    //BoltX BoltDB Store
    BoltX *BoltWrap
    boltXModels  = make(map[string]interface{}) // bolt models map
    boltXHasInit bool                           // bolt models has int
)
func init() {
    Registr2Bolt(&BoltIDIncrement{})
}
// initBoltDB 初始化 BoltDB 配置
func initBoltDB() {
    BoltX = NewBoltWrap(boltDBFileName)
    //INFO: bolthold 会为每个 typeOf(BoltModels) 创建一个 Bucket,
    // 对每个 Model 初始化一下, 避免没有对某 Model 进行过 `写` 操作之前 Get 的时候报 bucket nil 错误.
    for _, data := range boltXModels {
        if err := BoltX.Upsert("-", data); err != nil {
            panic(err)
        }
    }
}
// Registr2Bolt init Registr BoltDB Model
func Registr2Bolt(defaultData interface{}) {
    boltXModels[reflect.TypeOf(defaultData).String()] = defaultData
}
// GetBoltStore get bolt store
func GetBoltStore(dbFileName string) (*bolthold.Store, error) {
    if len(dbFileName) <= 0 {
        return nil, errors.New("dbFileName param nil")
    }
    //return bolthold.Open(dbFileName, 0666, nil)
    jsonOptions := &bolthold.Options{
        Encoder: DefaultEncode,
        Decoder: DefaultDecode,
    }
    return bolthold.Open(dbFileName, 0666, jsonOptions)
}
//----------------------BoltWrap----------------------
// BoltWrap bolt wrap
type BoltWrap struct {
    *bolthold.Store
    dbFileName string
}
// BoltIDIncrement ID Increment
type BoltIDIncrement struct {
    Current int64
}
// NewBoltWrap new BoltWrap
func NewBoltWrap(dbFileName string) *BoltWrap {
    db, err := GetBoltStore(dbFileName)
    if err != nil {
        panic("init BoltDB fail")
    }
    return &BoltWrap{
        dbFileName: dbFileName,
        Store:      db,
    }
}
// IsNotFoundError is error notFound
func (b *BoltWrap) IsNotFoundError(err error) bool {
    return err == bolthold.ErrNotFound
}
// GeNextBoltID Ge Next BoltID
func (b *BoltWrap) GeNextBoltID(tx *bolt.Tx, key interface{}) (int64, error) {
    if tx == nil || key == nil {
        return -1, errors.New("tx or key param nil")
    }
    id := &BoltIDIncrement{}
    err := b.TxGet(tx, key, id)
    if err != nil && !b.IsNotFoundError(err) {
        tx.Rollback()
        return -1, err
    }
    if id.Current <= 0 {
        id.Current = 1
        if err = b.TxUpsert(tx, key, id); err != nil {
            tx.Rollback()
            return -1, err
        }
    } else {
        id.Current++
        if err = b.TxUpdate(tx, key, id); err != nil {
            tx.Rollback()
            return -1, err
        }
    }
    return id.Current, nil
}
//-----------BoltDB JSON Encode/tDecode-------------
// DefaultEncode is the default encoding func for bolthold (Gob)
func DefaultEncode(value interface{}) ([]byte, error) {
    var buff bytes.Buffer
    en := json.NewEncoder(&buff)
    err := en.Encode(value)
    if err != nil {
        return nil, err
    }
    return buff.Bytes(), nil
}
// DefaultDecode is the default decoding func for bolthold (Gob)
func DefaultDecode(data []byte, value interface{}) error {
    var buff bytes.Buffer
    de := json.NewDecoder(&buff)
    _, err := buff.Write(data)
    if err != nil {
        return err
    }
    err = de.Decode(value)
    if err != nil {
        return err
    }
    return nil
}