From b32947f080de868a3c710341bb078c1b00bd7109 Mon Sep 17 00:00:00 2001 From: DarrylWong Date: Wed, 12 Jun 2024 14:51:16 -0400 Subject: [PATCH] roachtest: grafana annotations read creds from cloud storage As of #124099 we now store the service account creds in cloud storage. We already use this to access prometheus when generating dynamic configs. This change does the same for Grafana annotations by extracting the common logic into a helper. This will allow users to have access to Grafana annotations out of the box locally, and limit the amount of benign but potentially confusing warnings about invalid credentials. Epic: none Fixes: none Release note: none --- pkg/cmd/roachprod/BUILD.bazel | 1 + pkg/cmd/roachprod/grafana/BUILD.bazel | 1 + pkg/cmd/roachprod/grafana/annotations.go | 25 ++----- pkg/cmd/roachprod/main.go | 3 +- pkg/roachprod/promhelperclient/client.go | 21 +++--- pkg/roachprod/promhelperclient/client_test.go | 12 ++-- .../promhelperclient/promhelper_utils.go | 67 +++++++++++-------- 7 files changed, 64 insertions(+), 66 deletions(-) diff --git a/pkg/cmd/roachprod/BUILD.bazel b/pkg/cmd/roachprod/BUILD.bazel index 556b0e350f25..edd51e48081c 100644 --- a/pkg/cmd/roachprod/BUILD.bazel +++ b/pkg/cmd/roachprod/BUILD.bazel @@ -19,6 +19,7 @@ go_library( "//pkg/roachprod/fluentbit", "//pkg/roachprod/install", "//pkg/roachprod/opentelemetry", + "//pkg/roachprod/promhelperclient", "//pkg/roachprod/ssh", "//pkg/roachprod/ui", "//pkg/roachprod/vm", diff --git a/pkg/cmd/roachprod/grafana/BUILD.bazel b/pkg/cmd/roachprod/grafana/BUILD.bazel index 403485b37838..9132b0464d07 100644 --- a/pkg/cmd/roachprod/grafana/BUILD.bazel +++ b/pkg/cmd/roachprod/grafana/BUILD.bazel @@ -22,6 +22,7 @@ go_library( importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/grafana", visibility = ["//visibility:public"], deps = [ + "//pkg/roachprod/promhelperclient", "//pkg/util/httputil", "@com_github_cockroachdb_errors//:errors", "@com_github_go_openapi_strfmt//:strfmt", diff --git a/pkg/cmd/roachprod/grafana/annotations.go b/pkg/cmd/roachprod/grafana/annotations.go index 55788da13c89..5387a45e7acf 100644 --- a/pkg/cmd/roachprod/grafana/annotations.go +++ b/pkg/cmd/roachprod/grafana/annotations.go @@ -13,9 +13,9 @@ package grafana import ( "context" "fmt" - "os" "strings" + "github.com/cockroachdb/cockroach/pkg/roachprod/promhelperclient" "github.com/cockroachdb/cockroach/pkg/util/httputil" "github.com/cockroachdb/errors" "github.com/go-openapi/strfmt" @@ -24,9 +24,6 @@ import ( "google.golang.org/api/idtoken" ) -const ServiceAccountJson = "GRAFANA_SERVICE_ACCOUNT_JSON" -const ServiceAccountAudience = "GRAFANA_SERVICE_ACCOUNT_AUDIENCE" - // newGrafanaClient is a helper function that creates an HTTP client to // create grafana api calls with. If secure is true, it tries to get a // GCS identity token by using the service account specified by the env @@ -42,25 +39,15 @@ func newGrafanaClient( scheme = "https" // Read in the service account key and audience, so we can retrieve the identity token. - grafanaKey := os.Getenv(ServiceAccountJson) - if grafanaKey == "" { - return nil, errors.Newf("%s env variable was not found", ServiceAccountJson) - } - grafanaAudience := os.Getenv(ServiceAccountAudience) - if grafanaAudience == "" { - return nil, errors.Newf("%s env variable was not found", ServiceAccountAudience) + if _, err := promhelperclient.SetPromHelperCredsEnv(ctx, false); err != nil { + return nil, err } - ts, err := idtoken.NewTokenSource(ctx, grafanaAudience, idtoken.WithCredentialsJSON([]byte(grafanaKey))) + token, err := promhelperclient.GetToken(ctx, idtoken.NewTokenSource) if err != nil { - return nil, errors.Wrap(err, "Error creating GCS oauth token source from specified credential") + return nil, err } - token, err := ts.Token() - if err != nil { - return nil, errors.Wrap(err, "Error getting identity token") - } - - headers["Authorization"] = fmt.Sprintf("Bearer %s", token.AccessToken) + headers["Authorization"] = fmt.Sprintf("Bearer %s", token) } headers[httputil.ContentTypeHeader] = httputil.JSONContentType diff --git a/pkg/cmd/roachprod/main.go b/pkg/cmd/roachprod/main.go index b1c8804eb8d0..02bf4516dbe4 100644 --- a/pkg/cmd/roachprod/main.go +++ b/pkg/cmd/roachprod/main.go @@ -31,6 +31,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/roachprod/config" rperrors "github.com/cockroachdb/cockroach/pkg/roachprod/errors" "github.com/cockroachdb/cockroach/pkg/roachprod/install" + "github.com/cockroachdb/cockroach/pkg/roachprod/promhelperclient" "github.com/cockroachdb/cockroach/pkg/roachprod/ui" "github.com/cockroachdb/cockroach/pkg/roachprod/vm" "github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce" @@ -1358,7 +1359,7 @@ creates an annotation over time range. Example: # Create an annotation over time range 1-100 on the centralized grafana instance, which needs authentication. roachprod grafana-annotation grafana.testeng.crdb.io example-annotation-event --tags my-cluster --tags test-run-1 --dashboard-uid overview --time-range 1,100 -`, grafana.ServiceAccountJson, grafana.ServiceAccountAudience), +`, promhelperclient.ServiceAccountJson, promhelperclient.ServiceAccountAudience), Args: cobra.ExactArgs(2), Run: wrap(func(cmd *cobra.Command, args []string) error { req := grafana.AddAnnotationRequest{ diff --git a/pkg/roachprod/promhelperclient/client.go b/pkg/roachprod/promhelperclient/client.go index b864be16f792..7076d13e0db6 100644 --- a/pkg/roachprod/promhelperclient/client.go +++ b/pkg/roachprod/promhelperclient/client.go @@ -20,7 +20,6 @@ import ( "fmt" "io" "net/http" - "os" "strconv" "strings" @@ -37,8 +36,8 @@ import ( const ( resourceName = "instance-configs" resourceVersion = "v1" - serviceAccountJson = "PROM_HELPER_SERVICE_ACCOUNT_JSON" - serviceAccountAudience = "PROM_HELPER_SERVICE_ACCOUNT_AUDIENCE" + ServiceAccountJson = "PROM_HELPER_SERVICE_ACCOUNT_JSON" + ServiceAccountAudience = "PROM_HELPER_SERVICE_ACCOUNT_AUDIENCE" ) // SupportedPromProjects are the projects supported for prometheus target @@ -228,18 +227,14 @@ func (c *PromClient) getToken( return "", nil } // Read in the service account key and audience, so we can retrieve the identity token. - if _, err := setPromHelperCredsEnv(ctx, forceFetchCreds, l); err != nil { + if credSrc, err := SetPromHelperCredsEnv(ctx, forceFetchCreds); err != nil { return "", err + } else { + l.Printf("Prometheus credentials obtained from %s", credSrc) } - grafanaKey := os.Getenv(serviceAccountJson) - grafanaAudience := os.Getenv(serviceAccountAudience) - ts, err := c.newTokenSource(ctx, grafanaAudience, idtoken.WithCredentialsJSON([]byte(grafanaKey))) + token, err := GetToken(ctx, c.newTokenSource) if err != nil { - return "", errors.Wrap(err, "error creating GCS oauth token source from specified credential") - } - token, err := ts.Token() - if err != nil { - return "", errors.Wrap(err, "error getting identity token") + return "", err } - return fmt.Sprintf("Bearer %s", token.AccessToken), nil + return fmt.Sprintf("Bearer %s", token), nil } diff --git a/pkg/roachprod/promhelperclient/client_test.go b/pkg/roachprod/promhelperclient/client_test.go index e5c927067862..13f36b41537b 100644 --- a/pkg/roachprod/promhelperclient/client_test.go +++ b/pkg/roachprod/promhelperclient/client_test.go @@ -150,9 +150,9 @@ func Test_getToken(t *testing.T) { require.Empty(t, token) }) t.Run("invalid credentials", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return nil, fmt.Errorf("invalid") @@ -164,9 +164,9 @@ func Test_getToken(t *testing.T) { require.Equal(t, "error creating GCS oauth token source from specified credential: invalid", err.Error()) }) t.Run("invalid token", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return &mockToken{token: "", err: fmt.Errorf("failed")}, nil @@ -178,9 +178,9 @@ func Test_getToken(t *testing.T) { require.Equal(t, "error getting identity token: failed", err.Error()) }) t.Run("success", func(t *testing.T) { - err := os.Setenv(serviceAccountJson, "{}") + err := os.Setenv(ServiceAccountJson, "{}") require.Nil(t, err) - err = os.Setenv(serviceAccountAudience, "dummy_audience") + err = os.Setenv(ServiceAccountAudience, "dummy_audience") require.Nil(t, err) c.newTokenSource = func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error) { return &mockToken{token: "token"}, nil diff --git a/pkg/roachprod/promhelperclient/promhelper_utils.go b/pkg/roachprod/promhelperclient/promhelper_utils.go index 15010074ff4c..c3f5ee6e682d 100644 --- a/pkg/roachprod/promhelperclient/promhelper_utils.go +++ b/pkg/roachprod/promhelperclient/promhelper_utils.go @@ -12,14 +12,15 @@ package promhelperclient import ( "context" - "fmt" "io" "os" "path/filepath" "strings" "cloud.google.com/go/storage" - "github.com/cockroachdb/cockroach/pkg/roachprod/logger" + "github.com/cockroachdb/errors" + "golang.org/x/oauth2" + "google.golang.org/api/idtoken" "google.golang.org/api/option" ) @@ -47,9 +48,9 @@ const ( objectLocation = "promhelpers-secrets" ) -// setPromHelperCredsEnv sets the environment variables serviceAccountAudience and -// serviceAccountJson based on the following conditions: -// > check if forFetch is false +// SetPromHelperCredsEnv sets the environment variables ServiceAccountAudience and +// ServiceAccountJson based on the following conditions: +// > check if forceFetch is false // // > forceFetch is false // > if env is set return @@ -58,64 +59,76 @@ const ( // > read the creds from secrets manager // // > set the env variable and save the creds to the promCredFile -func setPromHelperCredsEnv( - ctx context.Context, forceFetch bool, l *logger.Logger, -) (FetchedFrom, error) { +func SetPromHelperCredsEnv(ctx context.Context, forceFetch bool) (FetchedFrom, error) { creds := "" fetchedFrom := Env if !forceFetch { // bypass environment and creds file if forceFetch is false // check if environment is set - audience := os.Getenv(serviceAccountAudience) - saJson := os.Getenv(serviceAccountJson) + audience := os.Getenv(ServiceAccountAudience) + saJson := os.Getenv(ServiceAccountJson) if audience != "" && saJson != "" { - l.Printf("Secrets obtained from environment.") return fetchedFrom, nil } // check if the secrets file is available b, err := os.ReadFile(promCredFile) if err == nil { - l.Printf("Secrets obtained from temp file: %s", promCredFile) creds = string(b) fetchedFrom = File } } if creds == "" { // creds == "" means (env is not set and the file does not have the creds) or forceFetch is true - l.Printf("creds need to be fetched from store.") options := []option.ClientOption{option.WithScopes(storage.ScopeReadOnly)} cj := os.Getenv("GOOGLE_EPHEMERAL_CREDENTIALS") - if len(cj) != 0 { - options = append(options, option.WithCredentialsJSON([]byte(cj))) - } else { - l.Printf("GOOGLE_EPHEMERAL_CREDENTIALS env is not set.") + if cj == "" { + return "", errors.New("SetPromHelperCredsEnv: GOOGLE_EPHEMERAL_CREDENTIALS env is not set") } + options = append(options, option.WithCredentialsJSON([]byte(cj))) + client, err := storage.NewClient(ctx, options...) if err != nil { - return fetchedFrom, err + return "", err } defer func() { _ = client.Close() }() fetchedFrom = Store obj := client.Bucket(bucket).Object(objectLocation) r, err := obj.NewReader(ctx) if err != nil { - return fetchedFrom, err + return "", err } defer func() { _ = r.Close() }() body, err := io.ReadAll(r) creds = string(body) if err != nil { - return fetchedFrom, err - } - err = os.WriteFile(promCredFile, body, 0700) - if err != nil { - l.Errorf("error writing to the credential file: %v", err) + return "", err } + _ = os.WriteFile(promCredFile, body, 0700) } secretValues := strings.Split(creds, secretsDelimiter) if len(secretValues) == 2 { - _ = os.Setenv(serviceAccountAudience, secretValues[0]) - _ = os.Setenv(serviceAccountJson, secretValues[1]) + _ = os.Setenv(ServiceAccountAudience, secretValues[0]) + _ = os.Setenv(ServiceAccountJson, secretValues[1]) return fetchedFrom, nil } - return fetchedFrom, fmt.Errorf("invalid secret values - %s", creds) + return "", errors.Newf("invalid secret values - %s", creds) +} + +// GetToken returns a GCS oauth token based on the service account key and audience +// set through the ServiceAccountJson and ServiceAccountAudience env vars. +// Assumes that the env vars have been set already, i.e. through SetPromHelperCredsEnv. +func GetToken( + ctx context.Context, + tokenSource func(ctx context.Context, audience string, opts ...idtoken.ClientOption) (oauth2.TokenSource, error), +) (string, error) { + key := os.Getenv(ServiceAccountJson) + audience := os.Getenv(ServiceAccountAudience) + ts, err := tokenSource(ctx, audience, idtoken.WithCredentialsJSON([]byte(key))) + if err != nil { + return "", errors.Wrap(err, "error creating GCS oauth token source from specified credential") + } + token, err := ts.Token() + if err != nil { + return "", errors.Wrap(err, "error getting identity token") + } + return token.AccessToken, nil }