github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/command/node_drain.go (about) 1 package command 2 3 import ( 4 "fmt" 5 "strings" 6 7 "github.com/hashicorp/nomad/api/contexts" 8 "github.com/posener/complete" 9 ) 10 11 type NodeDrainCommand struct { 12 Meta 13 } 14 15 func (c *NodeDrainCommand) Help() string { 16 helpText := ` 17 Usage: nomad node-drain [options] <node> 18 19 Toggles node draining on a specified node. It is required 20 that either -enable or -disable is specified, but not both. 21 The -self flag is useful to drain the local node. 22 23 General Options: 24 25 ` + generalOptionsUsage() + ` 26 27 Node Drain Options: 28 29 -disable 30 Disable draining for the specified node. 31 32 -enable 33 Enable draining for the specified node. 34 35 -self 36 Query the status of the local node. 37 38 -yes 39 Automatic yes to prompts. 40 ` 41 return strings.TrimSpace(helpText) 42 } 43 44 func (c *NodeDrainCommand) Synopsis() string { 45 return "Toggle drain mode on a given node" 46 } 47 48 func (c *NodeDrainCommand) AutocompleteFlags() complete.Flags { 49 return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), 50 complete.Flags{ 51 "-disable": complete.PredictNothing, 52 "-enable": complete.PredictNothing, 53 "-self": complete.PredictNothing, 54 "-yes": complete.PredictNothing, 55 }) 56 } 57 58 func (c *NodeDrainCommand) AutocompleteArgs() complete.Predictor { 59 return complete.PredictFunc(func(a complete.Args) []string { 60 client, err := c.Meta.Client() 61 if err != nil { 62 return nil 63 } 64 65 resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Nodes, nil) 66 if err != nil { 67 return []string{} 68 } 69 return resp.Matches[contexts.Nodes] 70 }) 71 } 72 73 func (c *NodeDrainCommand) Run(args []string) int { 74 var enable, disable, self, autoYes bool 75 76 flags := c.Meta.FlagSet("node-drain", FlagSetClient) 77 flags.Usage = func() { c.Ui.Output(c.Help()) } 78 flags.BoolVar(&enable, "enable", false, "Enable drain mode") 79 flags.BoolVar(&disable, "disable", false, "Disable drain mode") 80 flags.BoolVar(&self, "self", false, "") 81 flags.BoolVar(&autoYes, "yes", false, "Automatic yes to prompts.") 82 83 if err := flags.Parse(args); err != nil { 84 return 1 85 } 86 87 // Check that we got either enable or disable, but not both. 88 if (enable && disable) || (!enable && !disable) { 89 c.Ui.Error(c.Help()) 90 return 1 91 } 92 93 // Check that we got a node ID 94 args = flags.Args() 95 if l := len(args); self && l != 0 || !self && l != 1 { 96 c.Ui.Error(c.Help()) 97 return 1 98 } 99 100 // Get the HTTP client 101 client, err := c.Meta.Client() 102 if err != nil { 103 c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err)) 104 return 1 105 } 106 107 // If -self flag is set then determine the current node. 108 var nodeID string 109 if !self { 110 nodeID = args[0] 111 } else { 112 var err error 113 if nodeID, err = getLocalNodeID(client); err != nil { 114 c.Ui.Error(err.Error()) 115 return 1 116 } 117 } 118 119 // Check if node exists 120 if len(nodeID) == 1 { 121 c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters.")) 122 return 1 123 } 124 125 nodeID = sanatizeUUIDPrefix(nodeID) 126 nodes, _, err := client.Nodes().PrefixList(nodeID) 127 if err != nil { 128 c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) 129 return 1 130 } 131 // Return error if no nodes are found 132 if len(nodes) == 0 { 133 c.Ui.Error(fmt.Sprintf("No node(s) with prefix or id %q found", nodeID)) 134 return 1 135 } 136 if len(nodes) > 1 { 137 // Format the nodes list that matches the prefix so that the user 138 // can create a more specific request 139 out := make([]string, len(nodes)+1) 140 out[0] = "ID|Datacenter|Name|Class|Drain|Status" 141 for i, node := range nodes { 142 out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%v|%s", 143 node.ID, 144 node.Datacenter, 145 node.Name, 146 node.NodeClass, 147 node.Drain, 148 node.Status) 149 } 150 // Dump the output 151 c.Ui.Error(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out))) 152 return 1 153 } 154 155 // Prefix lookup matched a single node 156 node, _, err := client.Nodes().Info(nodes[0].ID, nil) 157 if err != nil { 158 c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) 159 return 1 160 } 161 162 // Confirm drain if the node was a prefix match. 163 if nodeID != node.ID && !autoYes { 164 verb := "enable" 165 if disable { 166 verb = "disable" 167 } 168 question := fmt.Sprintf("Are you sure you want to %s drain mode for node %q? [y/N]", verb, node.ID) 169 answer, err := c.Ui.Ask(question) 170 if err != nil { 171 c.Ui.Error(fmt.Sprintf("Failed to parse answer: %v", err)) 172 return 1 173 } 174 175 if answer == "" || strings.ToLower(answer)[0] == 'n' { 176 // No case 177 c.Ui.Output("Canceling drain toggle") 178 return 0 179 } else if strings.ToLower(answer)[0] == 'y' && len(answer) > 1 { 180 // Non exact match yes 181 c.Ui.Output("For confirmation, an exact ‘y’ is required.") 182 return 0 183 } else if answer != "y" { 184 c.Ui.Output("No confirmation detected. For confirmation, an exact 'y' is required.") 185 return 1 186 } 187 } 188 189 // Toggle node draining 190 if _, err := client.Nodes().ToggleDrain(node.ID, enable, nil); err != nil { 191 c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) 192 return 1 193 } 194 return 0 195 }