From e8d2bde5d66317bc12a8d232ac68e471e2616aa8 Mon Sep 17 00:00:00 2001 From: Arjun Mahishi Date: Mon, 23 Sep 2024 13:22:22 +0530 Subject: [PATCH] pkg/cli: update default behaviour to read datadog api/app keys from env 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 --- pkg/cli/debug.go | 11 ++++++----- pkg/cli/env.go | 10 ++++++++++ pkg/cli/zip_upload.go | 5 +++++ pkg/cli/zip_upload_test.go | 10 +++++++--- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go index 40ce18c4a959..3023f06c974d 100644 --- 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 index c7f0f1b19c39..87b40daf0bfc 100644 --- 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 index 31ee50698417..3577ace7e92a 100644 --- 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 index 36f56e0ccafe..6649edce2a27 100644 --- 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)