diff --git a/cmd/root.go b/cmd/root.go index dd6d70bf..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 { @@ -211,6 +213,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..357dfb4a 100644 --- a/docs/user-guide/configuration.md +++ b/docs/user-guide/configuration.md @@ -135,6 +135,13 @@ 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. +# 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 ``` ### Advanced flags diff --git a/internal/config/api.go b/internal/config/api.go index 6008e383..a5a2b094 100644 --- a/internal/config/api.go +++ b/internal/config/api.go @@ -69,6 +69,46 @@ 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). +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/api_test.go b/internal/config/api_test.go index 69546453..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" ) @@ -107,3 +108,113 @@ 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()) +} + +// 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) + }) + } +} diff --git a/internal/config/auth.go b/internal/config/auth.go index f573e868..f4d95cb1 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", "datarobot.cli.auth.verify") + } + 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..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 @@ -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.GetAPIConsumerTrace()) + } + if info != "" { log.Infof("Fetching %s from: %s", info, url) }