cmd/jobs: command for working with jobs
This is a CLI tool for working with ecosystem-metrics jobs.
Currently it supports displaying, listing and canceling jobs.
It will eventually support starting jobs.
Change-Id: Iae83c08ca600947b72076e8dc213c84eb797c0c6
Reviewed-on: https://go-review.googlesource.com/c/pkgsite-metrics/+/498265
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Zvonimir Pavlinovic <zpavlinovic@google.com>
diff --git a/cmd/jobs/main.go b/cmd/jobs/main.go
new file mode 100644
index 0000000..da13d7c
--- /dev/null
+++ b/cmd/jobs/main.go
@@ -0,0 +1,204 @@
+// Copyright 2023 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.
+
+// Command jobs supports jobs on ecosystem-metrics.
+package main
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "flag"
+ "fmt"
+ "io"
+ "net/http"
+ "os"
+ "reflect"
+
+ credsapi "cloud.google.com/go/iam/credentials/apiv1"
+ credspb "cloud.google.com/go/iam/credentials/apiv1/credentialspb"
+ "golang.org/x/pkgsite-metrics/internal/jobs"
+)
+
+var env = flag.String("env", "prod", "worker environment (dev or prod)")
+
+var commands = []command{
+ {"print-identity-token", "",
+ "print an identity token", doPrintToken},
+ {"list", "",
+ "list jobs", doList},
+ {"show", "jobID...",
+ "display information about jobs", doShow},
+ {"cancel", "jobID...",
+ "cancel the jobs", doCancel},
+}
+
+type command struct {
+ name string
+ argdoc string
+ desc string
+ run func(context.Context, []string) error
+}
+
+func main() {
+ flag.Usage = func() {
+ out := flag.CommandLine.Output()
+ fmt.Fprintln(out, "usage:")
+ for _, cmd := range commands {
+ fmt.Fprintf(out, " job %s %s\n", cmd.name, cmd.argdoc)
+ fmt.Fprintf(out, "\t%s\n", cmd.desc)
+ }
+ fmt.Fprintln(out, "\ncommon flags:")
+ flag.PrintDefaults()
+ }
+
+ flag.Parse()
+ if err := run(context.Background()); err != nil {
+ fmt.Fprintf(os.Stderr, "%v\n", err)
+ flag.Usage()
+ os.Exit(2)
+ }
+}
+
+var workerURL string
+
+func run(ctx context.Context) error {
+ wu := os.Getenv("GO_ECOSYSTEM_WORKER_URL_SUFFIX")
+ if wu == "" {
+ return errors.New("need GO_ECOSYSTEM_WORKER_URL_SUFFIX environment variable")
+ }
+ workerURL = fmt.Sprintf("https://%s-%s", *env, wu)
+ name := flag.Arg(0)
+ for _, cmd := range commands {
+ if cmd.name == name {
+ return cmd.run(ctx, flag.Args()[1:])
+ }
+ }
+ return fmt.Errorf("unknown command %q", name)
+}
+
+func doShow(ctx context.Context, args []string) error {
+ token, err := requestImpersonateIdentityToken(ctx)
+ if err != nil {
+ return err
+ }
+ for _, jobID := range args {
+ if err := showJob(ctx, jobID, token); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+func showJob(ctx context.Context, jobID, token string) error {
+ job, err := requestJSON[jobs.Job](ctx, "jobs/describe?jobid="+jobID, token)
+ if err != nil {
+ return err
+ }
+ rj := reflect.ValueOf(job).Elem()
+ rt := rj.Type()
+ for i := 0; i < rt.NumField(); i++ {
+ f := rt.Field(i)
+ if f.IsExported() {
+ v := rj.FieldByIndex(f.Index)
+ fmt.Printf("%s: %v\n", f.Name, v.Interface())
+ }
+ }
+ return nil
+}
+
+func doList(ctx context.Context, _ []string) error {
+ token, err := requestImpersonateIdentityToken(ctx)
+ if err != nil {
+ return err
+ }
+ body, err := httpGet(ctx, workerURL+"/jobs/list", token)
+ if err != nil {
+ return err
+ }
+ fmt.Printf("%s\n", body)
+ return nil
+}
+
+func doCancel(ctx context.Context, args []string) error {
+ token, err := requestImpersonateIdentityToken(ctx)
+ if err != nil {
+ return err
+ }
+ for _, jobID := range args {
+ if _, err := httpGet(ctx, workerURL+"/jobs/cancel?jobid="+jobID, token); err != nil {
+ return fmt.Errorf("canceling %q: %w", jobID, err)
+ }
+ }
+ return nil
+}
+
+// For testing. Can be used in place of `gcloud auth print-identity-token`.
+func doPrintToken(ctx context.Context, _ []string) error {
+ token, err := requestImpersonateIdentityToken(ctx)
+ if err != nil {
+ return err
+ }
+ fmt.Println(token)
+ return nil
+}
+
+// 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, token string) (*T, error) {
+ body, err := httpGet(ctx, workerURL+"/"+path, token)
+ if err != nil {
+ return nil, err
+ }
+ var t T
+ if err := json.Unmarshal(body, &t); err != nil {
+ return nil, err
+ }
+ return &t, nil
+}
+
+// httpGet makes a GET request to the given URL with the given identity token.
+// It reads the body and returns the HTTP response and the body.
+func httpGet(ctx context.Context, url, token string) (body []byte, err error) {
+ req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+ if err != nil {
+ return nil, err
+ }
+ req.Header.Add("Authorization", "Bearer "+token)
+ res, err := http.DefaultClient.Do(req)
+ if err != nil {
+ return nil, err
+ }
+ defer res.Body.Close()
+ body, err = io.ReadAll(res.Body)
+ if err != nil {
+ return nil, fmt.Errorf("reading body (%s): %v", res.Status, err)
+ }
+ if res.StatusCode != 200 {
+ return nil, fmt.Errorf("%s: %s", res.Status, body)
+ }
+ return body, nil
+}
+
+// requestImpersonateIdentityToken requests an identity token for a service
+// account to impersonate from the iamcredentials service.
+// See https://cloud.google.com/iam/docs/reference/credentials/rest.
+func requestImpersonateIdentityToken(ctx context.Context) (string, error) {
+ c, err := credsapi.NewIamCredentialsClient(ctx)
+ if err != nil {
+ return "", err
+ }
+ defer c.Close()
+ serviceAccountEmail := "impersonate@go-ecosystem.iam.gserviceaccount.com"
+ req := &credspb.GenerateIdTokenRequest{
+ Name: "projects/-/serviceAccounts/" + serviceAccountEmail,
+ Audience: workerURL,
+ IncludeEmail: true,
+ }
+ res, err := c.GenerateIdToken(ctx, req)
+ if err != nil {
+ return "", fmt.Errorf("GenerateIdToken: %w", err)
+ }
+ return res.Token, nil
+}
diff --git a/go.mod b/go.mod
index 9357cf3..ab648c9 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@
cloud.google.com/go/cloudtasks v1.8.0
cloud.google.com/go/errorreporting v0.3.0
cloud.google.com/go/firestore v1.9.0
+ cloud.google.com/go/iam v0.11.0
cloud.google.com/go/logging v1.7.0
cloud.google.com/go/secretmanager v1.9.0
cloud.google.com/go/storage v1.28.1
@@ -39,7 +40,6 @@
require (
cloud.google.com/go/compute v1.18.0 // indirect
cloud.google.com/go/compute/metadata v0.2.3 // indirect
- cloud.google.com/go/iam v0.11.0 // indirect
cloud.google.com/go/longrunning v0.3.0 // indirect
cloud.google.com/go/monitoring v1.8.0 // indirect
cloud.google.com/go/trace v1.4.0 // indirect
