open-nomad/command/quota_apply.go
Alex Dadgar c1cc51dbee sync
2017-10-13 14:36:02 -07:00

277 lines
6 KiB
Go

package command
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"os"
"strings"
multierror "github.com/hashicorp/go-multierror"
"github.com/hashicorp/hcl"
"github.com/hashicorp/hcl/hcl/ast"
"github.com/hashicorp/nomad/api"
"github.com/hashicorp/nomad/helper"
"github.com/mitchellh/mapstructure"
"github.com/posener/complete"
)
type QuotaApplyCommand struct {
Meta
}
func (c *QuotaApplyCommand) Help() string {
helpText := `
Usage: nomad quota apply [options] <input>
Apply is used to create or update a quota specification. The specification file
will be read from stdin by specifying "-", otherwise a path to the file is
expected.
General Options:
` + generalOptionsUsage() + `
Apply Options:
-json
Parse the input as a JSON quota specification.
`
return strings.TrimSpace(helpText)
}
func (c *QuotaApplyCommand) AutocompleteFlags() complete.Flags {
return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
complete.Flags{
"-json": complete.PredictNothing,
})
}
func (c *QuotaApplyCommand) AutocompleteArgs() complete.Predictor {
return complete.PredictFiles("*")
}
func (c *QuotaApplyCommand) Synopsis() string {
return "Create or update a quota specification"
}
func (c *QuotaApplyCommand) Run(args []string) int {
var jsonInput bool
flags := c.Meta.FlagSet("quota apply", FlagSetClient)
flags.Usage = func() { c.Ui.Output(c.Help()) }
flags.BoolVar(&jsonInput, "json", false, "")
if err := flags.Parse(args); err != nil {
return 1
}
// Check that we get exactly one argument
args = flags.Args()
if l := len(args); l != 1 {
c.Ui.Error(c.Help())
return 1
}
// Read the file contents
file := args[0]
var rawQuota []byte
var err error
if file == "-" {
rawQuota, err = ioutil.ReadAll(os.Stdin)
if err != nil {
c.Ui.Error(fmt.Sprintf("Failed to read stdin: %v", err))
return 1
}
} else {
rawQuota, err = ioutil.ReadFile(file)
if err != nil {
c.Ui.Error(fmt.Sprintf("Failed to read file: %v", err))
return 1
}
}
var spec *api.QuotaSpec
if jsonInput {
var jsonSpec api.QuotaSpec
dec := json.NewDecoder(bytes.NewBuffer(rawQuota))
if err := dec.Decode(&jsonSpec); err != nil {
c.Ui.Error(fmt.Sprintf("Failed to parse quota: %v", err))
return 1
}
spec = &jsonSpec
} else {
hclSpec, err := parseQuotaSpec(rawQuota)
if err != nil {
c.Ui.Error(fmt.Sprintf("Error parsing quota specification: %s", err))
return 1
}
spec = hclSpec
}
// Get the HTTP client
client, err := c.Meta.Client()
if err != nil {
c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
return 1
}
_, err = client.Quotas().Register(spec, nil)
if err != nil {
c.Ui.Error(fmt.Sprintf("Error applying quota specification: %s", err))
return 1
}
c.Ui.Output(fmt.Sprintf("Successfully applied quota specification %q!", spec.Name))
return 0
}
// parseQuotaSpec is used to parse the quota specification from HCL
func parseQuotaSpec(input []byte) (*api.QuotaSpec, error) {
root, err := hcl.ParseBytes(input)
if err != nil {
return nil, err
}
// Top-level item should be a list
list, ok := root.Node.(*ast.ObjectList)
if !ok {
return nil, fmt.Errorf("error parsing: root should be an object")
}
var spec api.QuotaSpec
if err := parseQuotaSpecImpl(&spec, list); err != nil {
return nil, err
}
return &spec, nil
}
// parseQuotaSpecImpl parses the quota spec taking as input the AST tree
func parseQuotaSpecImpl(result *api.QuotaSpec, list *ast.ObjectList) error {
// Check for invalid keys
valid := []string{
"name",
"description",
"limit",
}
if err := helper.CheckHCLKeys(list, valid); err != nil {
return err
}
// Decode the full thing into a map[string]interface for ease
var m map[string]interface{}
if err := hcl.DecodeObject(&m, list); err != nil {
return err
}
// Manually parse
delete(m, "limit")
// Decode the rest
if err := mapstructure.WeakDecode(m, result); err != nil {
return err
}
// Parse limits
if o := list.Filter("limit"); len(o.Items) > 0 {
if err := parseQuotaLimits(&result.Limits, o); err != nil {
return multierror.Prefix(err, "limit ->")
}
}
return nil
}
// parseQuotaLimits parses the quota limits
func parseQuotaLimits(result *[]*api.QuotaLimit, list *ast.ObjectList) error {
for _, o := range list.Elem().Items {
// Check for invalid keys
valid := []string{
"region",
"region_limit",
}
if err := helper.CheckHCLKeys(o.Val, valid); err != nil {
return err
}
var m map[string]interface{}
if err := hcl.DecodeObject(&m, o.Val); err != nil {
return err
}
// Manually parse
delete(m, "region_limit")
// Decode the rest
var limit api.QuotaLimit
if err := mapstructure.WeakDecode(m, &limit); err != nil {
return err
}
// We need this later
var listVal *ast.ObjectList
if ot, ok := o.Val.(*ast.ObjectType); ok {
listVal = ot.List
} else {
return fmt.Errorf("limit should be an object")
}
// Parse limits
if o := listVal.Filter("region_limit"); len(o.Items) > 0 {
limit.RegionLimit = new(api.Resources)
if err := parseQuotaResource(limit.RegionLimit, o); err != nil {
return multierror.Prefix(err, "region_limit ->")
}
}
*result = append(*result, &limit)
}
return nil
}
// parseQuotaResource parses the region_limit resources
func parseQuotaResource(result *api.Resources, list *ast.ObjectList) error {
list = list.Elem()
if len(list.Items) == 0 {
return nil
}
if len(list.Items) > 1 {
return fmt.Errorf("only one 'region_limit' block allowed per limit")
}
// Get our resource object
o := list.Items[0]
// We need this later
var listVal *ast.ObjectList
if ot, ok := o.Val.(*ast.ObjectType); ok {
listVal = ot.List
} else {
return fmt.Errorf("resource: should be an object")
}
// Check for invalid keys
valid := []string{
"cpu",
"memory",
}
if err := helper.CheckHCLKeys(listVal, valid); err != nil {
return multierror.Prefix(err, "resources ->")
}
var m map[string]interface{}
if err := hcl.DecodeObject(&m, o.Val); err != nil {
return err
}
if err := mapstructure.WeakDecode(m, result); err != nil {
return err
}
return nil
}