boltbrowser/models/bolt.go

package models

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

	"git.bullercodeworks.com/brian/boltbrowser/util"
	bolt "go.etcd.io/bbolt"
)

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

func NewBoltDB(db *bolt.DB) *BoltDB {
	b := &BoltDB{db: db}
	b.RefreshDatabase()
	return b
}

func (bd *BoltDB) RefreshDatabase() {
	// Reload the database from the file
	bd.buckets = []BoltBucket{}
	bd.db.View(func(tx *bolt.Tx) error {
		err := tx.ForEach(func(nm []byte, b *bolt.Bucket) error {
			bb, err := bd.ReadBucket(b)
			if err == nil {
				bb.name = string(nm)
				bb.expanded = false
				bd.buckets = append(bd.buckets, *bb)
				return nil
			}
			return err
		})
		if err != nil {
			// Check if there are key/values directly in the root
			bb, err := bd.ReadBucket(tx.Cursor().Bucket())
			if err == nil {
				bb.isRoot = true
				bb.expanded = true
				bd.buckets = append(bd.buckets, *bb)
				return nil
			}
			return err
		}
		return err
	})
}

func (bd *BoltDB) GetBuckets() []BoltBucket { return bd.buckets }

func (bd *BoltDB) GetGenericFromPath(path []string) (*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 []string) (*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 = bd.GetBucket(path[0])
		if err != nil {
			return nil, err
		}
		if len(path) > 1 {
			for p := 1; p < len(path); p++ {
				b, err = b.GetBucket(path[p])
				if err != nil {
					return nil, err
				}
			}
		}
		return b, nil
	}
	return nil, errors.New("Invalid Path")
}

func (bd *BoltDB) GetPairFromPath(path []string) (*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])
	return p, err
}

func (bd *BoltDB) GetVisibleItemCount(path []string) (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(filter string) ([][]string, error) {
	var retSlice [][]string
	var retErr error
	// The root path, recurse for root buckets
	for i := range bd.buckets {
		bktS, bktErr := bd.buckets[i].BuildVisiblePathSlice([]string{}, filter)
		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) IsVisiblePath(path []string, filter string) bool {
	visPaths, err := bd.BuildVisiblePathSlice(filter)
	if err != nil {
		return false
	}
	for _, pth := range visPaths {
		if len(pth) != len(path) {
			continue
		}
		isVisible := true
		for i := range path {
			if path[i] != pth[i] {
				isVisible = false
				break
			}
		}
		if isVisible {
			return true
		}
	}
	return false
}
func (bd *BoltDB) GetPrevVisiblePath(path []string, filter string) []string {
	visPaths, err := bd.BuildVisiblePathSlice(filter)
	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 || path[i] != pth[i] {
					isCurPath = false
					break
				}
			}
			if isCurPath && idx > 0 {
				return visPaths[idx-1]
			}
		}
	}
	return nil
}
func (bd *BoltDB) GetNextVisiblePath(path []string, filter string) []string {
	visPaths, err := bd.BuildVisiblePathSlice(filter)
	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 || path[i] != pth[i] {
					isCurPath = false
					break
				}
			}
			if isCurPath && len(visPaths) > idx+1 {
				return visPaths[idx+1]
			}
		}
	}
	return nil
}

func (bd *BoltDB) ToggleOpenBucket(path []string) 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 []string) 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 []string) 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 string) (*BoltBucket, error) {
	for i := range bd.buckets {
		if 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 bd.buckets[i].name == shadow.buckets[j].name {
				bd.buckets[i].SyncOpenBuckets(&shadow.buckets[j])
			}
		}
	}
}

func (bd *BoltDB) DeleteKey(path []string) error {
	err := bd.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([]byte(path[0]))
		}
		b := tx.Bucket([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("DeleteKey: Invalid Path")
					}
				}
			}
			// Now delete the last key in the path
			var err error
			if deleteBkt := b.Bucket([]byte(path[len(path)-1])); deleteBkt == nil {
				// Must be a pair
				err = b.Delete([]byte(path[len(path)-1]))
			} else {
				err = b.DeleteBucket([]byte(path[len(path)-1]))
			}
			return err
		}
		return errors.New("DeleteKey: Invalid Path")
	})
	return err
}

func (bd *BoltDB) ReadBucket(b *bolt.Bucket) (*BoltBucket, error) {
	bb := new(BoltBucket)
	if b == nil {
		return nil, errors.New("No bucket passed")
	}
	b.ForEach(func(k, v []byte) error {
		if v == nil {
			tb, err := bd.ReadBucket(b.Bucket(k))
			tb.parent = bb
			if err == nil {
				tb.name = string(k)
				bb.buckets = append(bb.buckets, *tb)
			}
		} else {
			tp := BoltPair{key: string(k), val: string(v)}
			tp.parent = bb
			bb.pairs = append(bb.pairs, tp)
		}
		return nil
	})
	return bb, nil
}

func (bd *BoltDB) RenameBucket(path []string, name string) error {
	if name == path[len(path)-1] {
		// No change requested
		return nil
	}
	var bb *BoltBucket // For caching the current bucket
	err := bd.db.View(func(tx *bolt.Tx) error {
		// len(b.path)-1 is the key we need to delete,
		// the rest are buckets leading to that key
		b := tx.Bucket([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 1 {
				for i := range path[1:] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("RenameBucket: Invalid Path")
					}
				}
			}
			var err error
			// Ok, cache b
			bb, err = bd.ReadBucket(b)
			if err != nil {
				return err
			}
		} else {
			return errors.New("RenameBucket: Invalid Bucket")
		}
		return nil
	})
	if err != nil {
		return err
	}
	if bb == nil {
		return errors.New("RenameBucket: Couldn't find Bucket")
	}

	// Ok, we have the bucket cached, now delete the current instance
	if err = bd.DeleteKey(path); err != nil {
		return err
	}
	// Rechristen our cached bucket
	bb.name = name
	// And re-add it

	parentPath := path[:len(path)-1]
	if err = bd.AddBucketFromBoltBucket(parentPath, bb); err != nil {
		return err
	}
	return nil
}

func (bd *BoltDB) UpdatePairKey(path []string, k string) error {
	err := bd.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([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 0 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("UpdatePairKey: Invalid Path")
					}
				}
			}
			bk := []byte(path[len(path)-1])
			v := b.Get(bk)
			err := b.Delete(bk)
			if err == nil {
				// Old pair has been deleted, now add the new one
				err = b.Put([]byte(k), v)
			}
			// Now update the last key in the path
			return err
		}
		return errors.New("UpdatePairKey: Invalid Path")
	})
	return err
}

func (bd *BoltDB) UpdatePairValue(path []string, v string) error {
	err := bd.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([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 0 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("UpdatePairValue: Invalid Path")
					}
				}
			}
			// Now update the last key in the path
			err := b.Put([]byte(path[len(path)-1]), []byte(v))
			return err
		}
		return errors.New("UpdatePairValue: Invalid Path")
	})
	return err
}

func (bd *BoltDB) InsertBucket(path []string, n string) error {
	// Inserts a new bucket named 'n' at 'path'
	err := bd.db.Update(func(tx *bolt.Tx) error {
		if len(path) == 0 || path[0] == "" {
			// insert at root
			_, err := tx.CreateBucket([]byte(n))
			if err != nil {
				return fmt.Errorf("InsertBucket: %s", err)
			}
		} else {
			rootBucket, path := path[0], path[1:]
			b := tx.Bucket([]byte(rootBucket))
			if b != nil {
				for len(path) > 0 {
					tstBucket := ""
					tstBucket, path = path[0], path[1:]
					nB := b.Bucket([]byte(tstBucket))
					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([]byte(n))
				return err
			}
			return fmt.Errorf("insertBucket: Invalid Path %s", rootBucket)
		}
		return nil
	})
	return err
}

func (bd *BoltDB) InsertPair(path []string, k string, v string) error {
	// Insert a new pair k => v at path
	err := bd.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([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 0 {
				for i := 1; i < len(path); i++ {
					b = b.Bucket([]byte(path[i]))
					if b == nil {
						return fmt.Errorf("InsertPair: %s", err)
					}
				}
			}
			err := b.Put([]byte(k), []byte(v))
			if err != nil {
				return fmt.Errorf("InsertPair: %s", err)
			}
		}
		return nil
	})
	return err
}

func (bd *BoltDB) exportValue(path []string, fName string) error {
	return bd.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([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("ExportValue: Invalid Path: " + strings.Join(path, "/"))
					}
				}
			}
			bk := []byte(path[len(path)-1])
			v := b.Get(bk)
			return util.WriteToFile(fName, string(v)+"\n", os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
		}
		return errors.New("ExportValue: Invalid Bucket")
	})
}

func (bd *BoltDB) ExportJSON(path []string, fName string) error {
	return bd.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([]byte(path[0]))
		if b == nil {
			// Invalid path, try for the root bucket
			b = tx.Cursor().Bucket()
		}
		if b != nil {
			if len(path) > 1 {
				for i := range path[1 : len(path)-1] {
					b = b.Bucket([]byte(path[i+1]))
					if b == nil {
						return errors.New("ExportJSON: Invalid Path: " + strings.Join(path, "/"))
					}
				}
			}
			bk := []byte(path[len(path)-1])
			if v := b.Get(bk); v != nil {
				return util.WriteToFile(fName, "{\""+string(bk)+"\":\""+string(v)+"\"}", os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
			}
			if b.Bucket(bk) != nil {
				return util.WriteToFile(fName, util.ToJsonString(b.Bucket(bk)), os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
			}
			return util.WriteToFile(fName, util.ToJsonString(b), os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
		}
		return errors.New("ExportJSON: Invalid Bucket")
	})
}

func (bd *BoltDB) AddBucketFromBoltBucket(path []string, bb *BoltBucket) error {
	if err := bd.InsertBucket(path, bb.name); err == nil {
		bucketPath := append(path, bb.name)
		for i := range bb.pairs {
			if err = bd.InsertPair(bucketPath, bb.pairs[i].key, bb.pairs[i].val); err != nil {
				return err
			}
		}
		for i := range bb.buckets {
			if err = bd.AddBucketFromBoltBucket(bucketPath, &bb.buckets[i]); err != nil {
				return err
			}
		}
	}
	return nil
}

func (bd BoltDB) Lines() []string {
	var ret []string
	for i := range bd.buckets {
		ret = append(ret, bd.buckets[i].Lines()...)
	}
	return ret
}