roachtest: grafana annotations read creds from cloud storage

As of cockroachdb#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

@@ -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",

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",

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

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{

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
 }

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

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
 }