github.com/kjdelisle/consul@v1.4.5/command/kv/imp/kv_import.go (about)

     1  package imp
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/base64"
     6  	"encoding/json"
     7  	"errors"
     8  	"flag"
     9  	"fmt"
    10  	"io"
    11  	"io/ioutil"
    12  	"os"
    13  
    14  	"github.com/hashicorp/consul/api"
    15  	"github.com/hashicorp/consul/command/flags"
    16  	"github.com/hashicorp/consul/command/kv/impexp"
    17  	"github.com/mitchellh/cli"
    18  )
    19  
    20  func New(ui cli.Ui) *cmd {
    21  	c := &cmd{UI: ui}
    22  	c.init()
    23  	return c
    24  }
    25  
    26  type cmd struct {
    27  	UI    cli.Ui
    28  	flags *flag.FlagSet
    29  	http  *flags.HTTPFlags
    30  	help  string
    31  
    32  	// testStdin is the input for testing.
    33  	testStdin io.Reader
    34  }
    35  
    36  func (c *cmd) init() {
    37  	c.flags = flag.NewFlagSet("", flag.ContinueOnError)
    38  	c.http = &flags.HTTPFlags{}
    39  	flags.Merge(c.flags, c.http.ClientFlags())
    40  	flags.Merge(c.flags, c.http.ServerFlags())
    41  	c.help = flags.Usage(help, c.flags)
    42  }
    43  
    44  func (c *cmd) Run(args []string) int {
    45  	if err := c.flags.Parse(args); err != nil {
    46  		return 1
    47  	}
    48  
    49  	// Check for arg validation
    50  	args = c.flags.Args()
    51  	data, err := c.dataFromArgs(args)
    52  	if err != nil {
    53  		c.UI.Error(fmt.Sprintf("Error! %s", err))
    54  		return 1
    55  	}
    56  
    57  	// Create and test the HTTP client
    58  	client, err := c.http.APIClient()
    59  	if err != nil {
    60  		c.UI.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
    61  		return 1
    62  	}
    63  
    64  	var entries []*impexp.Entry
    65  	if err := json.Unmarshal([]byte(data), &entries); err != nil {
    66  		c.UI.Error(fmt.Sprintf("Cannot unmarshal data: %s", err))
    67  		return 1
    68  	}
    69  
    70  	for _, entry := range entries {
    71  		value, err := base64.StdEncoding.DecodeString(entry.Value)
    72  		if err != nil {
    73  			c.UI.Error(fmt.Sprintf("Error base 64 decoding value for key %s: %s", entry.Key, err))
    74  			return 1
    75  		}
    76  
    77  		pair := &api.KVPair{
    78  			Key:   entry.Key,
    79  			Flags: entry.Flags,
    80  			Value: value,
    81  		}
    82  
    83  		if _, err := client.KV().Put(pair, nil); err != nil {
    84  			c.UI.Error(fmt.Sprintf("Error! Failed writing data for key %s: %s", pair.Key, err))
    85  			return 1
    86  		}
    87  
    88  		c.UI.Info(fmt.Sprintf("Imported: %s", pair.Key))
    89  	}
    90  
    91  	return 0
    92  }
    93  
    94  func (c *cmd) dataFromArgs(args []string) (string, error) {
    95  	var stdin io.Reader = os.Stdin
    96  	if c.testStdin != nil {
    97  		stdin = c.testStdin
    98  	}
    99  
   100  	switch len(args) {
   101  	case 0:
   102  		return "", errors.New("Missing DATA argument")
   103  	case 1:
   104  	default:
   105  		return "", fmt.Errorf("Too many arguments (expected 1, got %d)", len(args))
   106  	}
   107  
   108  	data := args[0]
   109  
   110  	if len(data) == 0 {
   111  		return "", errors.New("Empty DATA argument")
   112  	}
   113  
   114  	switch data[0] {
   115  	case '@':
   116  		data, err := ioutil.ReadFile(data[1:])
   117  		if err != nil {
   118  			return "", fmt.Errorf("Failed to read file: %s", err)
   119  		}
   120  		return string(data), nil
   121  	case '-':
   122  		if len(data) > 1 {
   123  			return data, nil
   124  		}
   125  		var b bytes.Buffer
   126  		if _, err := io.Copy(&b, stdin); err != nil {
   127  			return "", fmt.Errorf("Failed to read stdin: %s", err)
   128  		}
   129  		return b.String(), nil
   130  	default:
   131  		return data, nil
   132  	}
   133  }
   134  
   135  func (c *cmd) Synopsis() string {
   136  	return synopsis
   137  }
   138  
   139  func (c *cmd) Help() string {
   140  	return c.help
   141  }
   142  
   143  const synopsis = "Imports a tree stored as JSON to the KV store"
   144  const help = `
   145  Usage: consul kv import [DATA]
   146  
   147    Imports key-value pairs to the key-value store from the JSON representation
   148    generated by the "consul kv export" command.
   149  
   150    The data can be read from a file by prefixing the filename with the "@"
   151    symbol. For example:
   152  
   153        $ consul kv import @filename.json
   154  
   155    Or it can be read from stdin using the "-" symbol:
   156  
   157        $ cat filename.json | consul kv import -
   158  
   159    Alternatively the data may be provided as the final parameter to the command,
   160    though care must be taken with regards to shell escaping.
   161  
   162    For a full list of options and examples, please see the Consul documentation.
   163  `