diff --git a/cmd/bitcask/export.go b/cmd/bitcask/export.go new file mode 100644 index 0000000..1adbacd --- /dev/null +++ b/cmd/bitcask/export.go @@ -0,0 +1,142 @@ +package main + +import ( + "encoding/base64" + "encoding/json" + "errors" + "io" + "os" + + log "github.com/sirupsen/logrus" + "github.com/spf13/cobra" + "github.com/spf13/viper" + + "github.com/prologic/bitcask" +) + +var errNotAllDataWritten = errors.New("error: not all data written") + +var exportCmd = &cobra.Command{ + Use: "export", + Aliases: []string{"backup", "dump"}, + Short: "Export a database", + Long: `This command allows you to export or dump/backup a database's +key/values into a long-term portable archival format suitable for backup and +restore purposes or migrating from older on-disk formats of Bitcask. + +All key/value pairs are base64 encoded and serialized as JSON one pair per +line to form an output stream to either standard output or a file. You can +optionally compress the output with standard compression tools such as gzip.`, + Args: cobra.RangeArgs(0, 1), + Run: func(cmd *cobra.Command, args []string) { + var output string + + path := viper.GetString("path") + + if len(args) == 1 { + output = args[0] + } else { + output = "-" + } + + os.Exit(export(path, output)) + }, +} + +func init() { + RootCmd.AddCommand(exportCmd) + + exportCmd.PersistentFlags().IntP( + "with-max-datafile-size", "", bitcask.DefaultMaxDatafileSize, + "Maximum size of each datafile", + ) + exportCmd.PersistentFlags().IntP( + "with-max-key-size", "", bitcask.DefaultMaxKeySize, + "Maximum size of each key", + ) + exportCmd.PersistentFlags().IntP( + "with-max-value-size", "", bitcask.DefaultMaxValueSize, + "Maximum size of each value", + ) +} + +type kvPair struct { + Key string `json:"key"` + Value string `json:"value"` +} + +func export(path, output string) int { + var ( + err error + w io.WriteCloser + ) + + db, err := bitcask.Open(path) + if err != nil { + log.WithError(err).Error("error opening database") + return 1 + } + defer db.Close() + + if output == "-" { + w = os.Stdout + } else { + w, err = os.OpenFile(output, os.O_WRONLY|os.O_CREATE|os.O_EXCL|os.O_TRUNC, 0755) + if err != nil { + log.WithError(err). + WithField("output", output). + Error("error opening output for writing") + return 1 + } + } + + err = db.Fold(func(key string) error { + value, err := db.Get(key) + if err != nil { + log.WithError(err). + WithField("key", key). + Error("error reading key") + return err + } + + kv := kvPair{ + Key: base64.StdEncoding.EncodeToString([]byte(key)), + Value: base64.StdEncoding.EncodeToString(value), + } + + data, err := json.Marshal(&kv) + if err != nil { + log.WithError(err). + WithField("key", key). + Error("error serialzing key") + return err + } + + if n, err := w.Write(data); err != nil || n != len(data) { + if err == nil && n != len(data) { + err = errNotAllDataWritten + } + log.WithError(err). + WithField("key", key). + WithField("n", n). + Error("error writing key") + return err + } + + if _, err := w.Write([]byte("\n")); err != nil { + log.WithError(err).Error("error writing newline") + return err + } + + return nil + }) + if err != nil { + log.WithError(err). + WithField("path", path). + WithField("output", output). + Error("error exporting keys") + return 2 + } + + return 0 +} diff --git a/cmd/bitcask/import.go b/cmd/bitcask/import.go new file mode 100644 index 0000000..aa526be --- /dev/null +++ b/cmd/bitcask/import.go @@ -0,0 +1,106 @@ +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 +}