bitcask/cmd/bitcask/import.go

package main

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"io"
	"os"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/prologic/bitcask"
)

var importCmd = &cobra.Command{
	Use:     "import",
	Aliases: []string{"restore", "read"},
	Short:   "Import a database",
	Long: `This command allows you to import or restore a database from a
previous export/dump using the export command either creating a new database
or adding additional key/value pairs to an existing one.`,
	Args: cobra.RangeArgs(0, 1),
	Run: func(cmd *cobra.Command, args []string) {
		var input string

		path := viper.GetString("path")

		if len(args) == 1 {
			input = args[0]
		} else {
			input = "-"
		}

		os.Exit(_import(path, input))
	},
}

func init() {
	RootCmd.AddCommand(importCmd)
}

func _import(path, input string) int {
	var (
		err error
		r   io.ReadCloser
	)

	db, err := bitcask.Open(path)
	if err != nil {
		log.WithError(err).Error("error opening database")
		return 1
	}
	defer db.Close()

	if input == "-" {
		r = os.Stdin
	} else {
		r, err = os.Open(input)
		if err != nil {
			log.WithError(err).
				WithField("input", input).
				Error("error opening input for reading")
			return 1
		}
	}

	var kv kvPair

	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		if err := json.Unmarshal(scanner.Bytes(), &kv); err != nil {
			log.WithError(err).
				WithField("input", input).
				Error("error reading input")
			return 2
		}

		key, err := base64.StdEncoding.DecodeString(kv.Key)
		if err != nil {
			log.WithError(err).Error("error decoding key")
			return 2
		}

		value, err := base64.StdEncoding.DecodeString(kv.Value)
		if err != nil {
			log.WithError(err).Error("error decoding value")
			return 2
		}

		if err := db.Put(string(key), value); err != nil {
			log.WithError(err).Error("error writing key/value")
			return 2
		}
	}
	if err := scanner.Err(); err != nil {
		log.WithError(err).
			WithField("input", input).
			Error("error reading input")
		return 2

	}

	return 0
}