boltdb.go 2.8 KiB
Newer Older
Mayuri Shinde's avatar
Mayuri Shinde committed
package boltdb

import (
	"encoding/json"

	"corelab.mkcl.org/MKCLOS/coredevelopmentplatform/corepkgv2/errormdl"
	"github.com/boltdb/bolt"
)

//InitDB - To create database and bucket
func InitDB(boltDatabasePath, bucketName string) (*bolt.DB, error) {
	db, err := bolt.Open(boltDatabasePath, 0777, nil)
	if errormdl.CheckErr(err) != nil {
		return nil, errormdl.CheckErr(err)
	}
	err = db.Update(func(tx *bolt.Tx) error {
		_, bktErr := tx.CreateBucketIfNotExists([]byte(bucketName))
		if errormdl.CheckErr1(bktErr) != nil {
			return errormdl.CheckErr1(bktErr)
		}
		return nil
	})
	if errormdl.CheckErr(err) != nil {
		return nil, errormdl.Wrap("Could not set up bucket: " + bucketName)
	}
	return db, nil
}

// AddKey - To add key into bucket
func AddKey(db *bolt.DB, bucket string, key string, val string) error {
	return db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(bucket)).Put([]byte(key), []byte(val))
	})
}

// RemoveKey - To remove key from bucket
func RemoveKey(db *bolt.DB, bucket, key string) error {
	return db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(bucket)).Delete([]byte(key))
	})
}

// GetKey - To get data in bucket of given key
func GetKey(db *bolt.DB, bucket, key string) ([]byte, error) {
	var byteData []byte
	err := db.View(func(tx *bolt.Tx) error {
		bkt := tx.Bucket([]byte(bucket))
		if bkt == nil {
			return errormdl.Wrap("Bucket not found: " + bucket)
		}
		byteData = bkt.Get([]byte(key))
		return nil
	})
	if errormdl.CheckErr(err) != nil {
		return nil, errormdl.CheckErr(err)
	}
	return byteData, nil
}

//AddRecord - To Add,append or update data into a bucket
func AddRecord(db *bolt.DB, key string, data interface{}, bucket string) error {
	entryBytes, err := json.Marshal(data)
	if errormdl.CheckErr(err) != nil {
		return errormdl.Wrap("Could not marshal entry json")
	}
	err = db.Update(func(tx *bolt.Tx) error {
		err = tx.Bucket([]byte(bucket)).Put([]byte(key), entryBytes)
		if errormdl.CheckErr1(err) != nil {
			return errormdl.CheckErr1(err)
		}
		return nil
	})
	return errormdl.CheckErr(err)
}

// GetRecord - To get all data in bucket
func GetRecord(db *bolt.DB, bucket string) ([]byte, error) {
	var bucketData []interface{}
	err := db.View(func(tx *bolt.Tx) error {
		bkt := tx.Bucket([]byte(bucket))
		cursor := bkt.Cursor()
		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
			var interfaceObj interface{}
			unmarshallErr := json.Unmarshal(v, &interfaceObj)
			if errormdl.CheckErr(unmarshallErr) != nil {
				return errormdl.CheckErr(unmarshallErr)
			}
			bucketData = append(bucketData, interfaceObj)
		}
		return nil
	})
	if errormdl.CheckErr(err) != nil {
		return nil, errormdl.CheckErr(err)
	}
	byteData, marshallErr := json.Marshal(bucketData)
	if errormdl.CheckErr1(marshallErr) != nil {
		return nil, errormdl.CheckErr1(marshallErr)
	}
	return byteData, nil
}