diff --git a/command/alloc_exec.go b/command/alloc_exec.go index cfeb91d05..9555ee03d 100644 --- a/command/alloc_exec.go +++ b/command/alloc_exec.go @@ -100,25 +100,43 @@ func (l *AllocExecCommand) Name() string { return "alloc exec" } func (l *AllocExecCommand) Run(args []string) int { var job, stdinOpt, ttyOpt bool - var task string - var escapeChar string + var task, escapeChar string flags := l.Meta.FlagSet(l.Name(), FlagSetClient) flags.Usage = func() { l.Ui.Output(l.Help()) } flags.BoolVar(&job, "job", false, "") - flags.StringVar(&task, "task", "", "") - flags.StringVar(&escapeChar, "e", "~", "") - flags.BoolVar(&stdinOpt, "i", true, "") - - inferredTty := isTty() - flags.BoolVar(&ttyOpt, "t", inferredTty, "") + flags.BoolVar(&ttyOpt, "t", isTty(), "") + flags.StringVar(&escapeChar, "e", "~", "") + flags.StringVar(&task, "task", "", "") if err := flags.Parse(args); err != nil { return 1 } + args = flags.Args() + if len(args) < 1 { + if job { + l.Ui.Error("A job ID is required") + } else { + l.Ui.Error("An allocation ID is required") + } + l.Ui.Error(commandErrorText(l)) + return 1 + } + + if !job && len(args[0]) == 1 { + l.Ui.Error("Alloc ID must contain at least two characters") + return 1 + } + + if len(args) < 2 { + l.Ui.Error("A command is required") + l.Ui.Error(commandErrorText(l)) + return 1 + } + if ttyOpt && !stdinOpt { l.Ui.Error("-i must be enabled if running with tty") return 1 @@ -126,106 +144,83 @@ func (l *AllocExecCommand) Run(args []string) int { if escapeChar == "none" { escapeChar = "" - } else if len(escapeChar) > 1 { + } + + if len(escapeChar) > 1 { l.Ui.Error("-e requires 'none' or a single character") return 1 } - if numArgs := len(args); numArgs < 1 { - if job { - l.Ui.Error("A job ID is required") - } else { - l.Ui.Error("An allocation ID is required") - } - - l.Ui.Error(commandErrorText(l)) - return 1 - } else if numArgs < 2 { - l.Ui.Error("A command is required") - l.Ui.Error(commandErrorText(l)) - return 1 - } - - command := args[1:] - client, err := l.Meta.Client() if err != nil { l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err)) return 1 } - // If -job is specified, use random allocation, otherwise use provided allocation - allocID := args[0] + var allocStub *api.AllocationListStub if job { - allocID, err = getRandomJobAlloc(client, args[0]) + jobID := args[0] + allocStub, err = getRandomJobAlloc(client, jobID) if err != nil { l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err)) return 1 } + } else { + allocID := args[0] + allocs, _, err := client.Allocations().PrefixList(sanitizeUUIDPrefix(allocID)) + if err != nil { + l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err)) + return 1 + } + + if len(allocs) == 0 { + l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID)) + return 1 + } + + if len(allocs) > 1 { + out := formatAllocListStubs(allocs, false, shortId) + l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out)) + return 1 + } + + allocStub = allocs[0] } - length := shortId - - // Query the allocation info - if len(allocID) == 1 { - l.Ui.Error("Alloc ID must contain at least two characters.") - return 1 - } - - allocID = sanitizeUUIDPrefix(allocID) - allocs, _, err := client.Allocations().PrefixList(allocID) - if err != nil { - l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err)) - return 1 - } - if len(allocs) == 0 { - l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID)) - return 1 - } - if len(allocs) > 1 { - // Format the allocs - out := formatAllocListStubs(allocs, false, length) - l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out)) - return 1 - } - // Prefix lookup matched a single allocation - q := &api.QueryOptions{Namespace: allocs[0].Namespace} - alloc, _, err := client.Allocations().Info(allocs[0].ID, q) + q := &api.QueryOptions{Namespace: allocStub.Namespace} + alloc, _, err := client.Allocations().Info(allocStub.ID, q) if err != nil { l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err)) return 1 } - if task == "" { + if task != "" { + err = validateTaskExistsInAllocation(task, alloc) + } else { task, err = lookupAllocTask(alloc) - - if err != nil { - l.Ui.Error(err.Error()) - return 1 - } } - - if err := validateTaskExistsInAllocation(task, alloc); err != nil { + if err != nil { l.Ui.Error(err.Error()) return 1 } + if !stdinOpt { + l.Stdin = bytes.NewReader(nil) + } + if l.Stdin == nil { l.Stdin = os.Stdin } + if l.Stdout == nil { l.Stdout = os.Stdout } + if l.Stderr == nil { l.Stderr = os.Stderr } - var stdin io.Reader = l.Stdin - if !stdinOpt { - stdin = bytes.NewReader(nil) - } - - code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr) + code, err := l.execImpl(client, alloc, task, ttyOpt, args[1:], escapeChar, l.Stdin, l.Stdout, l.Stderr) if err != nil { l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err)) return 1 diff --git a/command/alloc_exec_test.go b/command/alloc_exec_test.go index dd6864936..f064aa404 100644 --- a/command/alloc_exec_test.go +++ b/command/alloc_exec_test.go @@ -30,39 +30,49 @@ func TestAllocExecCommand_Fails(t *testing.T) { expectedError string }{ { - "misuse", - []string{"bad"}, - commandErrorText(&AllocExecCommand{}), + "alloc id missing", + []string{}, + `An allocation ID is required`, }, { - "connection failure", - []string{"-address=nope", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, - "Error querying allocation", + "alloc id too short", + []string{"-address=" + url, "2", "/bin/bash"}, + `Alloc ID must contain at least two characters`, }, { - "not found alloc", + "alloc not found", []string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, - "No allocation(s) with prefix or id", + `No allocation(s) with prefix or id "26470238-5CF2-438F-8772-DC67CFB0705C"`, }, { - "not found job", + "alloc not found with odd-length prefix", + []string{"-address=" + url, "26470238-5CF", "/bin/bash"}, + `No allocation(s) with prefix or id "26470238-5CF"`, + }, + { + "job id missing", + []string{"-job"}, + `A job ID is required`, + }, + { + "job not found", []string{"-address=" + url, "-job", "example", "/bin/bash"}, `job "example" doesn't exist`, }, { - "too short allocis", - []string{"-address=" + url, "2", "/bin/bash"}, - "Alloc ID must contain at least two characters", - }, - { - "missing command", + "command missing", []string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C"}, - "A command is required", + `A command is required`, }, { - "long escaped char", - []string{"-address=" + url, "-e", "long_escape", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, - "a single character", + "connection failure", + []string{"-address=nope", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, + `Error querying allocation`, + }, + { + "escape char too long", + []string{"-address=" + url, "-e", "es", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, + `-e requires 'none' or a single character`, }, } diff --git a/command/alloc_fs.go b/command/alloc_fs.go index 932d3d5e1..a10750eb2 100644 --- a/command/alloc_fs.go +++ b/command/alloc_fs.go @@ -164,7 +164,7 @@ func (f *AllocFSCommand) Run(args []string) int { // If -job is specified, use random allocation, otherwise use provided allocation allocID := args[0] if job { - allocID, err = getRandomJobAlloc(client, args[0]) + allocID, err = getRandomJobAllocID(client, args[0]) if err != nil { f.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err)) return 1 @@ -376,15 +376,15 @@ func (f *AllocFSCommand) followFile(client *api.Client, alloc *api.Allocation, return r, nil } -// Get Random Allocation ID from a known jobID. Prefer to use a running allocation, +// Get Random Allocation from a known jobID. Prefer to use a running allocation, // but use a dead allocation if no running allocations are found -func getRandomJobAlloc(client *api.Client, jobID string) (string, error) { +func getRandomJobAlloc(client *api.Client, jobID string) (*api.AllocationListStub, error) { var runningAllocs []*api.AllocationListStub allocs, _, err := client.Jobs().Allocations(jobID, false, nil) // Check that the job actually has allocations if len(allocs) == 0 { - return "", fmt.Errorf("job %q doesn't exist or it has no allocations", jobID) + return nil, fmt.Errorf("job %q doesn't exist or it has no allocations", jobID) } for _, v := range allocs { @@ -398,6 +398,16 @@ func getRandomJobAlloc(client *api.Client, jobID string) (string, error) { } r := rand.New(rand.NewSource(time.Now().UnixNano())) - allocID := runningAllocs[r.Intn(len(runningAllocs))].ID - return allocID, err + alloc := runningAllocs[r.Intn(len(runningAllocs))] + + return alloc, err +} + +func getRandomJobAllocID(client *api.Client, jobID string) (string, error) { + alloc, err := getRandomJobAlloc(client, jobID) + if err != nil { + return "", err + } + + return alloc.ID, nil } diff --git a/command/alloc_logs.go b/command/alloc_logs.go index 6ba94260c..a5feba45b 100644 --- a/command/alloc_logs.go +++ b/command/alloc_logs.go @@ -144,7 +144,7 @@ func (l *AllocLogsCommand) Run(args []string) int { // If -job is specified, use random allocation, otherwise use provided allocation allocID := args[0] if job { - allocID, err = getRandomJobAlloc(client, args[0]) + allocID, err = getRandomJobAllocID(client, args[0]) if err != nil { l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err)) return 1