internal/worker: add streaming of results to ejobs - Usage: ejobs results -e -f -stream <job-id> - Perform a BigQuery call from the ejobs client. Change-Id: I6d8f6bcc9e7e3a8586324bc7bf5211d6dd64c73f Reviewed-on: https://go-review.googlesource.com/c/pkgsite-metrics/+/708735 LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Markus Kusano <kusano@google.com> Auto-Submit: Ethan Lee <ethanalee@google.com>
diff --git a/cmd/ejobs/main.go b/cmd/ejobs/main.go index 80d2f79..7316aab 100644 --- a/cmd/ejobs/main.go +++ b/cmd/ejobs/main.go
@@ -26,17 +26,21 @@ "time" "unicode" + bq "cloud.google.com/go/bigquery" "cloud.google.com/go/storage" "golang.org/x/oauth2" "golang.org/x/pkgsite-metrics/internal/analysis" + "golang.org/x/pkgsite-metrics/internal/bigquery" "golang.org/x/pkgsite-metrics/internal/jobs" "google.golang.org/api/impersonate" + "google.golang.org/api/iterator" "google.golang.org/api/option" ) const ( projectID = "go-ecosystem" uploaderMetadataKey = "uploader" + pollingInterval = 10 * time.Second ) // Common flags @@ -54,6 +58,7 @@ force bool // for results errs bool // for results outfile string // for results + stream bool // for results userFilter string // for list ) @@ -98,6 +103,7 @@ fs.BoolVar(&force, "f", false, "download even if unfinished") fs.BoolVar(&errs, "e", false, "also download error results (by default, only non-error results are downloaded)") fs.StringVar(&outfile, "o", "", "output filename") + fs.BoolVar(&stream, "stream", false, "stream output") }, }, } @@ -544,6 +550,12 @@ return dest.Close() } +type bqClient interface { + QueryWithParams(ctx context.Context, query string, params []bq.QueryParameter) (bigquery.RowIterator, error) + FullTableName(tableID string) string + Close() error +} + func doResults(ctx context.Context, args []string) (err error) { if len(args) == 0 { return errors.New("wrong number of args: want [-f] [-e] [-o FILE.json] JOB_ID") @@ -561,10 +573,6 @@ 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[[]*analysis.Result](ctx, fmt.Sprintf("jobs/results?jobid=%s&errors=%t", jobID, errs), ts) - if err != nil { - return err - } out := os.Stdout if outfile != "" { out, err = os.Create(outfile) @@ -573,9 +581,109 @@ } defer func() { err = errors.Join(err, out.Close()) }() } - enc := json.NewEncoder(out) - enc.SetIndent("", "\t") - return enc.Encode(results) + if !stream { + results, err := requestJSON[[]*analysis.Result](ctx, fmt.Sprintf("jobs/results?jobid=%s&errors=%t", jobID, errs), ts) + if err != nil { + return err + } + enc := json.NewEncoder(out) + enc.SetIndent("", "\t") + return enc.Encode(results) + } + + bqClient, err := bigquery.NewClientCreate(ctx, projectID, *env) + if err != nil { + return fmt.Errorf("creating BigQuery client: %w", err) + } + defer bqClient.Close() + + var lastCreatedAt time.Time + var totalResultsCount int + processedIDs := make(map[string]struct{}) + + fmt.Fprintf(os.Stderr, "Streaming results for job %s...\n", jobID) + + for { + job, err := requestJSON[jobs.Job](ctx, "jobs/describe?jobid="+jobID, ts) + if err != nil { + return fmt.Errorf("could not get job status: %w", err) + } + + newResultsCount, err := fetchAndPrintResults(ctx, out, bqClient, jobID, errs, lastCreatedAt, processedIDs) + if err != nil { + return err + } + totalResultsCount += newResultsCount + + if job.NumFinished() >= job.NumEnqueued { + fmt.Fprintf(os.Stderr, "\nJob finished (%d/%d completed). Found %d results. Stream complete.\n", job.NumFinished(), job.NumEnqueued, totalResultsCount) + break + } + + fmt.Fprintf(os.Stderr, "\rJob still running (%d/%d)... Found %d results... waiting %v", job.NumFinished(), job.NumEnqueued, totalResultsCount, pollingInterval) + select { + case <-ctx.Done(): + fmt.Fprintln(os.Stderr, "\nStream canceled by user.") + return ctx.Err() + case <-time.After(pollingInterval): + } + } + return nil +} + +func fetchAndPrintResults(ctx context.Context, out io.Writer, bqClient bqClient, jobID string, errs bool, lastCreatedAt time.Time, processedIDs map[string]struct{}) (int, error) { + queryStr := fmt.Sprintf(`SELECT * FROM %s WHERE job_id = @jobID`, bqClient.FullTableName(analysis.TableName)) + + params := []bq.QueryParameter{ + {Name: "jobID", Value: jobID}, + } + + if !errs { + queryStr += " AND error = ''" + } + if !(lastCreatedAt).IsZero() { + queryStr += " AND created_at >= @minCreatedAt" + params = append(params, bq.QueryParameter{Name: "minCreatedAt", Value: lastCreatedAt}) + } + queryStr += " ORDER BY created_at ASC" + + iter, err := bqClient.QueryWithParams(ctx, queryStr, params) + if err != nil { + return 0, fmt.Errorf("building BigQuery query failed: %w", err) + } + + count := 0 + for { + var r analysis.Result + err := iter.Next(&r) + if err == iterator.Done { + break + } + if err != nil { + return 0, fmt.Errorf("iterating BigQuery results failed: %w", err) + } + resultID := fmt.Sprintf("%s@%s", r.ModulePath, r.Version) + if _, ok := processedIDs[resultID]; ok { + continue + } + + count++ + b, err := json.MarshalIndent(&r, "", " ") + if err != nil { + return 0, fmt.Errorf("marshalling result: %w", err) + } + _, err = out.Write(b) + if err != nil { + return 0, err + } + fmt.Fprintln(out) + + processedIDs[resultID] = struct{}{} + if r.CreatedAt.After(lastCreatedAt) { + lastCreatedAt = r.CreatedAt + } + } + return count, nil } // requestJSON requests the path from the worker, then reads the returned body
diff --git a/cmd/ejobs/main_test.go b/cmd/ejobs/main_test.go new file mode 100644 index 0000000..f79bc54 --- /dev/null +++ b/cmd/ejobs/main_test.go
@@ -0,0 +1,80 @@ +// Copyright 2025 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package main + +import ( + "bytes" + "context" + "strings" + "testing" + "time" + + "golang.org/x/pkgsite-metrics/internal/analysis" + "golang.org/x/pkgsite-metrics/internal/bigquery" + "google.golang.org/api/iterator" + + bq "cloud.google.com/go/bigquery" +) + +type mockRowIterator struct { + results []*analysis.Result + index int +} + +func (it *mockRowIterator) Next(v interface{}) error { + if it.index >= len(it.results) { + return iterator.Done + } + res := it.results[it.index] + val := v.(*analysis.Result) + *val = *res + it.index++ + return nil +} + +type mockBQClient struct { + expectedResults []*analysis.Result +} + +func (c *mockBQClient) QueryWithParams(ctx context.Context, query string, params []bq.QueryParameter) (bigquery.RowIterator, error) { + return &mockRowIterator{results: c.expectedResults}, nil +} + +func (c *mockBQClient) FullTableName(tableID string) string { return "mock.table" } +func (c *mockBQClient) Close() error { return nil } + +func TestFetchAndPrintResults(t *testing.T) { + mockClient := &mockBQClient{ + expectedResults: []*analysis.Result{ + {ModulePath: "example.com/one", Version: "v1.0.0", CreatedAt: time.Now()}, + {ModulePath: "example.com/two", Version: "v1.2.3", CreatedAt: time.Now().Add(1 * time.Second)}, + }, + } + var out bytes.Buffer + var lastCreatedAt time.Time + processedIDs := make(map[string]struct{}) + + count, err := fetchAndPrintResults(context.Background(), &out, mockClient, "fake-job-id", false, lastCreatedAt, processedIDs) + if err != nil { + t.Fatalf("fetchAndPrintResults failed: %v", err) + } + + if count != len(mockClient.expectedResults) { + t.Errorf("expected to process %d new results, but got %d", len(mockClient.expectedResults), count) + } + got := out.String() + for _, result := range mockClient.expectedResults { + if !strings.Contains(got, result.ModulePath) { + t.Errorf("output did not contain first result; got:\n%s, wanted:%s", got, result.ModulePath) + } + if !strings.Contains(got, result.Version) { + t.Errorf("output did not contain first result; got:\n%s, wanted:%s", got, result.Version) + } + } + + if len(processedIDs) != len(mockClient.expectedResults) { + t.Errorf("expected %d processed IDs, but got %d", len(mockClient.expectedResults), len(processedIDs)) + } +}
diff --git a/internal/bigquery/bigquery.go b/internal/bigquery/bigquery.go index 48d5a2b..1fe3620 100644 --- a/internal/bigquery/bigquery.go +++ b/internal/bigquery/bigquery.go
@@ -32,6 +32,16 @@ deleteDatasetOnClose bool } +type RowIterator interface { + Next(v interface{}) error +} + +func (c *Client) QueryWithParams(ctx context.Context, query string, params []bq.QueryParameter) (RowIterator, error) { + q := c.client.Query(query) + q.Parameters = params + return q.Read(ctx) +} + // NewClientCreate creates a new client for connecting to BigQuery, referring // to a single dataset. It creates the dataset if it doesn't exist. func NewClientCreate(ctx context.Context, projectID, datasetID string) (_ *Client, err error) {