github.com/ryanslade/nomad@v0.2.4-0.20160128061903-fc95782f2089/command/node_drain.go (about)

     1  package command
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  )
     7  
     8  type NodeDrainCommand struct {
     9  	Meta
    10  }
    11  
    12  func (c *NodeDrainCommand) Help() string {
    13  	helpText := `
    14  Usage: nomad node-drain [options] <node>
    15  
    16    Toggles node draining on a specified node. It is required
    17    that either -enable or -disable is specified, but not both.
    18  
    19  General Options:
    20  
    21    ` + generalOptionsUsage() + `
    22  
    23  Node Drain Options:
    24  
    25    -disable
    26      Disable draining for the specified node.
    27  
    28    -enable
    29      Enable draining for the specified node.
    30  `
    31  	return strings.TrimSpace(helpText)
    32  }
    33  
    34  func (c *NodeDrainCommand) Synopsis() string {
    35  	return "Toggle drain mode on a given node"
    36  }
    37  
    38  func (c *NodeDrainCommand) Run(args []string) int {
    39  	var enable, disable bool
    40  
    41  	flags := c.Meta.FlagSet("node-drain", FlagSetClient)
    42  	flags.Usage = func() { c.Ui.Output(c.Help()) }
    43  	flags.BoolVar(&enable, "enable", false, "Enable drain mode")
    44  	flags.BoolVar(&disable, "disable", false, "Disable drain mode")
    45  
    46  	if err := flags.Parse(args); err != nil {
    47  		return 1
    48  	}
    49  
    50  	// Check that we got either enable or disable, but not both.
    51  	if (enable && disable) || (!enable && !disable) {
    52  		c.Ui.Error(c.Help())
    53  		return 1
    54  	}
    55  
    56  	// Check that we got a node ID
    57  	args = flags.Args()
    58  	if len(args) != 1 {
    59  		c.Ui.Error(c.Help())
    60  		return 1
    61  	}
    62  	nodeID := args[0]
    63  
    64  	// Get the HTTP client
    65  	client, err := c.Meta.Client()
    66  	if err != nil {
    67  		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
    68  		return 1
    69  	}
    70  
    71  	// Check if node exists
    72  	node, _, err := client.Nodes().Info(nodeID, nil)
    73  	if err != nil {
    74  		if len(nodeID) == 1 {
    75  			c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters."))
    76  			return 1
    77  		}
    78  		if len(nodeID)%2 == 1 {
    79  			// Identifiers must be of even length, so we strip off the last byte
    80  			// to provide a consistent user experience.
    81  			nodeID = nodeID[:len(nodeID)-1]
    82  		}
    83  
    84  		// Exact lookup failed, try with prefix based search
    85  		nodes, _, err := client.Nodes().PrefixList(nodeID)
    86  		if err != nil {
    87  			c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
    88  			return 1
    89  		}
    90  		// Return error if no nodes are found
    91  		if len(nodes) == 0 {
    92  			c.Ui.Error(fmt.Sprintf("No node(s) with prefix or id %q found", nodeID))
    93  			return 1
    94  		}
    95  		if len(nodes) > 1 {
    96  			// Format the nodes list that matches the prefix so that the user
    97  			// can create a more specific request
    98  			out := make([]string, len(nodes)+1)
    99  			out[0] = "ID|Datacenter|Name|Class|Drain|Status"
   100  			for i, node := range nodes {
   101  				out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%v|%s",
   102  					node.ID,
   103  					node.Datacenter,
   104  					node.Name,
   105  					node.NodeClass,
   106  					node.Drain,
   107  					node.Status)
   108  			}
   109  			// Dump the output
   110  			c.Ui.Output(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out)))
   111  			return 0
   112  		}
   113  		// Prefix lookup matched a single node
   114  		node, _, err = client.Nodes().Info(nodes[0].ID, nil)
   115  		if err != nil {
   116  			c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
   117  			return 1
   118  		}
   119  	}
   120  
   121  	// Toggle node draining
   122  	if _, err := client.Nodes().ToggleDrain(node.ID, enable, nil); err != nil {
   123  		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
   124  		return 1
   125  	}
   126  	return 0
   127  }