cli: Add KV export and import

command/kv_export.go

@@ -0,0 +1,123 @@
+package command
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"strings"
+
+	"github.com/hashicorp/consul/api"
+	"github.com/mitchellh/cli"
+)
+
+// KVExportCommand is a Command implementation that is used to export
+// a KV tree as JSON
+type KVExportCommand struct {
+	Ui cli.Ui
+}
+
+func (c *KVExportCommand) Synopsis() string {
+	return "Exports a tree from the KV store as JSON"
+}
+
+func (c *KVExportCommand) Help() string {
+	helpText := `
+Usage: consul kv export [KEY_OR_PREFIX]
+
+  Retrieves key-value pairs for the given prefix from Consul's key-value store,
+  and writes a JSON representation to stdout. This can be used with the command
+  "consul kv import" to move entire trees between Consul clusters.
+
+      $ consul kv export vault
+
+  For a full list of options and examples, please see the Consul documentation.
+
+` + apiOptsText + `
+
+KV Export Options:
+
+  None.
+`
+	return strings.TrimSpace(helpText)
+}
+
+func (c *KVExportCommand) Run(args []string) int {
+	cmdFlags := flag.NewFlagSet("export", flag.ContinueOnError)
+
+	datacenter := cmdFlags.String("datacenter", "", "")
+	token := cmdFlags.String("token", "", "")
+	stale := cmdFlags.Bool("stale", false, "")
+	httpAddr := HTTPAddrFlag(cmdFlags)
+	if err := cmdFlags.Parse(args); err != nil {
+		return 1
+	}
+
+	key := ""
+	// Check for arg validation
+	args = cmdFlags.Args()
+	switch len(args) {
+	case 0:
+		key = ""
+	case 1:
+		key = args[0]
+	default:
+		c.Ui.Error(fmt.Sprintf("Too many arguments (expected 1, got %d)", len(args)))
+		return 1
+	}
+
+	// This is just a "nice" thing to do. Since pairs cannot start with a /, but
+	// users will likely put "/" or "/foo", lets go ahead and strip that for them
+	// here.
+	if len(key) > 0 && key[0] == '/' {
+		key = key[1:]
+	}
+
+	// Create and test the HTTP client
+	conf := api.DefaultConfig()
+	conf.Address = *httpAddr
+	conf.Token = *token
+	client, err := api.NewClient(conf)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
+		return 1
+	}
+
+	pairs, _, err := client.KV().List(key, &api.QueryOptions{
+		Datacenter: *datacenter,
+		AllowStale: *stale,
+	})
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error querying Consul agent: %s", err))
+		return 1
+	}
+
+	exported := make([]*kvExportEntry, len(pairs))
+	for i, pair := range pairs {
+		exported[i] = toExportEntry(pair)
+	}
+
+	marshaled, err := json.MarshalIndent(exported, "", "\t")
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error exporting KV data: %s", err))
+		return 1
+	}
+
+	c.Ui.Info(string(marshaled))
+
+	return 0
+}
+
+type kvExportEntry struct {
+	Key   string `json:"key"`
+	Flags uint64 `json:"flags"`
+	Value string `json:"value"`
+}
+
+func toExportEntry(pair *api.KVPair) *kvExportEntry {
+	return &kvExportEntry{
+		Key:   pair.Key,
+		Flags: pair.Flags,
+		Value: base64.StdEncoding.EncodeToString(pair.Value),
+	}
+}

command/kv_export_test.go

@@ -0,0 +1,60 @@
+package command
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"testing"
+
+	"github.com/hashicorp/consul/api"
+	"github.com/mitchellh/cli"
+)
+
+func TestKVExportCommand_Run(t *testing.T) {
+	srv, client := testAgentWithAPIClient(t)
+	defer srv.Shutdown()
+	waitForLeader(t, srv.httpAddr)
+
+	ui := new(cli.MockUi)
+	c := &KVExportCommand{Ui: ui}
+
+	keys := map[string]string{
+		"foo/a": "a",
+		"foo/b": "b",
+		"foo/c": "c",
+		"bar":   "d",
+	}
+	for k, v := range keys {
+		pair := &api.KVPair{Key: k, Value: []byte(v)}
+		if _, err := client.KV().Put(pair, nil); err != nil {
+			t.Fatalf("err: %#v", err)
+		}
+	}
+
+	args := []string{
+		"-http-addr=" + srv.httpAddr,
+		"foo",
+	}
+
+	code := c.Run(args)
+	if code != 0 {
+		t.Fatalf("bad: %d. %#v", code, ui.ErrorWriter.String())
+	}
+
+	output := ui.OutputWriter.String()
+
+	var exported []*kvExportEntry
+	err := json.Unmarshal([]byte(output), &exported)
+	if err != nil {
+		t.Fatalf("bad: %s", code)
+	}
+
+	if len(exported) != 3 {
+		t.Fatalf("bad: expected 3, got %d", len(exported))
+	}
+
+	for _, entry := range exported {
+		if base64.StdEncoding.EncodeToString([]byte(keys[entry.Key])) != entry.Value {
+			t.Fatalf("bad: expected %s, got %s", keys[entry.Key], entry.Value)
+		}
+	}
+}

command/kv_import.go

@@ -0,0 +1,165 @@
+package command
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"strings"
+
+	"github.com/hashicorp/consul/api"
+	"github.com/mitchellh/cli"
+)
+
+// KVImportCommand is a Command implementation that is used to import
+// a KV tree stored as JSON
+type KVImportCommand struct {
+	Ui cli.Ui
+
+	// testStdin is the input for testing.
+	testStdin io.Reader
+}
+
+func (c *KVImportCommand) Synopsis() string {
+	return "Imports a tree stored as JSON to the KV store"
+}
+
+func (c *KVImportCommand) Help() string {
+	helpText := `
+Usage: consul kv import [DATA]
+
+  Imports key-value pairs to the key-value store from the JSON representation
+  generated by the "consul kv export" command.
+
+  The data can be read from a file by prefixing the filename with the "@"
+  symbol. For example:
+
+      $ consul kv import @filename.json
+
+  Or it can be read from stdin using the "-" symbol:
+
+      $ cat filename.json | consul kv import config/program/license -
+
+  Alternatively the data may be provided as the final parameter to the command,
+  though care must be taken with regards to shell escaping.
+
+  For a full list of options and examples, please see the Consul documentation.
+
+` + apiOptsText + `
+
+KV Import Options:
+
+  None.
+`
+	return strings.TrimSpace(helpText)
+}
+
+func (c *KVImportCommand) Run(args []string) int {
+	cmdFlags := flag.NewFlagSet("import", flag.ContinueOnError)
+
+	datacenter := cmdFlags.String("datacenter", "", "")
+	token := cmdFlags.String("token", "", "")
+	httpAddr := HTTPAddrFlag(cmdFlags)
+	if err := cmdFlags.Parse(args); err != nil {
+		return 1
+	}
+
+	// Check for arg validation
+	args = cmdFlags.Args()
+	data, err := c.dataFromArgs(args)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error! %s", err))
+		return 1
+	}
+
+	// Create and test the HTTP client
+	conf := api.DefaultConfig()
+	conf.Address = *httpAddr
+	conf.Token = *token
+	client, err := api.NewClient(conf)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
+		return 1
+	}
+
+	var entries []*kvExportEntry
+	if err := json.Unmarshal([]byte(data), &entries); err != nil {
+		c.Ui.Error(fmt.Sprintf("Cannot unmarshal data: %s", err))
+		return 1
+	}
+
+	for _, entry := range entries {
+		value, err := base64.StdEncoding.DecodeString(entry.Value)
+		if err != nil {
+			c.Ui.Error(fmt.Sprintf("Error base 64 decoding value for key %s: %s", entry.Key, err))
+			return 1
+		}
+
+		pair := &api.KVPair{
+			Key:   entry.Key,
+			Flags: entry.Flags,
+			Value: value,
+		}
+
+		wo := &api.WriteOptions{
+			Datacenter: *datacenter,
+			Token:      *token,
+		}
+
+		if _, err := client.KV().Put(pair, wo); err != nil {
+			c.Ui.Error(fmt.Sprintf("Error! Failed writing data for key %s: %s", pair.Key, err))
+			return 1
+		}
+
+		c.Ui.Info(fmt.Sprintf("Imported: %s", pair.Key))
+	}
+
+	return 0
+}
+
+func (c *KVImportCommand) dataFromArgs(args []string) (string, error) {
+	var stdin io.Reader = os.Stdin
+	if c.testStdin != nil {
+		stdin = c.testStdin
+	}
+
+	switch len(args) {
+	case 0:
+		return "", errors.New("Missing DATA argument")
+	case 1:
+	default:
+		return "", fmt.Errorf("Too many arguments (expected 1 or 2, got %d)", len(args))
+	}
+
+	data := args[0]
+
+	if len(data) == 0 {
+		return "", errors.New("Empty DATA argument")
+	}
+
+	switch data[0] {
+	case '@':
+		data, err := ioutil.ReadFile(data[1:])
+		if err != nil {
+			return "", fmt.Errorf("Failed to read file: %s", err)
+		}
+		return string(data), nil
+	case '-':
+		if len(data) > 1 {
+			return data, nil
+		} else {
+			var b bytes.Buffer
+			if _, err := io.Copy(&b, stdin); err != nil {
+				return "", fmt.Errorf("Failed to read stdin: %s", err)
+			}
+			return b.String(), nil
+		}
+	default:
+		return data, nil
+	}
+}