172 lines
4.2 KiB
Go
172 lines
4.2 KiB
Go
![]() |
package command
|
|||
|
|
|||
|
import (
|
|||
|
"fmt"
|
|||
|
"strings"
|
|||
|
)
|
|||
|
|
|||
|
type NodeDrainCommand struct {
|
|||
|
Meta
|
|||
|
}
|
|||
|
|
|||
|
func (c *NodeDrainCommand) Help() string {
|
|||
|
helpText := `
|
|||
|
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:
|
|||
|
|
|||
|
` + generalOptionsUsage() + `
|
|||
|
|
|||
|
Node Drain Options:
|
|||
|
|
|||
|
-disable
|
|||
|
Disable draining for the specified node.
|
|||
|
|
|||
|
-enable
|
|||
|
Enable draining for the specified node.
|
|||
|
|
|||
|
-self
|
|||
|
Query the status of the local node.
|
|||
|
|
|||
|
-yes
|
|||
|
Automatic yes to prompts.
|
|||
|
`
|
|||
|
return strings.TrimSpace(helpText)
|
|||
|
}
|
|||
|
|
|||
|
func (c *NodeDrainCommand) Synopsis() string {
|
|||
|
return "Toggle drain mode on a given node"
|
|||
|
}
|
|||
|
|
|||
|
func (c *NodeDrainCommand) Run(args []string) int {
|
|||
|
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 {
|
|||
|
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(c.Help())
|
|||
|
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.
|
|||
|
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."))
|
|||
|
return 1
|
|||
|
}
|
|||
|
if len(nodeID)%2 == 1 {
|
|||
|
// Identifiers must be of even length, so we strip off the last byte
|
|||
|
// to provide a consistent user experience.
|
|||
|
nodeID = nodeID[:len(nodeID)-1]
|
|||
|
}
|
|||
|
|
|||
|
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.Output(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out)))
|
|||
|
return 0
|
|||
|
}
|
|||
|
|
|||
|
// 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
|
|||
|
}
|
|||
|
|
|||
|
// Confirm drain if the node was a prefix match.
|
|||
|
if nodeID != node.ID && !autoYes {
|
|||
|
verb := "enable"
|
|||
|
if disable {
|
|||
|
verb = "disable"
|
|||
|
}
|
|||
|
question := fmt.Sprintf("Are you sure you want to %s drain mode for node %q? [y/N]", verb, node.ID)
|
|||
|
answer, err := c.Ui.Ask(question)
|
|||
|
if err != nil {
|
|||
|
c.Ui.Error(fmt.Sprintf("Failed to parse answer: %v", err))
|
|||
|
return 1
|
|||
|
}
|
|||
|
|
|||
|
if answer == "" || strings.ToLower(answer)[0] == 'n' {
|
|||
|
// No case
|
|||
|
c.Ui.Output("Canceling drain toggle")
|
|||
|
return 0
|
|||
|
} else if strings.ToLower(answer)[0] == 'y' && len(answer) > 1 {
|
|||
|
// Non exact match yes
|
|||
|
c.Ui.Output("For confirmation, an exact ‘y’ is required.")
|
|||
|
return 0
|
|||
|
} else if answer != "y" {
|
|||
|
c.Ui.Output("No confirmation detected. For confirmation, an exact 'y' is required.")
|
|||
|
return 1
|
|||
|
}
|
|||
|
}
|
|||
|
|
|||
|
// Toggle node draining
|
|||
|
if _, err := client.Nodes().ToggleDrain(node.ID, enable, nil); err != nil {
|
|||
|
c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
|
|||
|
return 1
|
|||
|
}
|
|||
|
return 0
|
|||
|
}
|