diff --git a/command/helpers.go b/command/helpers.go index 8cafabc3c32f..dfc8eea00387 100644 --- a/command/helpers.go +++ b/command/helpers.go @@ -1,8 +1,10 @@ package command import ( + "fmt" "time" + "github.com/hashicorp/nomad/api" "github.com/ryanuber/columnize" ) @@ -45,3 +47,25 @@ func limit(s string, length int) string { func formatTime(t time.Time) string { return t.Format("02/01/06 15:04:05 MST") } + +// getLocalNodeID returns the node ID of the local Nomad Client and an error if +// it couldn't be determined or the Agent is not running in Client mode. +func getLocalNodeID(client *api.Client) (string, error) { + info, err := client.Agent().Self() + if err != nil { + return "", fmt.Errorf("Error querying agent info: %s", err) + } + var stats map[string]interface{} + stats, _ = info["stats"] + clientStats, ok := stats["client"].(map[string]interface{}) + if !ok { + return "", fmt.Errorf("Nomad not running in client mode") + } + + nodeID, ok := clientStats["node_id"].(string) + if !ok { + return "", fmt.Errorf("Failed to determine node ID") + } + + return nodeID, nil +} diff --git a/command/helpers_test.go b/command/helpers_test.go index 327c4171277b..276719c33f56 100644 --- a/command/helpers_test.go +++ b/command/helpers_test.go @@ -2,6 +2,8 @@ package command import ( "testing" + + "github.com/mitchellh/cli" ) func TestHelpers_FormatKV(t *testing.T) { @@ -27,3 +29,19 @@ func TestHelpers_FormatList(t *testing.T) { t.Fatalf("expect: %s, got: %s", expect, out) } } + +func TestHelpers_NodeID(t *testing.T) { + srv, _, _ := testServer(t, nil) + defer srv.Stop() + + meta := Meta{Ui: new(cli.MockUi)} + client, err := meta.Client() + if err != nil { + t.FailNow() + } + + // This is because there is no client + if _, err := getLocalNodeID(client); err == nil { + t.Fatalf("getLocalNodeID() should fail") + } +} diff --git a/command/node_drain.go b/command/node_drain.go index 9880cb3db709..b5f814880c1b 100644 --- a/command/node_drain.go +++ b/command/node_drain.go @@ -15,6 +15,7 @@ Usage: nomad node-drain [options] Toggles node draining on a specified node. It is required that either -enable or -disable is specified, but not both. + The -self flag is useful to drain the local node. General Options: @@ -28,6 +29,9 @@ Node Drain Options: -enable Enable draining for the specified node. + -self + Query the status of the local node. + -yes Automatic yes to prompts. ` @@ -39,12 +43,13 @@ func (c *NodeDrainCommand) Synopsis() string { } func (c *NodeDrainCommand) Run(args []string) int { - var enable, disable, autoYes bool + var enable, disable, self, autoYes bool flags := c.Meta.FlagSet("node-drain", FlagSetClient) flags.Usage = func() { c.Ui.Output(c.Help()) } flags.BoolVar(&enable, "enable", false, "Enable drain mode") flags.BoolVar(&disable, "disable", false, "Disable drain mode") + flags.BoolVar(&self, "self", false, "") flags.BoolVar(&autoYes, "yes", false, "Automatic yes to prompts.") if err := flags.Parse(args); err != nil { @@ -59,11 +64,10 @@ func (c *NodeDrainCommand) Run(args []string) int { // Check that we got a node ID args = flags.Args() - if len(args) != 1 { + if l := len(args); self && l != 0 || !self && l != 1 { c.Ui.Error(c.Help()) return 1 } - nodeID := args[0] // Get the HTTP client client, err := c.Meta.Client() @@ -72,6 +76,18 @@ func (c *NodeDrainCommand) Run(args []string) int { return 1 } + // If -self flag is set then determine the current node. + nodeID := "" + if !self { + nodeID = args[0] + } else { + var err error + if nodeID, err = getLocalNodeID(client); err != nil { + c.Ui.Error(err.Error()) + return 1 + } + } + // Check if node exists if len(nodeID) == 1 { c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters.")) @@ -83,7 +99,6 @@ func (c *NodeDrainCommand) Run(args []string) int { nodeID = nodeID[:len(nodeID)-1] } - // Exact lookup failed, try with prefix based search nodes, _, err := client.Nodes().PrefixList(nodeID) if err != nil { c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) diff --git a/command/node_status.go b/command/node_status.go index e15eee8edbb0..66e14335fdad 100644 --- a/command/node_status.go +++ b/command/node_status.go @@ -85,30 +85,6 @@ func (c *NodeStatusCommand) Run(args []string) int { return 1 } - // If -self flag is set then determine the current node. - nodeID := "" - if self { - info, err := client.Agent().Self() - if err != nil { - c.Ui.Error(fmt.Sprintf("Error querying agent info: %s", err)) - return 1 - } - var stats map[string]interface{} - stats, _ = info["stats"] - clientStats, ok := stats["client"].(map[string]interface{}) - if !ok { - c.Ui.Error("Nomad not running in client mode") - return 1 - } - - nodeID, ok = clientStats["node_id"].(string) - if !ok { - c.Ui.Error("Failed to determine node ID") - return 1 - } - - } - // Use list mode if no node name was provided if len(args) == 0 && !self { // Query the node info @@ -162,8 +138,15 @@ func (c *NodeStatusCommand) Run(args []string) int { } // Query the specific node + nodeID := "" if !self { nodeID = args[0] + } else { + var err error + if nodeID, err = getLocalNodeID(client); err != nil { + c.Ui.Error(err.Error()) + return 1 + } } if len(nodeID) == 1 { c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters.")) @@ -175,7 +158,6 @@ func (c *NodeStatusCommand) Run(args []string) int { nodeID = nodeID[:len(nodeID)-1] } - // Exact lookup failed, try with prefix based search nodes, _, err := client.Nodes().PrefixList(nodeID) if err != nil { c.Ui.Error(fmt.Sprintf("Error querying node info: %s", err)) diff --git a/website/source/docs/commands/node-drain.html.md.erb b/website/source/docs/commands/node-drain.html.md.erb index 919989f4261b..64096e55b624 100644 --- a/website/source/docs/commands/node-drain.html.md.erb +++ b/website/source/docs/commands/node-drain.html.md.erb @@ -21,9 +21,10 @@ nicely by providing the current drain status of a given node. nomad node-drain [options] ``` -A node ID or prefix must be provided. If there is an exact match, the -drain mode will be adjusted for that node. Otherwise, a list of matching -nodes and information will be displayed. +A `-self` flag can be used to drain the local node. If this is not supplied, a +node ID or prefix must be provided. If there is an exact match, the drain mode +will be adjusted for that node. Otherwise, a list of matching nodes and +information will be displayed. It is also required to pass one of `-enable` or `-disable`, depending on which operation is desired. @@ -36,11 +37,19 @@ operation is desired. * `-enable`: Enable node drain mode. * `-disable`: Disable node drain mode. +* `-self`: Drain the local node. +* `-yes`: Automtic yes to prompts. ## Examples -Enable drain mode on node1: +Enable drain mode on node with ID prefix "4d2ba53b": ``` -$ nomad node-drain -enable node1 +$ nomad node-drain -enable 4d2ba53b +``` + +Enable drain mode on the local node: + +``` +$ nomad node-drain -enable -self ```