node eligibility command

command/node_drain.go

@@ -269,5 +269,7 @@ func (c *NodeDrainCommand) Run(args []string) int {
 		c.Ui.Error(fmt.Sprintf("Error updating drain specification: %s", err))
 		return 1
 	}
+
+	c.Ui.Output(fmt.Sprintf("Node %q drain strategy set", node.ID))
 	return 0
 }

command/node_eligibility.go

@@ -0,0 +1,168 @@
+package command
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/hashicorp/nomad/api/contexts"
+	"github.com/posener/complete"
+)
+
+type NodeEligibilityCommand struct {
+	Meta
+}
+
+func (c *NodeEligibilityCommand) Help() string {
+	helpText := `
+Usage: nomad node eligibility [options] <node>
+
+  Toggles the nodes scheduling eligibility. When a node is marked as ineligible,
+  no new allocations will be placed on it but existing allocations will remain.
+  To remove existing allocations, use the node drain command.
+
+  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:
+
+  ` + generalOptionsUsage() + `
+
+Node Eligibility Options:
+
+  -disable
+    Mark the specified node as ineligible for new allocations.
+
+  -enable
+    Mark the specified node as eligible for new allocations.
+
+  -self
+    Set the eligibility of the local node.
+`
+	return strings.TrimSpace(helpText)
+}
+
+func (c *NodeEligibilityCommand) Synopsis() string {
+	return "Toggle scheduling eligibility for a given node"
+}
+
+func (c *NodeEligibilityCommand) AutocompleteFlags() complete.Flags {
+	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
+		complete.Flags{
+			"-disable": complete.PredictNothing,
+			"-enable":  complete.PredictNothing,
+			"-self":    complete.PredictNothing,
+		})
+}
+
+func (c *NodeEligibilityCommand) AutocompleteArgs() complete.Predictor {
+	return complete.PredictFunc(func(a complete.Args) []string {
+		client, err := c.Meta.Client()
+		if err != nil {
+			return nil
+		}
+
+		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Nodes, nil)
+		if err != nil {
+			return []string{}
+		}
+		return resp.Matches[contexts.Nodes]
+	})
+}
+
+func (c *NodeEligibilityCommand) Run(args []string) int {
+	var enable, disable, self bool
+
+	flags := c.Meta.FlagSet("node-eligibility", FlagSetClient)
+	flags.Usage = func() { c.Ui.Output(c.Help()) }
+	flags.BoolVar(&enable, "enable", false, "Mark node as eligibile for scheduling")
+	flags.BoolVar(&disable, "disable", false, "Mark node as ineligibile for scheduling")
+	flags.BoolVar(&self, "self", false, "")
+
+	if err := flags.Parse(args); err != nil {
+		return 1
+	}
+
+	// Check that we got either enable or disable, but not both.
+	if (enable && disable) || (!enable && !disable) {
+		c.Ui.Error(c.Help())
+		return 1
+	}
+
+	// Check that we got a node ID
+	args = flags.Args()
+	if l := len(args); self && l != 0 || !self && l != 1 {
+		c.Ui.Error("Node ID must be specified if -self isn't being used")
+		return 1
+	}
+
+	// Get the HTTP client
+	client, err := c.Meta.Client()
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
+		return 1
+	}
+
+	// If -self flag is set then determine the current node.
+	var nodeID string
+	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."))
+		return 1
+	}
+
+	nodeID = sanatizeUUIDPrefix(nodeID)
+	nodes, _, err := client.Nodes().PrefixList(nodeID)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
+		return 1
+	}
+	// Return error if no nodes are found
+	if len(nodes) == 0 {
+		c.Ui.Error(fmt.Sprintf("No node(s) with prefix or id %q found", nodeID))
+		return 1
+	}
+	if len(nodes) > 1 {
+		// Format the nodes list that matches the prefix so that the user
+		// can create a more specific request
+		out := make([]string, len(nodes)+1)
+		out[0] = "ID|Datacenter|Name|Class|Drain|Status"
+		for i, node := range nodes {
+			out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%v|%s",
+				node.ID,
+				node.Datacenter,
+				node.Name,
+				node.NodeClass,
+				node.Drain,
+				node.Status)
+		}
+		// Dump the output
+		c.Ui.Error(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out)))
+		return 1
+	}
+
+	// Prefix lookup matched a single node
+	node, _, err := client.Nodes().Info(nodes[0].ID, nil)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
+		return 1
+	}
+
+	// Toggle node eligibility
+	if _, err := client.Nodes().ToggleEligibility(node.ID, enable, nil); err != nil {
+		c.Ui.Error(fmt.Sprintf("Error updating scheduling eligibility: %s", err))
+		return 1
+	}
+
+	c.Ui.Output(fmt.Sprintf("Node %q scheduling eligibility set", node.ID))
+	return 0
+}

command/node_eligibility_test.go

@@ -0,0 +1,124 @@
+package command
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/hashicorp/nomad/testutil"
+	"github.com/mitchellh/cli"
+	"github.com/posener/complete"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNodeEligibilityCommand_Implements(t *testing.T) {
+	t.Parallel()
+	var _ cli.Command = &NodeEligibilityCommand{}
+}
+
+func TestNodeEligibilityCommand_Fails(t *testing.T) {
+	t.Parallel()
+	srv, _, url := testServer(t, false, nil)
+	defer srv.Shutdown()
+
+	ui := new(cli.MockUi)
+	cmd := &NodeEligibilityCommand{Meta: Meta{Ui: ui}}
+
+	// Fails on misuse
+	if code := cmd.Run([]string{"some", "bad", "args"}); code != 1 {
+		t.Fatalf("expected exit code 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) {
+		t.Fatalf("expected help output, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails on connection failure
+	if code := cmd.Run([]string{"-address=nope", "-enable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+		t.Fatalf("expected exit code 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error toggling") {
+		t.Fatalf("expected failed toggle error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails on non-existent node
+	if code := cmd.Run([]string{"-address=" + url, "-enable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "No node(s) with prefix or id") {
+		t.Fatalf("expected not exist error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails if both enable and disable specified
+	if code := cmd.Run([]string{"-enable", "-disable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) {
+		t.Fatalf("expected help output, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails if neither enable or disable specified
+	if code := cmd.Run([]string{"12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) {
+		t.Fatalf("expected help output, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fail on identifier with too few characters
+	if code := cmd.Run([]string{"-address=" + url, "-enable", "1"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "must contain at least two characters.") {
+		t.Fatalf("expected too few characters error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Identifiers with uneven length should produce a query result
+	if code := cmd.Run([]string{"-address=" + url, "-enable", "123"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "No node(s) with prefix or id") {
+		t.Fatalf("expected not exist error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+}
+
+func TestNodeEligibilityCommand_AutocompleteArgs(t *testing.T) {
+	assert := assert.New(t)
+	t.Parallel()
+
+	srv, client, url := testServer(t, true, nil)
+	defer srv.Shutdown()
+
+	// Wait for a node to appear
+	var nodeID string
+	testutil.WaitForResult(func() (bool, error) {
+		nodes, _, err := client.Nodes().List(nil)
+		if err != nil {
+			return false, err
+		}
+		if len(nodes) == 0 {
+			return false, fmt.Errorf("missing node")
+		}
+		nodeID = nodes[0].ID
+		return true, nil
+	}, func(err error) {
+		t.Fatalf("err: %s", err)
+	})
+
+	ui := new(cli.MockUi)
+	cmd := &NodeEligibilityCommand{Meta: Meta{Ui: ui, flagAddress: url}}
+
+	prefix := nodeID[:len(nodeID)-5]
+	args := complete.Args{Last: prefix}
+	predictor := cmd.AutocompleteArgs()
+
+	res := predictor.Predict(args)
+	assert.Equal(1, len(res))
+	assert.Equal(nodeID, res[0])
+}

commands.go

@@ -273,6 +273,11 @@ func Commands(metaPtr *command.Meta) map[string]cli.CommandFactory {
 				Meta: meta,
 			}, nil
 		},
+		"node eligibility": func() (cli.Command, error) {
+			return &command.NodeEligibilityCommand{
+				Meta: meta,
+			}, nil
+		},
 		"node-status": func() (cli.Command, error) {
 			return &command.NodeStatusCommand{
 				Meta: meta,

main.go

@@ -47,7 +47,7 @@ func RunCustom(args []string, commands map[string]cli.CommandFactory) int {
 	// users should not be running should be placed here, versus hiding
 	// subcommands from the main help, which should be filtered out of the
 	// commands above.
-	hidden := []string{"check", "executor", "syslog"}
+	hidden := []string{"check", "executor", "syslog", "node-drain", "node-status"}
 
 	cli := &cli.CLI{
 		Name:           "nomad",