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/BUILD.bazel

@@ -1593,6 +1593,7 @@ GO_TARGETS = [
     "//pkg/roachprod/prometheus:prometheus_test",
     "//pkg/roachprod/promhelperclient:promhelperclient",
     "//pkg/roachprod/promhelperclient:promhelperclient_test",
+    "//pkg/roachprod/roachprodutil:roachprodutil",
     "//pkg/roachprod/ssh:ssh",
     "//pkg/roachprod/ssh:ssh_test",
     "//pkg/roachprod/ui:ui",

pkg/cmd/roachprod/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
         "//pkg/roachprod/fluentbit",
         "//pkg/roachprod/install",
         "//pkg/roachprod/opentelemetry",
+        "//pkg/roachprod/roachprodutil",
         "//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/roachprodutil",
         "//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/roachprodutil"
 	"github.com/cockroachdb/cockroach/pkg/util/httputil"
 	"github.com/cockroachdb/errors"
 	"github.com/go-openapi/strfmt"
@@ -24,14 +24,10 @@ 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
-// variable ServiceAccountJson. This identity token is passed in
-// every request to authenticate with grafana.
+// GCS identity token by using the service account helpers in `roachprodutil/identity.go`.
+// This identity token is passed in every request to authenticate with grafana.
 func newGrafanaClient(
 	ctx context.Context, host string, secure bool,
 ) (*grafana.GrafanaHTTPAPI, error) {
@@ -42,25 +38,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 := roachprodutil.SetServiceAccountCredsEnv(ctx, false); err != nil {
+			return nil, err
 		}
 
-		ts, err := idtoken.NewTokenSource(ctx, grafanaAudience, idtoken.WithCredentialsJSON([]byte(grafanaKey)))
+		token, err := roachprodutil.GetServiceAccountToken(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/roachprodutil"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/ui"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce"
@@ -1372,7 +1373,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),
+`, roachprodutil.ServiceAccountJson, roachprodutil.ServiceAccountAudience),
 	Args: cobra.ExactArgs(2),
 	Run: wrap(func(cmd *cobra.Command, args []string) error {
 		req := grafana.AddAnnotationRequest{

pkg/roachprod/promhelperclient/BUILD.bazel

@@ -2,22 +2,18 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "promhelperclient",
-    srcs = [
-        "client.go",
-        "promhelper_utils.go",
-    ],
+    srcs = ["client.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/roachprod/promhelperclient",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/roachprod/config",
         "//pkg/roachprod/logger",
+        "//pkg/roachprod/roachprodutil",
         "//pkg/roachprod/vm/gce",
         "//pkg/util/httputil",
         "@com_github_cockroachdb_errors//:errors",
-        "@com_google_cloud_go_storage//:storage",
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@org_golang_google_api//idtoken",
-        "@org_golang_google_api//option",
         "@org_golang_x_oauth2//:oauth2",
     ],
 )
@@ -28,6 +24,7 @@ go_test(
     embed = [":promhelperclient"],
     deps = [
         "//pkg/roachprod/logger",
+        "//pkg/roachprod/roachprodutil",
         "@com_github_stretchr_testify//require",
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@org_golang_google_api//idtoken",

pkg/roachprod/promhelperclient/client.go

@@ -20,12 +20,12 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"os"
 	"strconv"
 	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/roachprod/config"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce"
 	"github.com/cockroachdb/cockroach/pkg/util/httputil"
 	"github.com/cockroachdb/errors"
@@ -35,10 +35,8 @@ import (
 )
 
 const (
-	resourceName           = "instance-configs"
-	resourceVersion        = "v1"
-	serviceAccountJson     = "PROM_HELPER_SERVICE_ACCOUNT_JSON"
-	serviceAccountAudience = "PROM_HELPER_SERVICE_ACCOUNT_AUDIENCE"
+	resourceName    = "instance-configs"
+	resourceVersion = "v1"
 )
 
 // SupportedPromProjects are the projects supported for prometheus target
@@ -231,18 +229,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 := roachprodutil.SetServiceAccountCredsEnv(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 := roachprodutil.GetServiceAccountToken(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

@@ -23,6 +23,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
 	"google.golang.org/api/idtoken"
@@ -161,9 +162,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(roachprodutil.ServiceAccountJson, "{}")
 		require.Nil(t, err)
-		err = os.Setenv(serviceAccountAudience, "dummy_audience")
+		err = os.Setenv(roachprodutil.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")
@@ -175,9 +176,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(roachprodutil.ServiceAccountJson, "{}")
 		require.Nil(t, err)
-		err = os.Setenv(serviceAccountAudience, "dummy_audience")
+		err = os.Setenv(roachprodutil.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
@@ -189,9 +190,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(roachprodutil.ServiceAccountJson, "{}")
 		require.Nil(t, err)
-		err = os.Setenv(serviceAccountAudience, "dummy_audience")
+		err = os.Setenv(roachprodutil.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/roachprodutil/BUILD.bazel

@@ -0,0 +1,15 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "roachprodutil",
+    srcs = ["identity.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_google_cloud_go_storage//:storage",
+        "@org_golang_google_api//idtoken",
+        "@org_golang_google_api//option",
+        "@org_golang_x_oauth2//:oauth2",
+    ],
+)