Merge cockroachdb#131178

131178: pkg/cli: read datadog api/app keys from env also r=arjunmahishi a=arjunmahishi Currently, the datadog keys are set only via cli flags (for the tsdump and the debug zip upload commands). This commit introduces the ability to also set them via env vars Part of: CRDB-42143 Release note: None Co-authored-by: Arjun Mahishi <[email protected]>
RaduBerinde · Sep 25, 2024 · 8846663 · 8846663
2 parents a55b017 + e8d2bde
commit 8846663
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 8 deletions.
diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go
@@ -1574,11 +1574,11 @@ func init() {
 		"tenant IDs to filter logs by")
 
 	f = debugZipUploadCmd.Flags()
-	f.StringVar(&debugZipUploadOpts.ddAPIKey, "dd-api-key", "",
+	f.StringVar(&debugZipUploadOpts.ddAPIKey, "dd-api-key", getEnvOrDefault(datadogAPIKeyEnvVar, ""),
 		"Datadog API key to use to send debug.zip artifacts to datadog")
-	f.StringVar(&debugZipUploadOpts.ddAPPKey, "dd-app-key", "",
+	f.StringVar(&debugZipUploadOpts.ddAPPKey, "dd-app-key", getEnvOrDefault(datadogAPPKeyEnvVar, ""),
 		"Datadog APP key to use to send debug.zip artifacts to datadog")
-	f.StringVar(&debugZipUploadOpts.ddSite, "dd-site", defaultDDSite,
+	f.StringVar(&debugZipUploadOpts.ddSite, "dd-site", getEnvOrDefault(datadogSiteEnvVar, defaultDDSite),
 		"Datadog site to use to send debug.zip artifacts to datadog")
 	f.StringSliceVar(&debugZipUploadOpts.include, "include", nil,
 		"The debug zip artifacts to include. Possible values: "+strings.Join(zipArtifactTypes, ", "))
@@ -1621,9 +1621,10 @@ func init() {
 		"", "prometheus label for cluster name")
 	f.StringVar(&debugTimeSeriesDumpOpts.yaml, "yaml", debugTimeSeriesDumpOpts.yaml, "full path to create the tsdump.yaml with storeID: nodeID mappings (raw format only). This file is required when loading the raw tsdump for troubleshooting.")
 	f.StringVar(&debugTimeSeriesDumpOpts.targetURL, "target-url", "", "target URL to send openmetrics data over HTTP")
-	f.StringVar(&debugTimeSeriesDumpOpts.ddSite, "dd-site", "us5",
+	f.StringVar(&debugTimeSeriesDumpOpts.ddSite, "dd-site", getEnvOrDefault(datadogSiteEnvVar, defaultDDSite),
 		"Datadog site to use to send tsdump artifacts to datadog")
-	f.StringVar(&debugTimeSeriesDumpOpts.ddApiKey, "dd-api-key", "", "Datadog API key to use to send to the datadog formatter")
+	f.StringVar(&debugTimeSeriesDumpOpts.ddApiKey, "dd-api-key", getEnvOrDefault(datadogAPIKeyEnvVar, ""),
+		"Datadog API key to use to send to the datadog formatter")
 	f.StringVar(&debugTimeSeriesDumpOpts.httpToken, "http-token", "", "HTTP header to use with the json export format")
 
 	f = debugSendKVBatchCmd.Flags()

diff --git a/pkg/cli/env.go b/pkg/cli/env.go
@@ -22,3 +22,13 @@ func getDefaultHost() string {
 	}
 	return ""
 }
+
+// getEnvOrDefault is a helper function that returns the value of an environment
+// variable or a default value if the environment variable is not set.
+func getEnvOrDefault(key, def string) string {
+	if v := os.Getenv(key); v != "" {
+		return v
+	}
+
+	return def
+}
diff --git a/pkg/cli/zip_upload.go b/pkg/cli/zip_upload.go
@@ -59,6 +59,11 @@ const (
 	defaultDDSite       = "us5"
 	defaultGCPProjectID = "arjun-sandbox-424904" // TODO: change this project ID
 
+	// datadog env vars
+	datadogSiteEnvVar   = "DD_SITE"
+	datadogAPIKeyEnvVar = "DD_API_KEY"
+	datadogAPPKeyEnvVar = "DD_APP_KEY"
+
 	// datadog HTTP headers
 	datadogAPIKeyHeader = "DD-API-KEY"
 	datadogAppKeyHeader = "DD-APPLICATION-KEY"

diff --git a/pkg/cli/zip_upload_test.go b/pkg/cli/zip_upload_test.go
@@ -199,10 +199,7 @@ func TestUploadZipEndToEnd(t *testing.T) {
 			require.NoError(t, err)
 
 			lines := strings.Split(finaloutput.String(), "\n")
-			// if d.Cmd == "upload-profiles" {
-			// sort the lines to avoid flakiness in the test
 			sort.Strings(lines)
-			// }
 
 			// replace the debugDir with a constant string to avoid flakiness in the test
 			return strings.ReplaceAll(strings.TrimSpace(strings.Join(lines, "\n")), debugDir, "debugDir")
@@ -274,6 +271,9 @@ func uploadProfileHook(t *testing.T, req *http.Request) ([]byte, error) {
 	_, params, _ := mime.ParseMediaType(req.Header.Get("Content-Type"))
 	reader := multipart.NewReader(req.Body, params["boundary"])
 
+	// validate the headers
+	require.Equal(t, "dd-api-key", req.Header.Get("DD-API-KEY"))
+
 	// find the "event" part in the multipart request and copy it to the final output
 	for {
 		part, err := reader.NextPart()
@@ -317,6 +317,10 @@ func uploadProfileHook(t *testing.T, req *http.Request) ([]byte, error) {
 func setupDDArchiveHook(t *testing.T, req *http.Request) ([]byte, error) {
 	t.Helper()
 
+	// validate the headers
+	require.Equal(t, "dd-api-key", req.Header.Get("DD-API-KEY"))
+	require.Equal(t, "dd-app-key", req.Header.Get("DD-APPLICATION-KEY"))
+
 	var body bytes.Buffer
 	_, err := body.ReadFrom(req.Body)
 	require.NoError(t, err)