cli: use Datadog API client in tsdump uploader

Previously, we used raw POST requests in the tsdump uploader. Now, we
use the Datadog API client which is already linked in elsewhere in the
CockroachDB binary, for use in the cockroach workload command.

Epic: none
Release note: None

Author: jordanlewis

pkg/cli/BUILD.bazel

@@ -256,6 +256,8 @@ go_library(
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_cockroachdb_ttycolor//:ttycolor",
+        "@com_github_datadog_datadog_api_client_go_v2//api/datadog",
+        "@com_github_datadog_datadog_api_client_go_v2//api/datadogV2",
         "@com_github_dustin_go_humanize//:go-humanize",
         "@com_github_fsnotify_fsnotify//:fsnotify",
         "@com_github_gogo_protobuf//jsonpb",
@@ -467,6 +469,7 @@ go_test(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_errors//oserror",
         "@com_github_cockroachdb_pebble//vfs",
+        "@com_github_datadog_datadog_api_client_go_v2//api/datadogV2",
         "@com_github_google_pprof//profile",
         "@com_github_pmezard_go_difflib//difflib",
         "@com_github_spf13_cobra//:cobra",

pkg/cli/debug.go

@@ -1615,6 +1615,7 @@ func init() {
 	f.StringVar(&debugTimeSeriesDumpOpts.organizationName, "org-name", "", "organization name to use in datadog upload")
 	f.StringVar(&debugTimeSeriesDumpOpts.userName, "user-name", "", "name of the user to perform datadog upload")
 	f.StringVar(&debugTimeSeriesDumpOpts.storeToNodeMapYAMLFile, "store-to-node-map-file", "", "yaml file path which contains the mapping of store ID to node ID for datadog upload.")
+	f.BoolVar(&debugTimeSeriesDumpOpts.dryRun, "dry-run", false, "run in dry-run mode without making any actual uploads")
 
 	f = debugSendKVBatchCmd.Flags()
 	f.StringVar(&debugSendKVBatchContext.traceFormat, "trace", debugSendKVBatchContext.traceFormat,

pkg/cli/testdata/tsdump/json

@@ -5,4 +5,4 @@ cr.node.sql.query.count,2021-01-01T00:00:00Z,1,100.5
 cr.node.sql.query.count,2021-01-01T00:00:10Z,1,102.3
 cr.store.rocksdb.block.cache.usage,2021-01-01T00:00:00Z,2,75.2
 ----
-{"series":[{"metric":"crdb.tsdump.admission.admitted.elastic-cpu","type":0,"points":[{"timestamp":1748248320,"value":1}],"interval":10,"resources":null,"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"]},{"metric":"crdb.tsdump.sql.query.count","type":0,"points":[{"timestamp":1609459200,"value":100.5}],"interval":10,"resources":null,"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"]},{"metric":"crdb.tsdump.sql.query.count","type":0,"points":[{"timestamp":1609459210,"value":102.3}],"interval":10,"resources":null,"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"]},{"metric":"crdb.tsdump.rocksdb.block.cache.usage","type":0,"points":[{"timestamp":1609459200,"value":75.2}],"interval":10,"resources":null,"tags":["store:2","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"]}]}
+{"series":[{"interval":10,"metric":"crdb.tsdump.admission.admitted.elastic-cpu","points":[{"timestamp":1748248320,"value":1}],"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0},{"interval":10,"metric":"crdb.tsdump.sql.query.count","points":[{"timestamp":1609459200,"value":100.5}],"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0},{"interval":10,"metric":"crdb.tsdump.sql.query.count","points":[{"timestamp":1609459210,"value":102.3}],"tags":["node_id:1","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0},{"interval":10,"metric":"crdb.tsdump.rocksdb.block.cache.usage","points":[{"timestamp":1609459200,"value":75.2}],"tags":["store:2","cluster_type:SELF_HOSTED","cluster_label:\"test-cluster\"","cluster_id:test-cluster-id","zendesk_ticket:zd-test","org_name:test-org","user_name:test-user","upload_id:\"test-cluster\"-20241114000000","upload_timestamp:2024-11-14 00:00:00","upload_year:2024","upload_month:11","upload_day:14"],"type":0}]}

pkg/cli/testdata/tsdump_upload_e2e

@@ -51,6 +51,7 @@ var debugTimeSeriesDumpOpts = struct {
 	organizationName       string
 	userName               string
 	storeToNodeMapYAMLFile string
+	dryRun                 bool
 }{
 	format:       tsDumpText,
 	from:         timestampValue{},
@@ -59,6 +60,9 @@ var debugTimeSeriesDumpOpts = struct {
 	yaml:         "/tmp/tsdump.yaml",
 }
 
+// hostNameOverride is used to override the hostname for testing purpose.
+var hostNameOverride string
+
 var debugTimeSeriesDumpCmd = &cobra.Command{
 	Use:   "tsdump",
 	Short: "dump all the raw timeseries values in a cluster",
@@ -111,18 +115,16 @@ will then convert it to the --format requested in the current invocation.
 				return errors.New("no input file provided")
 			}
 
-			targetURL, err := getDatadogTargetURL(debugTimeSeriesDumpOpts.ddSite)
-			if err != nil {
-				return err
-			}
-
-			var datadogWriter = makeDatadogWriter(
-				targetURL,
+			datadogWriter, err := makeDatadogWriter(
+				debugTimeSeriesDumpOpts.ddSite,
 				cmd == tsDumpDatadogInit,
 				debugTimeSeriesDumpOpts.ddApiKey,
 				100,
-				doDDRequest,
+				hostNameOverride,
 			)
+			if err != nil {
+				return err
+			}
 			return datadogWriter.upload(args[0])
 		case tsDumpOpenMetrics:
 			if debugTimeSeriesDumpOpts.targetURL != "" {
@@ -259,15 +261,6 @@ will then convert it to the --format requested in the current invocation.
 	}),
 }
 
-func getDatadogTargetURL(site string) (string, error) {
-	host, ok := ddSiteToHostMap[site]
-	if !ok {
-		return "", fmt.Errorf("unsupported datadog site '%s'", site)
-	}
-	targetURL := fmt.Sprintf(targetURLFormat, host)
-	return targetURL, nil
-}
-
 func doRequest(req *http.Request) error {
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {

pkg/cli/tsdump.go

@@ -13,12 +13,14 @@ import (
 	"math/rand"
 	"net"
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/DataDog/datadog-api-client-go/v2/api/datadogV2"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
@@ -145,7 +147,7 @@ func parseTSInput(t *testing.T, input string, w tsWriter) {
 }
 
 func parseDDInput(t *testing.T, input string, w *datadogWriter) {
-	var data *DatadogSeries
+	var data *datadogV2.MetricSeries
 	var source, storeNodeKey string
 
 	for _, s := range strings.Split(input, "\n") {
@@ -164,10 +166,10 @@ func parseDDInput(t *testing.T, input string, w *datadogWriter) {
 			(data != nil && data.Metric != metricName ||
 				(data != nil && source != nameValueTimestamp[1])) {
 			if data != nil {
-				_, err := w.emitDataDogMetrics([]DatadogSeries{*data})
+				_, err := w.emitDataDogMetrics([]datadogV2.MetricSeries{*data})
 				require.NoError(t, err)
 			}
-			data = &DatadogSeries{
+			data = &datadogV2.MetricSeries{
 				Metric: metricName,
 				Type:   w.resolveMetricType(metricName),
 			}
@@ -178,12 +180,12 @@ func parseDDInput(t *testing.T, input string, w *datadogWriter) {
 		require.NoError(t, err)
 		ts, err := strconv.ParseInt(nameValueTimestamp[3], 10, 64)
 		require.NoError(t, err)
-		data.Points = append(data.Points, DatadogPoint{
-			Value:     value,
-			Timestamp: ts,
+		data.Points = append(data.Points, datadogV2.MetricPoint{
+			Value:     &value,
+			Timestamp: &ts,
 		})
 	}
-	_, err := w.emitDataDogMetrics([]DatadogSeries{*data})
+	_, err := w.emitDataDogMetrics([]datadogV2.MetricSeries{*data})
 	require.NoError(t, err)
 }
 
@@ -193,7 +195,17 @@ func TestTsDumpFormatsDataDriven(t *testing.T) {
 	defer testutils.TestingHook(&getCurrentTime, func() time.Time {
 		return time.Date(2024, 11, 14, 0, 0, 0, 0, time.UTC)
 	})()
-
+	var testReqs []*http.Request
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		r2 := r.Clone(r.Context())
+		// Clone the body so that it can be read again
+		body, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		r2.Body = io.NopCloser(bytes.NewReader(body))
+		testReqs = append(testReqs, r2)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
 	datadriven.Walk(t, "testdata/tsdump", func(t *testing.T, path string) {
 		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
 			var w tsWriter
@@ -204,28 +216,25 @@ func TestTsDumpFormatsDataDriven(t *testing.T) {
 				debugTimeSeriesDumpOpts.zendeskTicket = "zd-test"
 				debugTimeSeriesDumpOpts.organizationName = "test-org"
 				debugTimeSeriesDumpOpts.userName = "test-user"
-				var testReqs []*http.Request
 				var series int
 				d.ScanArgs(t, "series-threshold", &series)
-				var ddwriter = makeDatadogWriter(
-					"https://example.com/data", d.Cmd == "format-datadog-init", "api-key", series, func(req *http.Request,
-					) error {
-						testReqs = append(testReqs, req)
-						return nil
-					})
+				var ddwriter, err = makeDatadogWriter(
+					defaultDDSite, d.Cmd == "format-datadog-init", "api-key", series,
+					server.Listener.Addr().String(),
+				)
+				require.NoError(t, err)
 
 				parseDDInput(t, d.Input, ddwriter)
 
 				out := strings.Builder{}
 				for _, tr := range testReqs {
-					rc, err := tr.GetBody()
-					require.NoError(t, err)
-					zipR, err := gzip.NewReader(rc)
+					reader, err := gzip.NewReader(tr.Body)
 					require.NoError(t, err)
-					body, err := io.ReadAll(zipR)
+					body, err := io.ReadAll(reader)
 					require.NoError(t, err)
 					out.WriteString(fmt.Sprintf("%s: %s\nDD-API-KEY: %s\nBody: %s", tr.Method, tr.URL, tr.Header.Get("DD-API-KEY"), body))
 				}
+				testReqs = testReqs[:0] // reset the slice
 				return out.String()
 			case "format-json":
 				debugTimeSeriesDumpOpts.clusterLabel = "test-cluster"