luxolus
/
open-nomad
2
0
Fork 0
Code Issues 1 Pull Requests Packages Releases Wiki Activity

Refactor alloc exec command (#9718) 

This re-arranges the alloc exec Run implementation to have validation
hoisted as high as possible.
This commit is contained in:
Kris Hicks 2021-01-05 09:33:04 -08:00 committed by GitHub
parent af5fd4bb89
commit be6580bd28
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 111 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

111
command/alloc_exec.go
View File

@ -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
}
}
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)
} 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 {
// Format the allocs
out := formatAllocListStubs(allocs, false, length)
out := formatAllocListStubs(allocs, false, shortId)
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)
allocStub = allocs[0]
}
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 {
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

48
command/alloc_exec_test.go
View File

@ -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`,
},
}

22
command/alloc_fs.go
View File

@ -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
}

2
command/alloc_logs.go
View File

@ -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