cmd/ejobs: support for getting results

Add a results subcommand to fetch results of a job.

This won't work until we add the jobs/results endpoint to the worker,
in a forthcoming CL.

Change-Id: I213ec6d3020ee37ee19896c82b68366f8415f08e
Reviewed-on: https://go-review.googlesource.com/c/pkgsite-metrics/+/511757
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Maceo Thompson <maceothompson@google.com>

cmd/ejobs/main.go

diff --git a/cmd/ejobs/main.go b/cmd/ejobs/main.go
index b423569..7277cfa 100644
--- a/cmd/ejobs/main.go
+++ b/cmd/ejobs/main.go
@@ -50,6 +50,10 @@
 
 	waitFlagSet  = flag.NewFlagSet("wait", flag.ContinueOnError)
 	waitInterval = waitFlagSet.Duration("i", 0, "display updates at this interval")
+
+	resultsFlagSet = flag.NewFlagSet("results", flag.ContinueOnError)
+	force          = resultsFlagSet.Bool("f", false, "download even if unfinished")
+	outfile        = resultsFlagSet.String("o", "", "output filename")
 )
 
 var commands = []command{
@@ -62,12 +66,15 @@
 	{"cancel", "JOBID...",
 		"cancel the jobs",
 		doCancel, nil},
-	{"start", "-min [MIN_IMPORTERS] BINARY ARGS...",
+	{"start", "[-min MIN_IMPORTERS] BINARY ARGS...",
 		"start a job",
 		doStart, startFlagSet},
 	{"wait", "JOBID",
 		"do not exit until JOBID is done",
 		doWait, nil},
+	{"results", "[-f] [-o FILE.json] JOBID",
+		"download results as JSON",
+		doResults, nil},
 }
 
 type command struct {
@@ -219,7 +226,7 @@
 		if err != nil {
 			return err
 		}
-		done := job.NumSkipped + job.NumFailed + job.NumErrored + job.NumSucceeded
+		done := job.NumFinished()
 		if done >= job.NumEnqueued {
 			break
 		}
@@ -426,6 +433,44 @@
 	return dest.Close()
 }
 
+func doResults(ctx context.Context, args []string) (err error) {
+	fs := resultsFlagSet
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+	if fs.NArg() == 0 {
+		return errors.New("wrong number of args: want [-f] [-o FILE.json] JOB_ID")
+	}
+	jobID := fs.Arg(0)
+	ts, err := identityTokenSource(ctx)
+	if err != nil {
+		return err
+	}
+	job, err := requestJSON[jobs.Job](ctx, "jobs/describe?jobid="+jobID, ts)
+	if err != nil {
+		return err
+	}
+	done := job.NumFinished()
+	if !*force && done < job.NumEnqueued {
+		return fmt.Errorf("job not finished (%d/%d completed); use -f for partial results", done, job.NumEnqueued)
+	}
+	results, err := requestJSON[jobs.Results](ctx, "jobs/results?jobid="+jobID, ts)
+	if err != nil {
+		return err
+	}
+	out := os.Stdout
+	if *outfile != "" {
+		out, err = os.Create(*outfile)
+		if err != nil {
+			return err
+		}
+		defer func() { err = errors.Join(err, out.Close()) }()
+	}
+	enc := json.NewEncoder(out)
+	enc.SetIndent("", "\t")
+	return enc.Encode(results)
+}
+
 // requestJSON requests the path from the worker, then reads the returned body
 // and unmarshals it as JSON.
 func requestJSON[T any](ctx context.Context, path string, ts oauth2.TokenSource) (*T, error) {

internal/bigquery/bigquery.go

diff --git a/internal/bigquery/bigquery.go b/internal/bigquery/bigquery.go
index dc503eb..d5ff368 100644
--- a/internal/bigquery/bigquery.go
+++ b/internal/bigquery/bigquery.go
@@ -321,7 +321,7 @@
 }
 
 // PartitionQuery describes a query that returns one row for each distinct value
-// of the partition column in the given table.
+// of the partition columns in the given table.
 //
 // The selected row will be the first one according to the OrderBy clauses.
 //

internal/jobs/job.go

diff --git a/internal/jobs/job.go b/internal/jobs/job.go
index 2c0169c..5e09912 100644
--- a/internal/jobs/job.go
+++ b/internal/jobs/job.go
@@ -7,6 +7,8 @@
 
 import (
 	"time"
+
+	"golang.org/x/pkgsite-metrics/internal/analysis"
 )
 
 // A Job is a set of related scan tasks enqueued at the same time.
@@ -45,3 +47,14 @@
 func (j *Job) ID() string {
 	return j.User + "-" + j.StartedAt.In(time.UTC).Format(startTimeFormat)
 }
+
+func (j *Job) NumFinished() int {
+	return j.NumSkipped + j.NumFailed + j.NumErrored + j.NumSucceeded
+}
+
+// Results hold the results of a job.
+type Results struct {
+	JobID   string
+	Table   string // bigquery table containing results
+	Results []*analysis.Result
+}