open-consul/command/kvput/kv_put.go

package kvput

import (
	"bytes"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/command/flags"
	"github.com/mitchellh/cli"
)

func New(ui cli.Ui) *cmd {
	c := &cmd{UI: ui}
	c.init()
	return c
}

type cmd struct {
	UI    cli.Ui
	flags *flag.FlagSet
	http  *flags.HTTPFlags
	usage string

	// flags
	cas           bool
	kvflags       uint64
	base64encoded bool
	modifyIndex   uint64
	session       string
	acquire       bool
	release       bool

	// testStdin is the input for testing.
	testStdin io.Reader
}

func (c *cmd) init() {
	c.flags = flag.NewFlagSet("", flag.ContinueOnError)
	c.flags.BoolVar(&c.cas, "cas", false,
		"Perform a Check-And-Set operation. Specifying this value also "+
			"requires the -modify-index flag to be set. The default value "+
			"is false.")
	c.flags.Uint64Var(&c.kvflags, "flags", 0,
		"Unsigned integer value to assign to this key-value pair. This "+
			"value is not read by Consul, so clients can use this value however "+
			"makes sense for their use case. The default value is 0 (no flags).")
	c.flags.BoolVar(&c.base64encoded, "base64", false,
		"Treat the data as base 64 encoded. The default value is false.")
	c.flags.Uint64Var(&c.modifyIndex, "modify-index", 0,
		"Unsigned integer representing the ModifyIndex of the key. This is "+
			"used in combination with the -cas flag.")
	c.flags.StringVar(&c.session, "session", "",
		"User-defined identifer for this session as a string. This is commonly "+
			"used with the -acquire and -release operations to build robust locking, "+
			"but it can be set on any key. The default value is empty (no session).")
	c.flags.BoolVar(&c.acquire, "acquire", false,
		"Obtain a lock on the key. If the key does not exist, this operation "+
			"will create the key and obtain the lock. The session must already "+
			"exist and be specified via the -session flag. The default value is false.")
	c.flags.BoolVar(&c.release, "release", false,
		"Forfeit the lock on the key at the given path. This requires the "+
			"-session flag to be set. The key must be held by the session in order to "+
			"be unlocked. The default value is false.")

	c.http = &flags.HTTPFlags{}
	flags.Merge(c.flags, c.http.ClientFlags())
	flags.Merge(c.flags, c.http.ServerFlags())
	c.usage = flags.Usage(usage, c.flags, c.http.ClientFlags(), c.http.ServerFlags())
}

func (c *cmd) Run(args []string) int {
	if err := c.flags.Parse(args); err != nil {
		return 1
	}

	// Check for arg validation
	args = c.flags.Args()
	key, data, err := c.dataFromArgs(args)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error! %s", err))
		return 1
	}

	dataBytes := []byte(data)
	if c.base64encoded {
		dataBytes, err = base64.StdEncoding.DecodeString(data)
		if err != nil {
			c.UI.Error(fmt.Sprintf("Error! Cannot base 64 decode data: %s", err))
		}
	}

	// Session is reauired for release or acquire
	if (c.release || c.acquire) && c.session == "" {
		c.UI.Error("Error! Missing -session (required with -acquire and -release)")
		return 1
	}

	// ModifyIndex is required for CAS
	if c.cas && c.modifyIndex == 0 {
		c.UI.Error("Must specify -modify-index with -cas!")
		return 1
	}

	// Create and test the HTTP client
	client, err := c.http.APIClient()
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
		return 1
	}

	pair := &api.KVPair{
		Key:         key,
		ModifyIndex: c.modifyIndex,
		Flags:       c.kvflags,
		Value:       dataBytes,
		Session:     c.session,
	}

	switch {
	case c.cas:
		ok, _, err := client.KV().CAS(pair, nil)
		if err != nil {
			c.UI.Error(fmt.Sprintf("Error! Did not write to %s: %s", key, err))
			return 1
		}
		if !ok {
			c.UI.Error(fmt.Sprintf("Error! Did not write to %s: CAS failed", key))
			return 1
		}

		c.UI.Info(fmt.Sprintf("Success! Data written to: %s", key))
		return 0
	case c.acquire:
		ok, _, err := client.KV().Acquire(pair, nil)
		if err != nil {
			c.UI.Error(fmt.Sprintf("Error! Failed writing data: %s", err))
			return 1
		}
		if !ok {
			c.UI.Error("Error! Did not acquire lock")
			return 1
		}

		c.UI.Info(fmt.Sprintf("Success! Lock acquired on: %s", key))
		return 0
	case c.release:
		ok, _, err := client.KV().Release(pair, nil)
		if err != nil {
			c.UI.Error(fmt.Sprintf("Error! Failed writing data: %s", key))
			return 1
		}
		if !ok {
			c.UI.Error("Error! Did not release lock")
			return 1
		}

		c.UI.Info(fmt.Sprintf("Success! Lock released on: %s", key))
		return 0
	default:
		if _, err := client.KV().Put(pair, nil); err != nil {
			c.UI.Error(fmt.Sprintf("Error! Failed writing data: %s", err))
			return 1
		}

		c.UI.Info(fmt.Sprintf("Success! Data written to: %s", key))
		return 0
	}
}

func (c *cmd) Synopsis() string {
	return "Sets or updates data in the KV store"
}

func (c *cmd) Help() string {
	return c.usage
}

func (c *cmd) dataFromArgs(args []string) (string, string, error) {
	var stdin io.Reader = os.Stdin
	if c.testStdin != nil {
		stdin = c.testStdin
	}

	switch len(args) {
	case 0:
		return "", "", fmt.Errorf("Missing KEY argument")
	case 1:
		return args[0], "", nil
	case 2:
	default:
		return "", "", fmt.Errorf("Too many arguments (expected 1 or 2, got %d)", len(args))
	}

	key := args[0]
	data := args[1]

	// Handle empty quoted shell parameters
	if len(data) == 0 {
		return key, "", nil
	}

	switch data[0] {
	case '@':
		data, err := ioutil.ReadFile(data[1:])
		if err != nil {
			return "", "", fmt.Errorf("Failed to read file: %s", err)
		}
		return key, string(data), nil
	case '-':
		if len(data) > 1 {
			return key, data, nil
		}
		var b bytes.Buffer
		if _, err := io.Copy(&b, stdin); err != nil {
			return "", "", fmt.Errorf("Failed to read stdin: %s", err)
		}
		return key, b.String(), nil
	default:
		return key, data, nil
	}
}

const usage = `Usage: consul kv put [options] KEY [DATA]

  Writes the data to the given path in the key-value store. The data can be of
  any type.

      $ consul kv put config/redis/maxconns 5

  The data can also be consumed from a file on disk by prefixing with the "@"
  symbol. For example:

      $ consul kv put config/program/license @license.lic

  Or it can be read from stdin using the "-" symbol:

      $ echo "abcd1234" | consul kv put config/program/license -

  The DATA argument itself is optional. If omitted, this will create an empty
  key-value pair at the specified path:

      $ consul kv put webapp/beta/active

  If the -base64 flag is specified, the data will be treated as base 64
  encoded.

  To perform a Check-And-Set operation, specify the -cas flag with the
  appropriate -modify-index flag corresponding to the key you want to perform
  the CAS operation on:

      $ consul kv put -cas -modify-index=844 config/redis/maxconns 5

  Additional flags and more advanced use cases are detailed below.`