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) }