improve wander exec command

cmd/exec.go

@@ -9,23 +9,41 @@ import (
 
 var (
 	execCmd = &cobra.Command{
-		Use:               "exec",
-		Short:             "Exec into a running task",
-		Long:              `Exec into a running nomad task`,
+		Use:   "exec",
+		Short: "Exec into a running task",
+		Long:  `Exec into a running nomad task`,
+		Example: `
+  # specify job and task, assuming single allocation
+  wander exec alright_stop --task redis echo "hi" 
+
+  # specify allocation, assuming single task
+  wander exec 3dca0982 echo "hi" 
+
+  # use prefixes of jobs or allocation ids
+  wander exec al echo "hi"
+  wander exec 3d echo "hi"
+
+  # specify flags for the exec command with --
+  wander exec alright_stop --task redis -- echo -n "hi"
+`,
 		Run:               execEntrypoint,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error { return nil },
 	}
 )
 
 func execEntrypoint(cmd *cobra.Command, args []string) {
+	task := cmd.Flags().Lookup("task").Value.String()
 	client, err := getConfig(cmd, "").Client()
 	if err != nil {
 		fmt.Println(fmt.Errorf("could not get client: %v", err))
 		os.Exit(1)
 	}
 	allocID := args[0]
-	task := args[1]
-	execArgs := args[2:]
+	execArgs := args[1:]
+	if len(execArgs) == 0 {
+		fmt.Println("no command specified")
+		os.Exit(1)
+	}
 	_, err = nomad.AllocExec(client, allocID, task, execArgs)
 	if err != nil {
 		fmt.Println(fmt.Errorf("could not exec into task: %v", err))

cmd/root.go

@@ -268,6 +268,9 @@ func init() {
 		viper.BindPFlag(cliLong, serveCmd.PersistentFlags().Lookup(c.cfgFileEnvVar))
 	}
 
+	// exec
+	execCmd.PersistentFlags().StringP("task", "", "", "Sets the task to exec command in")
+
 	rootCmd.AddCommand(serveCmd)
 	rootCmd.AddCommand(execCmd)
 }

internal/tui/components/app/app.go

@@ -266,7 +266,7 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 
 	case message.PageInputReceivedMsg:
 		if m.currentPage == nomad.ExecPage {
-			c := exec.Command("wander", "exec", m.alloc.ID, m.taskName, msg.Input)
+			c := exec.Command("wander", "exec", m.alloc.ID, "--task", m.taskName, msg.Input)
 			stdoutProxy := &nomad.StdoutProxy{}
 			c.Stdout = stdoutProxy
 			m.getCurrentPageModel().SetDoesNeedNewInput()
@@ -622,7 +622,7 @@ func (m Model) getCurrentPageCmd() tea.Cmd {
 				}
 				allPageRows = append(allPageRows, page.Row{Row: formatter.StripOSCommandSequences(formatter.StripANSI(row))})
 			}
-			return nomad.PageLoadedMsg{Page: nomad.ExecCompletePage, TableHeader: []string{"Output"}, AllPageRows: allPageRows}
+			return nomad.PageLoadedMsg{Page: nomad.ExecCompletePage, TableHeader: []string{"Exec Session Output"}, AllPageRows: allPageRows}
 		}
 	case nomad.AllocSpecPage:
 		return nomad.FetchAllocSpec(m.client, m.alloc.ID)

internal/tui/nomad/exec.go

@@ -7,6 +7,7 @@ import (
 	"github.com/hashicorp/nomad/api"
 	"github.com/moby/term"
 	"github.com/robinovitch61/wander/internal/tui/formatter"
+	"golang.org/x/exp/maps"
 	"golang.org/x/sys/unix"
 	"io"
 	"os"
@@ -15,14 +16,15 @@ import (
 )
 
 // TODO:
-// - [ ] code review this
+// - [x] try with bash
+// - [x] improve help text on wander exec
 // - [x] clear screen on start of exec
 // - [x] don't print config file used on exec
+// - [x] warning message that "you're in wander"
+// - [x] cmd human friendly (not full 36char id)
+// - [x] code review this
 // - [ ] update gif
-// - [ ] cmd human friendly (not full 36char id)
-// - [ ] warning message that "you're in wander"
-// - [ ] try with bash
-// - [ ] improve help text on wander exec
+// - [ ] fix wander serve exec
 
 type ExecCompleteMsg struct {
 	Output string
@@ -37,15 +39,79 @@ func (so *StdoutProxy) Write(p []byte) (n int, err error) {
 	return os.Stdout.Write(p)
 }
 
+func findAllocsForJobPrefix(client api.Client, jobName string) map[string][]*api.AllocationListStub {
+	allocs := make(map[string][]*api.AllocationListStub)
+	jobs, _, err := client.Jobs().PrefixList(jobName)
+	if err != nil || len(jobs) == 0 {
+		return allocs
+	}
+	for _, job := range jobs {
+		jobAllocs, _, err := client.Jobs().Allocations(job.ID, true, &api.QueryOptions{Namespace: job.Namespace})
+		if err != nil || len(jobAllocs) == 0 {
+			continue
+		}
+		allocs[job.ID] = jobAllocs
+	}
+	return allocs
+}
+
 func AllocExec(client *api.Client, allocID, task string, args []string) (int, error) {
 	alloc, _, err := client.Allocations().Info(allocID, nil)
 	if err != nil {
-		return 1, fmt.Errorf("error querying allocation: %s", err)
+		// maybe allocID is actually a job name
+		foundAllocs := findAllocsForJobPrefix(*client, allocID)
+		if len(foundAllocs) > 0 {
+			if len(foundAllocs) == 1 && len(maps.Values(foundAllocs)[0]) == 1 && maps.Values(foundAllocs)[0][0] != nil {
+				// only one job with one allocation found, use that
+				alloc, _, err = client.Allocations().Info(maps.Values(foundAllocs)[0][0].ID, nil)
+			} else {
+				// multiple jobs and/or allocations found, print them and exit
+				for job, jobAllocs := range foundAllocs {
+					fmt.Printf("allocations for job %s:\n", job)
+					for _, alloc := range jobAllocs {
+						fmt.Printf("  %s (%s in %s)\n", formatter.ShortAllocID(alloc.ID), alloc.Name, alloc.Namespace)
+					}
+				}
+				return 1, nil
+			}
+		} else {
+			// maybe allocID is short form of id
+			shortIDAllocs, _, err := client.Allocations().List(&api.QueryOptions{Prefix: allocID})
+			if err != nil {
+				return 1, fmt.Errorf("no jobs or allocation id for %s found: %v", allocID, err)
+			}
+			if len(shortIDAllocs) > 1 {
+				// rare but possible that uuid prefixes match
+				fmt.Printf("prefix %s matched multiple allocations:\n", allocID)
+				for _, alloc := range shortIDAllocs {
+					fmt.Printf("  %s (%s in %s)\n", formatter.ShortAllocID(alloc.ID), alloc.Name, alloc.Namespace)
+				}
+				return 1, err
+			} else if len(shortIDAllocs) == 1 {
+				alloc, _, err = client.Allocations().Info(shortIDAllocs[0].ID, nil)
+			} else {
+				return 1, fmt.Errorf("no allocations found for alloc id %s", allocID)
+			}
+		}
+	}
+
+	// if task is blank, user has assumed that there is only one task in the allocation and wants to
+	// use that
+	if task == "" {
+		if len(alloc.TaskStates) == 1 {
+			for taskName := range alloc.TaskStates {
+				task = taskName
+			}
+		} else {
+			fmt.Printf("multiple tasks found in allocation %s (%s in %s)\n", formatter.ShortAllocID(alloc.ID), alloc.Name, alloc.Namespace)
+			for taskName := range alloc.TaskStates {
+				fmt.Printf("  %s\n", taskName)
+			}
+		}
 	}
 
 	code, err := execImpl(client, alloc, task, args, "~", os.Stdin, os.Stdout, os.Stderr)
 	if err != nil {
-		fmt.Printf("failed to exec into task: %v", err)
 		return 1, err
 	}
 	return code, nil

scripts/Dockerfile

@@ -0,0 +1,3 @@
+FROM chentex/random-logger:v1.0.1
+
+RUN apk add bash

scripts/generate.sh

@@ -2,6 +2,8 @@
 
 set -e
 
+docker build . -t random-logger:local
+
 nomad namespace apply -description "My super cool namespace" my-namespace
 rm -f example.nomad.hcl
 nomad job init -short
@@ -17,7 +19,7 @@ sed -i "" -E "s/memory = 256/memory = 32/g" example.nomad.hcl
 run_job() {
     sed -i "" -E "s/job \".*\"/job \"${1:-example}\"/" example.nomad.hcl
     sed -i "" -E "s/namespace = \".*\"/namespace = \"${2:-default}\"/" example.nomad.hcl
-    sed -i "" -E "s/image          = \"redis.*\"/image          = \"chentex\/random-logger:v1.0.1\"\\nargs = ["50", "100"]/" example.nomad.hcl
+    sed -i "" -E "s/image          = \"redis.*\"/image          = \"random-logger:local\"\\nargs = ["50", "100"]/" example.nomad.hcl
     cat example.nomad.hcl
     nomad job run -detach example.nomad.hcl
 }