github.com/anuvu/nomad@v0.8.7-atom1/command/quota_apply.go (about) 1 package command 2 3 import ( 4 "bytes" 5 "encoding/json" 6 "fmt" 7 "io/ioutil" 8 "os" 9 "strings" 10 11 multierror "github.com/hashicorp/go-multierror" 12 "github.com/hashicorp/hcl" 13 "github.com/hashicorp/hcl/hcl/ast" 14 "github.com/hashicorp/nomad/api" 15 "github.com/hashicorp/nomad/helper" 16 "github.com/mitchellh/mapstructure" 17 "github.com/posener/complete" 18 ) 19 20 type QuotaApplyCommand struct { 21 Meta 22 } 23 24 func (c *QuotaApplyCommand) Help() string { 25 helpText := ` 26 Usage: nomad quota apply [options] <input> 27 28 Apply is used to create or update a quota specification. The specification file 29 will be read from stdin by specifying "-", otherwise a path to the file is 30 expected. 31 32 General Options: 33 34 ` + generalOptionsUsage() + ` 35 36 Apply Options: 37 38 -json 39 Parse the input as a JSON quota specification. 40 ` 41 42 return strings.TrimSpace(helpText) 43 } 44 45 func (c *QuotaApplyCommand) AutocompleteFlags() complete.Flags { 46 return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), 47 complete.Flags{ 48 "-json": complete.PredictNothing, 49 }) 50 } 51 52 func (c *QuotaApplyCommand) AutocompleteArgs() complete.Predictor { 53 return complete.PredictFiles("*") 54 } 55 56 func (c *QuotaApplyCommand) Synopsis() string { 57 return "Create or update a quota specification" 58 } 59 60 func (c *QuotaApplyCommand) Name() string { return "quota apply" } 61 62 func (c *QuotaApplyCommand) Run(args []string) int { 63 var jsonInput bool 64 flags := c.Meta.FlagSet(c.Name(), FlagSetClient) 65 flags.Usage = func() { c.Ui.Output(c.Help()) } 66 flags.BoolVar(&jsonInput, "json", false, "") 67 68 if err := flags.Parse(args); err != nil { 69 return 1 70 } 71 72 // Check that we get exactly one argument 73 args = flags.Args() 74 if l := len(args); l != 1 { 75 c.Ui.Error("This command takes one argument: <input>") 76 c.Ui.Error(commandErrorText(c)) 77 return 1 78 } 79 80 // Read the file contents 81 file := args[0] 82 var rawQuota []byte 83 var err error 84 if file == "-" { 85 rawQuota, err = ioutil.ReadAll(os.Stdin) 86 if err != nil { 87 c.Ui.Error(fmt.Sprintf("Failed to read stdin: %v", err)) 88 return 1 89 } 90 } else { 91 rawQuota, err = ioutil.ReadFile(file) 92 if err != nil { 93 c.Ui.Error(fmt.Sprintf("Failed to read file: %v", err)) 94 return 1 95 } 96 } 97 98 var spec *api.QuotaSpec 99 if jsonInput { 100 var jsonSpec api.QuotaSpec 101 dec := json.NewDecoder(bytes.NewBuffer(rawQuota)) 102 if err := dec.Decode(&jsonSpec); err != nil { 103 c.Ui.Error(fmt.Sprintf("Failed to parse quota: %v", err)) 104 return 1 105 } 106 spec = &jsonSpec 107 } else { 108 hclSpec, err := parseQuotaSpec(rawQuota) 109 if err != nil { 110 c.Ui.Error(fmt.Sprintf("Error parsing quota specification: %s", err)) 111 return 1 112 } 113 114 spec = hclSpec 115 } 116 117 // Get the HTTP client 118 client, err := c.Meta.Client() 119 if err != nil { 120 c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err)) 121 return 1 122 } 123 124 _, err = client.Quotas().Register(spec, nil) 125 if err != nil { 126 c.Ui.Error(fmt.Sprintf("Error applying quota specification: %s", err)) 127 return 1 128 } 129 130 c.Ui.Output(fmt.Sprintf("Successfully applied quota specification %q!", spec.Name)) 131 return 0 132 } 133 134 // parseQuotaSpec is used to parse the quota specification from HCL 135 func parseQuotaSpec(input []byte) (*api.QuotaSpec, error) { 136 root, err := hcl.ParseBytes(input) 137 if err != nil { 138 return nil, err 139 } 140 141 // Top-level item should be a list 142 list, ok := root.Node.(*ast.ObjectList) 143 if !ok { 144 return nil, fmt.Errorf("error parsing: root should be an object") 145 } 146 147 var spec api.QuotaSpec 148 if err := parseQuotaSpecImpl(&spec, list); err != nil { 149 return nil, err 150 } 151 152 return &spec, nil 153 } 154 155 // parseQuotaSpecImpl parses the quota spec taking as input the AST tree 156 func parseQuotaSpecImpl(result *api.QuotaSpec, list *ast.ObjectList) error { 157 // Check for invalid keys 158 valid := []string{ 159 "name", 160 "description", 161 "limit", 162 } 163 if err := helper.CheckHCLKeys(list, valid); err != nil { 164 return err 165 } 166 167 // Decode the full thing into a map[string]interface for ease 168 var m map[string]interface{} 169 if err := hcl.DecodeObject(&m, list); err != nil { 170 return err 171 } 172 173 // Manually parse 174 delete(m, "limit") 175 176 // Decode the rest 177 if err := mapstructure.WeakDecode(m, result); err != nil { 178 return err 179 } 180 181 // Parse limits 182 if o := list.Filter("limit"); len(o.Items) > 0 { 183 if err := parseQuotaLimits(&result.Limits, o); err != nil { 184 return multierror.Prefix(err, "limit ->") 185 } 186 } 187 188 return nil 189 } 190 191 // parseQuotaLimits parses the quota limits 192 func parseQuotaLimits(result *[]*api.QuotaLimit, list *ast.ObjectList) error { 193 for _, o := range list.Elem().Items { 194 // Check for invalid keys 195 valid := []string{ 196 "region", 197 "region_limit", 198 } 199 if err := helper.CheckHCLKeys(o.Val, valid); err != nil { 200 return err 201 } 202 203 var m map[string]interface{} 204 if err := hcl.DecodeObject(&m, o.Val); err != nil { 205 return err 206 } 207 208 // Manually parse 209 delete(m, "region_limit") 210 211 // Decode the rest 212 var limit api.QuotaLimit 213 if err := mapstructure.WeakDecode(m, &limit); err != nil { 214 return err 215 } 216 217 // We need this later 218 var listVal *ast.ObjectList 219 if ot, ok := o.Val.(*ast.ObjectType); ok { 220 listVal = ot.List 221 } else { 222 return fmt.Errorf("limit should be an object") 223 } 224 225 // Parse limits 226 if o := listVal.Filter("region_limit"); len(o.Items) > 0 { 227 limit.RegionLimit = new(api.Resources) 228 if err := parseQuotaResource(limit.RegionLimit, o); err != nil { 229 return multierror.Prefix(err, "region_limit ->") 230 } 231 } 232 233 *result = append(*result, &limit) 234 } 235 236 return nil 237 } 238 239 // parseQuotaResource parses the region_limit resources 240 func parseQuotaResource(result *api.Resources, list *ast.ObjectList) error { 241 list = list.Elem() 242 if len(list.Items) == 0 { 243 return nil 244 } 245 if len(list.Items) > 1 { 246 return fmt.Errorf("only one 'region_limit' block allowed per limit") 247 } 248 249 // Get our resource object 250 o := list.Items[0] 251 252 // We need this later 253 var listVal *ast.ObjectList 254 if ot, ok := o.Val.(*ast.ObjectType); ok { 255 listVal = ot.List 256 } else { 257 return fmt.Errorf("resource: should be an object") 258 } 259 260 // Check for invalid keys 261 valid := []string{ 262 "cpu", 263 "memory", 264 } 265 if err := helper.CheckHCLKeys(listVal, valid); err != nil { 266 return multierror.Prefix(err, "resources ->") 267 } 268 269 var m map[string]interface{} 270 if err := hcl.DecodeObject(&m, o.Val); err != nil { 271 return err 272 } 273 274 if err := mapstructure.WeakDecode(m, result); err != nil { 275 return err 276 } 277 278 return nil 279 }