boltbrowser/bolt_model.go

package main

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"strings"

	"github.com/boltdb/bolt"
)

/*
BoltDB A Database, basically a collection of buckets
*/
type BoltDB struct {
	buckets []BoltBucket
}

type PathNode struct {
	name     []byte
	dataType Datatype
}

/*
BoltBucket is just a struct representation of a Bucket in the Bolt DB
*/
type BoltBucket struct {
	name         []byte
	nameDatatype Datatype
	pairs        []BoltPair
	buckets      []BoltBucket
	parent       *BoltBucket
	expanded     bool
	errorFlag    bool
}

func NewBoltBucket(parent *BoltBucket, name []byte) *BoltBucket {
	ret := new(BoltBucket)
	ret.parent = parent
	ret.name = name
	for _, dtk := range dataTypeNameSlice {
		if _, err := datatypes[dtk].ToString(name); err == nil {
			ret.nameDatatype = datatypes[dtk]
			break
		}
	}
	return ret
}

func (b *BoltBucket) GetPathNode() PathNode {
	return PathNode{
		name:     b.name,
		dataType: b.nameDatatype,
	}
}

func (b *BoltBucket) SetName(nm []byte) {
	b.name = nm
}

/*
BoltPair is just a struct representation of a Pair in the Bolt DB
*/
type BoltPair struct {
	parent      *BoltBucket
	key         []byte
	val         []byte
	keyDatatype Datatype
	valDatatype Datatype
}

func NewBoltPair(parent *BoltBucket, key, val []byte) *BoltPair {
	tp := BoltPair{
    parent: parent,
		key: key,
		val: val,
	}
	for _, dtk := range dataTypeNameSlice {
		if _, err := datatypes[dtk].ToString(key); err == nil {
			tp.keyDatatype = datatypes[dtk]
			break
		}
	}
	for _, dtk := range dataTypeNameSlice {
		if _, err := datatypes[dtk].ToString(val); err == nil {
			tp.valDatatype = datatypes[dtk]
			break
		}
	}
	return &tp
}

func (p *BoltPair) GetPathNode() PathNode {
	return PathNode{
		name:     p.key,
		dataType: p.keyDatatype,
	}
}

/*
GetStringKey returns the key of the pair after passing it through it's
Datatype ToString function
*/
func (p *BoltPair) GetStringKey() string {
	ret, err := p.keyDatatype.ToString(p.key)
	if err != nil {
		return "ERROR"
	}
	return ret
}


/*
GetStringVal returns the val of the pair after passing it through it's
Datatype ToString function
*/
func (p *BoltPair) GetStringVal() string {
	ret, err := p.valDatatype.ToString(p.val)
	if err != nil {
    return "ERROR"
	}
	return ret
}

func (bd *BoltDB) getGenericFromStringPath(path []PathNode) (*BoltBucket, *BoltPair, error) {
	return bd.getGenericFromPath(path)
}

func (bd *BoltDB) getGenericFromPath(path []PathNode) (*BoltBucket, *BoltPair, error) {
	// Check if 'path' leads to a pair
	p, err := bd.getPairFromPath(path)
	if err == nil {
		return nil, p, nil
	}
	// Nope, check if it leads to a bucket
	b, err := bd.getBucketFromPath(path)
	if err == nil {
		return b, nil, nil
	}
	// Nope, error
	return nil, nil, errors.New("Invalid Path")
}

func (bd *BoltDB) getBucketFromPath(path []PathNode) (*BoltBucket, error) {
	if len(path) > 0 {
		// Find the BoltBucket with a path == path
		var b *BoltBucket
		var err error
		// Find the root bucket
		b, err = memBolt.getBucket(path[0].name)
		if err != nil {
			return nil, err
		}
		if len(path) > 1 {
			for p := 1; p < len(path); p++ {
				b, err = b.getBucket(path[p].name)
				if err != nil {
					return nil, err
				}
			}
		}
		return b, nil
	}
	return nil, errors.New("Invalid Path")
}

func (bd *BoltDB) getPairFromPath(path []PathNode) (*BoltPair, error) {
	if len(path) <= 0 {
		return nil, errors.New("No Path")
	}
	b, err := bd.getBucketFromPath(path[:len(path)-1])
	if err != nil {
		return nil, err
	}
	// Found the bucket, pull out the pair
	p, err := b.getPair(path[len(path)-1].name)
	return p, err
}

func (bd *BoltDB) getVisibleItemCount(path []PathNode) (int, error) {
	vis := 0
	var retErr error
	if len(path) == 0 {
		for i := range bd.buckets {
			n, err := bd.getVisibleItemCount(bd.buckets[i].GetPath())
			if err != nil {
				return 0, err
			}
			vis += n
		}
	} else {
		b, err := bd.getBucketFromPath(path)
		if err != nil {
			return 0, err
		}
		// 1 for the bucket
		vis++
		if b.expanded {
			// This bucket is expanded, add up it's children
			// * 1 for each pair
			vis += len(b.pairs)
			// * recurse for buckets
			for i := range b.buckets {
				n, err := bd.getVisibleItemCount(b.buckets[i].GetPath())
				if err != nil {
					return 0, err
				}
				vis += n
			}
		}
	}
	return vis, retErr
}

func (bd *BoltDB) buildVisiblePathSlice() ([][]PathNode, error) {
	var retSlice [][]PathNode
	var retErr error
	// The root path, recurse for root buckets
	for i := range bd.buckets {
		bktS, bktErr := bd.buckets[i].buildVisiblePathSlice([]PathNode{})
		if bktErr == nil {
			retSlice = append(retSlice, bktS...)
		} else {
			// Something went wrong, set the error flag
			bd.buckets[i].errorFlag = true
		}
	}
	return retSlice, retErr
}

func (bd *BoltDB) getPrevVisiblePath(path []PathNode) []PathNode {
	visPaths, err := bd.buildVisiblePathSlice()
	if path == nil {
		if len(visPaths) > 0 {
			return visPaths[len(visPaths)-1]
		}
		return nil
	}
	if err == nil {
		for idx, pth := range visPaths {
			isCurPath := true
			for i := range path {
				if len(pth) <= i || !bytes.Equal(path[i].name, pth[i].name) {
					isCurPath = false
					break
				}
			}
			if isCurPath && idx > 0 {
				return visPaths[idx-1]
			}
		}
	}
	return nil
}
func (bd *BoltDB) getNextVisiblePath(path []PathNode) []PathNode {
	visPaths, err := bd.buildVisiblePathSlice()
	if path == nil {
		if len(visPaths) > 0 {
			return visPaths[0]
		}
		return nil
	}
	if err == nil {
		for idx, pth := range visPaths {
			isCurPath := true
			for i := range path {
				if len(pth) <= i || !bytes.Equal(path[i].name, pth[i].name) {
					isCurPath = false
					break
				}
			}
			if isCurPath && len(visPaths) > idx+1 {
				return visPaths[idx+1]
			}
		}
	}
	return nil
}

func (bd *BoltDB) toggleOpenBucket(path []PathNode) error {
	// Find the BoltBucket with a path == path
	b, err := bd.getBucketFromPath(path)
	if err == nil {
		b.expanded = !b.expanded
	}
	return err
}

func (bd *BoltDB) closeBucket(path []PathNode) error {
	// Find the BoltBucket with a path == path
	b, err := bd.getBucketFromPath(path)
	if err == nil {
		b.expanded = false
	}
	return err
}

func (bd *BoltDB) openBucket(path []PathNode) error {
	// Find the BoltBucket with a path == path
	b, err := bd.getBucketFromPath(path)
	if err == nil {
		b.expanded = true
	}
	return err
}

func (bd *BoltDB) getBucket(k []byte) (*BoltBucket, error) {
	for i := range bd.buckets {
		if bytes.Equal(bd.buckets[i].name, k) {
			return &bd.buckets[i], nil
		}
	}
	return nil, errors.New("Bucket Not Found")
}

func (bd *BoltDB) openAllBuckets() {
	for i := range bd.buckets {
		bd.buckets[i].openAllBuckets()
		bd.buckets[i].expanded = true
	}
}

func (bd *BoltDB) syncOpenBuckets(shadow *BoltDB) {
	// First test this bucket
	for i := range bd.buckets {
		for j := range shadow.buckets {
			if bytes.Equal(bd.buckets[i].name, shadow.buckets[j].name) {
				bd.buckets[i].syncOpenBuckets(&shadow.buckets[j])
			}
		}
	}
}

func (bd *BoltDB) refreshDatabase() *BoltDB {
	// Reload the database into memBolt
	memBolt = new(BoltDB)
	db.View(func(tx *bolt.Tx) error {
		return tx.ForEach(func(nm []byte, b *bolt.Bucket) error {
			bb, err := readBucket(nil, b, nm)
			if err == nil {
				bb.SetName(nm)
				bb.expanded = false
				memBolt.buckets = append(memBolt.buckets, *bb)
				return nil
			}
			return err
		})
	})
	return memBolt
}

/*
GetStringName returns the name of the bucket after passing it through it's
Datatype ToString function
*/
func (b *BoltBucket) GetStringName() string {
	ret, err := b.nameDatatype.ToString(b.name)
	if err != nil {
		return "ERROR"
	}
	return ret
}

/*
GetPath returns the database path leading to this BoltBucket
*/
func (b *BoltBucket) GetPath() []PathNode {
	bktPath := b.GetPathNode()
	if b.parent != nil {
		return append(b.parent.GetPath(), bktPath)
	}
	return []PathNode{bktPath}
}

/*
buildVisiblePathSlice builds a slice of PathNode slices containing all visible paths in this bucket
The passed prefix is the path leading to the current bucket
*/
func (b *BoltBucket) buildVisiblePathSlice(prefix []PathNode) ([][]PathNode, error) {
	var retSlice [][]PathNode
	var retErr error
	bucketNode := b.GetPathNode()
	retSlice = append(retSlice, append(prefix, bucketNode))
	if b.expanded {
		// Add subbuckets
		for i := range b.buckets {
			bktS, bktErr := b.buckets[i].buildVisiblePathSlice(append(prefix, bucketNode))
			if bktErr != nil {
				return retSlice, bktErr
			}
			retSlice = append(retSlice, bktS...)
		}
		// Add pairs
		for i := range b.pairs {
			retSlice = append(retSlice, append(append(prefix, bucketNode, b.pairs[i].GetPathNode())))
		}
	}
	return retSlice, retErr
}

func (b *BoltBucket) syncOpenBuckets(shadow *BoltBucket) {
	// First test this bucket
	b.expanded = shadow.expanded
	for i := range b.buckets {
		for j := range shadow.buckets {
			if bytes.Equal(b.buckets[i].name, shadow.buckets[j].name) {
				b.buckets[i].syncOpenBuckets(&shadow.buckets[j])
			}
		}
	}
}

func (b *BoltBucket) openAllBuckets() {
	for i := range b.buckets {
		b.buckets[i].openAllBuckets()
		b.buckets[i].expanded = true
	}
}

func (b *BoltBucket) getBucket(k []byte) (*BoltBucket, error) {
	for i := range b.buckets {
		if bytes.Equal(b.buckets[i].name, k) {
			return &b.buckets[i], nil
		}
	}
	return nil, errors.New("Bucket Not Found")
}

func (b *BoltBucket) getPair(k []byte) (*BoltPair, error) {
	for i := range b.pairs {
		if bytes.Equal(b.pairs[i].key, k) {
			return &b.pairs[i], nil
		}
	}
	return nil, errors.New("Pair Not Found")
}

/*
GetPath Returns the path of the BoltPair
*/
func (p *BoltPair) GetPath() []PathNode {
  if p.parent == nil {
    return []PathNode{p.GetPathNode()}
  }
	return append(
		p.parent.GetPath(),
		p.GetPathNode(),
	)
}

/* This is a go-between function (between the boltbrowser structs
 * above, and the bolt convenience functions below)
 * for taking a boltbrowser bucket and recursively adding it
 * and all of it's children into the database.
 * Mainly used for moving a bucket from one path to another
 * as in the 'renameBucket' function below.
 */
func addBucketFromBoltBucket(path []PathNode, bb *BoltBucket) error {
	if err := insertBucket(path, bb.name); err == nil {
		bucketPath := append(path, bb.GetPathNode())
		for i := range bb.pairs {
			if err = insertPair(bucketPath, bb.pairs[i].key, bb.pairs[i].val); err != nil {
				return err
			}
		}
		for i := range bb.buckets {
			if err = addBucketFromBoltBucket(bucketPath, &bb.buckets[i]); err != nil {
				return err
			}
		}
	}
	return nil
}

func deleteKey(path []PathNode) error {
	if AppArgs.ReadOnly {
		return errors.New("DB is in Read-Only Mode")
	}
	err := db.Update(func(tx *bolt.Tx) error {
		// len(b.path)-1 is the key we need to delete,
		// the rest are buckets leading to that key
		if len(path) == 1 {
			// Deleting a root bucket
			return tx.DeleteBucket(path[0].name)
		}
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket(path[i+1].name)
					if b == nil {
						return errors.New("deleteKey: Invalid Path")
					}
				}
			}
			// Now delete the last key in the path
			var err error
			if deleteBkt := b.Bucket(path[len(path)-1].name); deleteBkt == nil {
				// Must be a pair
				err = b.Delete(path[len(path)-1].name)
			} else {
				err = b.DeleteBucket(path[len(path)-1].name)
			}
			return err
		}
		return errors.New("deleteKey: Invalid Path")
	})
	return err
}

// Recursively read a bucket and pairs from the DB
func readBucket(parent *BoltBucket, b *bolt.Bucket, nm []byte) (*BoltBucket, error) {
  bb := NewBoltBucket(parent, nm)
	b.ForEach(func(k, v []byte) error {
		if v == nil {
			tb, err := readBucket(bb, b.Bucket(k), k)
			if err == nil {
				bb.buckets = append(bb.buckets, *tb)
			}
		} else {
			bb.pairs = append(bb.pairs, *NewBoltPair(bb, k, v))
		}
		return nil
	})
	return bb, nil
}

func renameBucket(path []PathNode, name []byte) error {
	if bytes.Equal(name, path[len(path)-1].name) {
		// No change requested
		return nil
	}
	var bb *BoltBucket // For caching the current bucket
  bb, err := memBolt.getBucketFromPath(path)
  if err != nil {
    return err
  }

	// Ok, we have the bucket, now delete the current instance
	if err = deleteKey(path); err != nil {
		return err
	}
	// Rechristen our cached bucket
	bb.SetName(name)
	// And re-add it
	parentPath := path[:len(path)-1]
	if err = addBucketFromBoltBucket(parentPath, bb); err != nil {
		return err
	}
	return nil
}

func updatePairKey(path []PathNode, k []byte) error {
	if AppArgs.ReadOnly {
		return errors.New("DB is in Read-Only Mode")
	}
	err := db.Update(func(tx *bolt.Tx) error {
		// len(b.path)-1 is the key for the pair we're updating,
		// the rest are buckets leading to that key
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 0 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket(path[i+1].name)
					if b == nil {
						return errors.New("updatePairValue: Invalid Path")
					}
				}
			}
			bk := path[len(path)-1]
			v := b.Get(bk.name)
			err := b.Delete(bk.name)
			if err == nil {
				// Old pair has been deleted, now add the new one
				err = b.Put(k, v)
			}
			// Now update the last key in the path
			return err
		}
		return errors.New("updatePairValue: Invalid Path")
	})
	return err
}

func updatePairValue(path []PathNode, v []byte) error {
	if AppArgs.ReadOnly {
		return errors.New("DB is in Read-Only Mode")
	}
	err := db.Update(func(tx *bolt.Tx) error {
		// len(b.GetPath())-1 is the key for the pair we're updating,
		// the rest are buckets leading to that key
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 0 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket(path[i+1].name)
					if b == nil {
						return errors.New("updatePairValue: Invalid Path")
					}
				}
			}
			// Now update the last key in the path
			err := b.Put(path[len(path)-1].name, v)
			return err
		}
		return errors.New("updatePairValue: Invalid Path")
	})
	return err
}

func insertBucket(path []PathNode, n []byte) error {
	if AppArgs.ReadOnly {
		return errors.New("DB is in Read-Only Mode")
	}
	// Inserts a new bucket named 'n' at 'path'
	err := db.Update(func(tx *bolt.Tx) error {
		if len(path) == 0 {
			// insert at root
			_, err := tx.CreateBucket(n)
			if err != nil {
				return fmt.Errorf("insertBucket: %s", err)
			}
		} else {
			rootBucket, path := path[0], path[1:]
			b := tx.Bucket(rootBucket.name)
			if b != nil {
				for len(path) > 0 {
					var tstBucket PathNode
					tstBucket, path = path[0], path[1:]
					nB := b.Bucket(tstBucket.name)
					if nB == nil {
						// Not a bucket, if we're out of path, just move on
						if len(path) != 0 {
							// Out of path, error
							return errors.New("insertBucket: Invalid Path 1")
						}
					} else {
						b = nB
					}
				}
				_, err := b.CreateBucket(n)
				return err
			}
			return fmt.Errorf("insertBucket: Invalid Path %s", rootBucket)
		}
		return nil
	})
	return err
}

func insertPair(path []PathNode, k []byte, v []byte) error {
	if AppArgs.ReadOnly {
		return errors.New("DB is in Read-Only Mode")
	}
	// Insert a new pair k => v at path
	err := db.Update(func(tx *bolt.Tx) error {
		if len(path) == 0 {
			// We cannot insert a pair at root
			return errors.New("insertPair: Cannot insert pair at root")
		}
		var err error
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 0 {
				for i := 1; i < len(path); i++ {
					b = b.Bucket(path[i].name)
					if b == nil {
						return fmt.Errorf("insertPair: %s", err)
					}
				}
			}
			err := b.Put(k, v)
			if err != nil {
				return fmt.Errorf("insertPair: %s", err)
			}
		}
		return nil
	})
	return err
}

func stringPathToBytePath(path []string) [][]byte {
	var ret [][]byte
	for _, v := range path {
		ret = append(ret, []byte(v))
	}
	return ret
}

func bytePathToStringPath(path [][]byte) []string {
	var ret []string
	for _, v := range path {
		ret = append(ret, string(v))
	}
	return ret
}

func nodePathToStringPath(path []PathNode) []string {
	var ret []string
	for _, v := range path {
		add, err := v.dataType.ToString(v.name)
		if err != nil {
			add = "ERROR"
		}
		ret = append(ret, add)
	}
	return ret
}

func exportValue(path []PathNode, fName string) error {
	return db.View(func(tx *bolt.Tx) error {
		// len(b.path)-1 is the key whose value we want to export
		// the rest are buckets leading to that key
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket(path[i+1].name)
					if b == nil {
						return errors.New("exportValue: Invalid Path: " + strings.Join(nodePathToStringPath(path), "/"))
					}
				}
			}
			bk := path[len(path)-1]
			v := b.Get(bk.name)
			return writeToFile(fName, string(v)+"\n", os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
		}
		return errors.New("exportValue: Invalid Bucket")
	})
}

func exportJSON(path []PathNode, fName string) error {
	return db.View(func(tx *bolt.Tx) error {
		// len(b.path)-1 is the key whose value we want to export
		// the rest are buckets leading to that key
		b := tx.Bucket(path[0].name)
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket(path[i+1].name)
					if b == nil {
						return errors.New("exportValue: Invalid Path: " + strings.Join(nodePathToStringPath(path), "/"))
					}
				}
			}
			bk := path[len(path)-1]
			if v := b.Get(bk.name); v != nil {
				str, err := bk.dataType.ToString(bk.name)
				if err != nil {
					str = "ERROR"
				}
				return writeToFile(fName, "{\""+str+"\":\""+string(v)+"\"}", os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
			}
			if b.Bucket(bk.name) != nil {
				return writeToFile(fName, genJSONString(b.Bucket(bk.name)), os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
			}
			return writeToFile(fName, genJSONString(b), os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
		}
		return errors.New("exportValue: Invalid Bucket")
	})
}

func genJSONString(b *bolt.Bucket) string {
	ret := "{"
	b.ForEach(func(k, v []byte) error {
		ret = fmt.Sprintf("%s\"%s\":", ret, string(k))
		if v == nil {
			ret = fmt.Sprintf("%s%s,", ret, genJSONString(b.Bucket(k)))
		} else {
			ret = fmt.Sprintf("%s\"%s\",", ret, string(v))
		}
		return nil
	})
	ret = fmt.Sprintf("%s}", ret[:len(ret)-1])
	return ret
}

func logToFile(s string) error {
	return writeToFile("bolt-log", s+"\n", os.O_RDWR|os.O_APPEND)
}

func writeToFile(fn, s string, mode int) error {
	var f *os.File
	var err error
	if f == nil {
		f, err = os.OpenFile(fn, mode, 0660)
	}
	defer f.Close()
	if err != nil {
		return err
	}
	if _, err = f.WriteString(s); err != nil {
		return err
	}
	if err = f.Sync(); err != nil {
		return err
	}
	return nil
}