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