boltprint/boltprint.go

package main

import (
	"fmt"
	"github.com/br0xen/bolt"
	"os"
	"strings"
)

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		fmt.Println("Usage: printbolt <Database File> [build]")
		os.Exit(1)
	}
	filename := args[0]

	var db *bolt.DB
	var err error

	if len(args) > 1 && args[1] == "build" {
		if _, stat_err := os.Stat(filename); stat_err == nil {
			fmt.Printf("Cowardly refusing to overwrite file: %s\n", filename)
			return
		}
	}
	if err != nil {
		fmt.Println(err)
	}

	//opts := bolt.Options{0, true}
	//db, err := bolt.Open(filename, 0600, &opts)
	db, err = bolt.Open(filename, 0600, nil)
	if err != nil {
		fmt.Printf("Error opening DB: %s\n", err)
		return
	}
	defer db.Close()

	if len(args) > 1 {
		if args[1] == "build" {
			createTestDatabase(db)
		} else if args[1] == "json" {
			jsonDatabase(db)
		} else if args[1] == "go" {
			fmt.Print(db.GoString())
		}
	} else {
		printDatabase(db)
	}
}

func printDatabase(db *bolt.DB) error {
	db.View(func(tx *bolt.Tx) error {
		tx.ForEach(func(nm []byte, b *bolt.Bucket) error {
			fmt.Printf("%s => \n", string(nm))
			printBucket(b, 0)
			return nil
		})
		return nil
	})
	return nil
}

func printBucket(b *bolt.Bucket, lvl int) {
	b.ForEach(func(k, v []byte) error {
		fmt.Printf("%s. %s => %s\n", strings.Repeat(".", lvl), k, v)
		if v == nil {
			printBucket(b.Bucket(k), lvl+1)
		}
		return nil
	})
}

func jsonDatabase(db *bolt.DB) error {
	out := "{"
	db.View(func(tx *bolt.Tx) error {
		tx.ForEach(func(nm []byte, b *bolt.Bucket) error {
			bkt := jsonBucket(b)
			out = fmt.Sprintf("%s\"%s\":%s,", out, string(nm), bkt[:len(bkt)-1])
			return nil
		})
		return nil
	})
	out = fmt.Sprintf("%s}", out[:len(out)-1])
	fmt.Print(out)
	return nil
}

func jsonBucket(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, jsonBucket(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 createTestDatabase(db *bolt.DB) error {
	fmt.Println("Generating Database")
	for i := 0; i < 10; i++ {
		fmt.Printf("Creating Bucket %d & Values", i)
		db.Update(func(tx *bolt.Tx) error {
			b, err := tx.CreateBucketIfNotExists([]byte(fmt.Sprintf("Bucket%d", i)))
			fmt.Print(".")
			ba, err := b.CreateBucketIfNotExists([]byte(fmt.Sprintf("Bucket%da", i)))
			bb, err := b.CreateBucketIfNotExists([]byte(fmt.Sprintf("Bucket%db", i)))
			for j := 0; j < 10; j++ {
				fmt.Print(".")
				b.Put([]byte(fmt.Sprintf("key%d", j)), []byte(fmt.Sprintf("value%d", j)))
				ba.Put([]byte(fmt.Sprintf("key%d", j)), []byte(fmt.Sprintf("value%d", j)))
				bb.Put([]byte(fmt.Sprintf("key%d", j)), []byte(fmt.Sprintf("value%d", j)))
			}

			if err != nil {
				return err
			}
			return nil
		})
		fmt.Printf("\n")
	}
	fmt.Print("\n")
	return nil
}