Add -self flag to node-drain

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
+}

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")
+	}
+}

command/node_drain.go

@@ -15,6 +15,7 @@ Usage: nomad node-drain [options] <node>
 
   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))

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))

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] <node>
 ```
 
-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
 ```