From 161bb5e1a7ee1a632ee9ada859b454feb382e916 Mon Sep 17 00:00:00 2001 From: Volodymyr Onofriychuk Date: Fri, 10 Apr 2026 12:10:05 +0300 Subject: [PATCH 1/4] feat: enable API consumer tracking by default and add related configuration --- cmd/root.go | 7 +++++++ docs/user-guide/configuration.md | 5 +++++ internal/config/api.go | 7 +++++++ internal/config/auth.go | 4 ++++ internal/config/constants.go | 2 ++ internal/drapi/get.go | 4 ++++ 6 files changed, 29 insertions(+) diff --git a/cmd/root.go b/cmd/root.go index dd6d70bf..6bdf0821 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -211,6 +211,13 @@ func initializeConfig(cmd *cobra.Command) error { _ = viper.BindEnv("external-editor", "VISUAL", "EDITOR") + // API consumer tracking is enabled by default. + // Set DATAROBOT_API_CONSUMER_TRACKING_ENABLED=false to opt out, + // matching the Python SDK convention. + viper.SetDefault(config.APIConsumerTrackingEnabled, true) + + _ = viper.BindEnv(config.APIConsumerTrackingEnabled, "DATAROBOT_API_CONSUMER_TRACKING_ENABLED") + // If DATAROBOT_CLI_CONFIG is set and no explicit --config flag was provided, // use the environment variable value if configFilePath == "" { diff --git a/docs/user-guide/configuration.md b/docs/user-guide/configuration.md index 704f9301..9d0e0e6b 100644 --- a/docs/user-guide/configuration.md +++ b/docs/user-guide/configuration.md @@ -135,6 +135,11 @@ export EDITOR=nano # Force setup wizard to run even if already completed export DATAROBOT_CLI_FORCE_INTERACTIVE=true + +# API consumer tracking (default: true) +# Set to false to disable the X-DataRobot-Api-Consumer-Trace header on API requests. +# Matches the Python SDK's DATAROBOT_API_CONSUMER_TRACKING_ENABLED behavior. +export DATAROBOT_API_CONSUMER_TRACKING_ENABLED=false ``` ### Advanced flags diff --git a/internal/config/api.go b/internal/config/api.go index 6008e383..9bde5b65 100644 --- a/internal/config/api.go +++ b/internal/config/api.go @@ -69,6 +69,13 @@ func GetUserAgentHeader() string { return version.GetAppNameVersionText() } +// IsAPIConsumerTrackingEnabled returns true if the X-DataRobot-Api-Consumer-Trace +// header should be sent with API requests. Controlled by the +// DATAROBOT_API_CONSUMER_TRACKING_ENABLED environment variable (default: true). +func IsAPIConsumerTrackingEnabled() bool { + return viper.GetBool(APIConsumerTrackingEnabled) +} + func RedactedReqInfo(req *http.Request) string { // Dump the request to a byte slice after cloning and removing Auth header dumpReq := req.Clone(req.Context()) diff --git a/internal/config/auth.go b/internal/config/auth.go index f573e868..0e06d7a0 100644 --- a/internal/config/auth.go +++ b/internal/config/auth.go @@ -62,6 +62,10 @@ func VerifyToken(ctx context.Context, datarobotEndpoint, token string) error { req.Header.Add("Authorization", bearer) req.Header.Add("User-Agent", getUserAgent(ctx)) + if IsAPIConsumerTrackingEnabled() { + req.Header.Add("X-DataRobot-Api-Consumer-Trace", getUserAgent(ctx)) + } + log.Debug("Request Info: \n" + RedactedReqInfo(req)) client := &http.Client{ diff --git a/internal/config/constants.go b/internal/config/constants.go index 309e51c9..105898f9 100644 --- a/internal/config/constants.go +++ b/internal/config/constants.go @@ -17,4 +17,6 @@ package config const ( DataRobotURL = "endpoint" DataRobotAPIKey = "token" + + APIConsumerTrackingEnabled = "api-consumer-tracking-enabled" ) diff --git a/internal/drapi/get.go b/internal/drapi/get.go index 3cb869d3..7cee48eb 100644 --- a/internal/drapi/get.go +++ b/internal/drapi/get.go @@ -46,6 +46,10 @@ func Get(url, info string) (*http.Response, error) { req.Header.Add("Authorization", "Bearer "+token) req.Header.Add("User-Agent", config.GetUserAgentHeader()) + if config.IsAPIConsumerTrackingEnabled() { + req.Header.Add("X-DataRobot-Api-Consumer-Trace", config.GetUserAgentHeader()) + } + if info != "" { log.Infof("Fetching %s from: %s", info, url) } From b417cec860c9e7579a254ed074b31c4beecbf4b5 Mon Sep 17 00:00:00 2001 From: Volodymyr Onofriychuk Date: Fri, 10 Apr 2026 14:42:04 +0300 Subject: [PATCH 2/4] implement API consumer tracking with command path tracing --- cmd/root.go | 2 ++ docs/user-guide/configuration.md | 2 ++ internal/config/api.go | 33 +++++++++++++++++ internal/config/api_test.go | 62 ++++++++++++++++++++++++++++++++ internal/config/auth.go | 2 +- internal/drapi/get.go | 2 +- 6 files changed, 101 insertions(+), 2 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 6bdf0821..0f69e4f8 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -90,6 +90,8 @@ using pre-built templates. Get from idea to production in minutes, not hours. // Store telemetry client in context for use by commands cmd.SetContext(context.WithValue(cmd.Context(), telemetryClientKey{}, client)) + config.SetAPIConsumerTrace(config.CommandPathToTrace(cmd.CommandPath())) + return nil }, PersistentPostRunE: func(cmd *cobra.Command, _ []string) error { diff --git a/docs/user-guide/configuration.md b/docs/user-guide/configuration.md index 9d0e0e6b..357dfb4a 100644 --- a/docs/user-guide/configuration.md +++ b/docs/user-guide/configuration.md @@ -139,6 +139,8 @@ export DATAROBOT_CLI_FORCE_INTERACTIVE=true # API consumer tracking (default: true) # Set to false to disable the X-DataRobot-Api-Consumer-Trace header on API requests. # Matches the Python SDK's DATAROBOT_API_CONSUMER_TRACKING_ENABLED behavior. +# When enabled, the header value identifies the command being run using dot-notation: +# datarobot.cli.. (e.g. datarobot.cli.templates.setup) export DATAROBOT_API_CONSUMER_TRACKING_ENABLED=false ``` diff --git a/internal/config/api.go b/internal/config/api.go index 9bde5b65..a5a2b094 100644 --- a/internal/config/api.go +++ b/internal/config/api.go @@ -69,6 +69,39 @@ func GetUserAgentHeader() string { return version.GetAppNameVersionText() } +// apiConsumerTrace holds the dot-notation command path set at startup. +// Example: "datarobot.cli.templates.setup" +var apiConsumerTrace string + +// SetAPIConsumerTrace stores the dot-notation trace value for the running command. +// Called once during PersistentPreRunE with the result of CommandPathToTrace. +func SetAPIConsumerTrace(trace string) { + apiConsumerTrace = trace +} + +// GetAPIConsumerTrace returns the dot-notation trace value for the running command. +// Falls back to "datarobot.cli" if SetAPIConsumerTrace has not been called. +func GetAPIConsumerTrace() string { + if apiConsumerTrace == "" { + return "datarobot.cli" + } + + return apiConsumerTrace +} + +// CommandPathToTrace converts a cobra command path (e.g. "dr templates setup") +// to the canonical dot-notation trace format (e.g. "datarobot.cli.templates.setup"). +func CommandPathToTrace(commandPath string) string { + parts := strings.Fields(commandPath) + if len(parts) == 0 { + return "datarobot.cli" + } + + parts[0] = "datarobot.cli" + + return strings.Join(parts, ".") +} + // IsAPIConsumerTrackingEnabled returns true if the X-DataRobot-Api-Consumer-Trace // header should be sent with API requests. Controlled by the // DATAROBOT_API_CONSUMER_TRACKING_ENABLED environment variable (default: true). diff --git a/internal/config/api_test.go b/internal/config/api_test.go index 69546453..e0978d1f 100644 --- a/internal/config/api_test.go +++ b/internal/config/api_test.go @@ -107,3 +107,65 @@ func (suite *APITestSuite) TestSetURLToConfigDoesNotWriteFile() { configFile := filepath.Join(suite.tempDir, ".config/datarobot/drconfig.yaml") suite.NoFileExists(configFile, "SetURLToConfig must not write the config file to disk") } + +func (suite *APITestSuite) TestCommandPathToTrace() { + tests := []struct { + name string + input string + expected string + }{ + { + name: "empty string", + input: "", + expected: "datarobot.cli", + }, + { + name: "root command only", + input: "dr", + expected: "datarobot.cli", + }, + { + name: "single subcommand", + input: "dr start", + expected: "datarobot.cli.start", + }, + { + name: "nested subcommand", + input: "dr templates setup", + expected: "datarobot.cli.templates.setup", + }, + { + name: "deeply nested subcommand", + input: "dr self plugin add", + expected: "datarobot.cli.self.plugin.add", + }, + } + + for _, tc := range tests { + suite.Run(tc.name, func() { + suite.Equal(tc.expected, CommandPathToTrace(tc.input)) + }) + } +} + +func (suite *APITestSuite) TestGetSetAPIConsumerTrace() { + SetAPIConsumerTrace("") + + suite.Equal("datarobot.cli", GetAPIConsumerTrace(), "should fall back to datarobot.cli when unset") + + SetAPIConsumerTrace("datarobot.cli.templates.list") + suite.Equal("datarobot.cli.templates.list", GetAPIConsumerTrace()) + + SetAPIConsumerTrace("datarobot.cli.start") + suite.Equal("datarobot.cli.start", GetAPIConsumerTrace()) +} + +func (suite *APITestSuite) TestIsAPIConsumerTrackingEnabled() { + suite.False(IsAPIConsumerTrackingEnabled(), "should be false when viper has no config set") + + viper.Set(APIConsumerTrackingEnabled, true) + suite.True(IsAPIConsumerTrackingEnabled()) + + viper.Set(APIConsumerTrackingEnabled, false) + suite.False(IsAPIConsumerTrackingEnabled()) +} diff --git a/internal/config/auth.go b/internal/config/auth.go index 0e06d7a0..f4d95cb1 100644 --- a/internal/config/auth.go +++ b/internal/config/auth.go @@ -63,7 +63,7 @@ func VerifyToken(ctx context.Context, datarobotEndpoint, token string) error { req.Header.Add("User-Agent", getUserAgent(ctx)) if IsAPIConsumerTrackingEnabled() { - req.Header.Add("X-DataRobot-Api-Consumer-Trace", getUserAgent(ctx)) + req.Header.Add("X-DataRobot-Api-Consumer-Trace", "datarobot.cli.auth.verify") } log.Debug("Request Info: \n" + RedactedReqInfo(req)) diff --git a/internal/drapi/get.go b/internal/drapi/get.go index 7cee48eb..9f1d5324 100644 --- a/internal/drapi/get.go +++ b/internal/drapi/get.go @@ -47,7 +47,7 @@ func Get(url, info string) (*http.Response, error) { req.Header.Add("User-Agent", config.GetUserAgentHeader()) if config.IsAPIConsumerTrackingEnabled() { - req.Header.Add("X-DataRobot-Api-Consumer-Trace", config.GetUserAgentHeader()) + req.Header.Add("X-DataRobot-Api-Consumer-Trace", config.GetAPIConsumerTrace()) } if info != "" { From 30c00d39e18753ff1a7ce238f8a50bb80b4d7287 Mon Sep 17 00:00:00 2001 From: Volodymyr Onofriychuk Date: Tue, 14 Apr 2026 18:21:44 +0300 Subject: [PATCH 3/4] feat: add test for command path resolution with aliases in cobra --- internal/config/api_test.go | 49 +++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/internal/config/api_test.go b/internal/config/api_test.go index e0978d1f..d40c20c4 100644 --- a/internal/config/api_test.go +++ b/internal/config/api_test.go @@ -19,6 +19,7 @@ import ( "path/filepath" "testing" + "github.com/spf13/cobra" "github.com/spf13/viper" "github.com/stretchr/testify/suite" ) @@ -169,3 +170,51 @@ func (suite *APITestSuite) TestIsAPIConsumerTrackingEnabled() { viper.Set(APIConsumerTrackingEnabled, false) suite.False(IsAPIConsumerTrackingEnabled()) } + +// TestCommandPathToTraceWithAliases verifies that cobra always resolves +// command aliases to their canonical Use name before CommandPath() is called. +// This means CommandPathToTrace never receives an alias string — cobra +// normalises it first — so the trace always uses the canonical command name. +func (suite *APITestSuite) TestCommandPathToTraceWithAliases() { + // Build a small cobra command tree that mirrors real CLI aliases: + // root (Use: "dr") + // └─ run (Use: "run", Aliases: ["r"]) → like dr run / dr r + // └─ start (Use: "start", Aliases: ["quickstart"]) → like dr start / dr quickstart + // └─ templates (Use: "templates", Aliases: ["template"]) + // └─ list (Use: "list") + root := &cobra.Command{Use: "dr"} + + run := &cobra.Command{Use: "run [tasks]", Aliases: []string{"r"}} + start := &cobra.Command{Use: "start", Aliases: []string{"quickstart"}} + templates := &cobra.Command{Use: "templates", Aliases: []string{"template"}} + list := &cobra.Command{Use: "list"} + + templates.AddCommand(list) + root.AddCommand(run, start, templates) + + tests := []struct { + name string + args []string + expectedTrace string + }{ + // Canonical invocations + {"canonical: dr run", []string{"run"}, "datarobot.cli.run"}, + {"canonical: dr start", []string{"start"}, "datarobot.cli.start"}, + {"canonical: dr templates list", []string{"templates", "list"}, "datarobot.cli.templates.list"}, + // Cobra command aliases — CommandPath() must return the canonical name + {"alias: dr r → dr run", []string{"r"}, "datarobot.cli.run"}, + {"alias: dr quickstart → dr start", []string{"quickstart"}, "datarobot.cli.start"}, + {"alias: dr template list → dr templates list", []string{"template", "list"}, "datarobot.cli.templates.list"}, + } + + for _, tc := range tests { + suite.Run(tc.name, func() { + cmd, _, err := root.Find(tc.args) + suite.Require().NoError(err) + suite.Require().NotNil(cmd) + + trace := CommandPathToTrace(cmd.CommandPath()) + suite.Equal(tc.expectedTrace, trace) + }) + } +} From 3a4afa819fb9d66ce76b54fc8f11ac5916764652 Mon Sep 17 00:00:00 2001 From: Volodymyr Onofriychuk Date: Tue, 14 Apr 2026 18:21:51 +0300 Subject: [PATCH 4/4] fix: update import statement for logging package in drapi/get.go --- internal/drapi/get.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/drapi/get.go b/internal/drapi/get.go index 9f1d5324..7d4b14f0 100644 --- a/internal/drapi/get.go +++ b/internal/drapi/get.go @@ -21,8 +21,8 @@ import ( "net/http" "time" - "github.com/charmbracelet/log" "github.com/datarobot/cli/internal/config" + "github.com/datarobot/cli/internal/log" ) var token string