Add async logger for telemetry

acceptance/acceptance_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/databricks/cli/libs/testdiff"
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -156,6 +157,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
 	testdiff.PrepareReplacementsUUID(t, &repls)
 	testdiff.PrepareReplacementsDevVersion(t, &repls)
+	testdiff.PrepareReplaceOS(t, &repls)
 
 	testDirs := getTests(t)
 	require.NotEmpty(t, testDirs)
@@ -281,6 +283,21 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	cmd.Dir = tmpDir
 	err = cmd.Run()
 
+	// Wait for the files to appear before starting assertion on the output.
+	// This is useful for concurrency control when the CLI spawns subprocesses.
+	missingFiles := config.EventuallyFiles
+	assert.Eventually(t, func() bool {
+		for _, file := range missingFiles {
+			_, err := os.Stat(filepath.Join(tmpDir, file))
+			if err == nil {
+				missingFiles = slices.DeleteFunc(missingFiles, func(n string) bool {
+					return n == file
+				})
+			}
+		}
+		return len(missingFiles) == 0
+	}, 10*time.Second, 100*time.Millisecond, "Files did not appear: %v", missingFiles)
+
 	// Write the requests made to the server to a output file if the test is
 	// configured to record requests.
 	if config.RecordRequests {

acceptance/config_test.go

@@ -44,6 +44,10 @@ type TestConfig struct {
 	// Record the requests made to the server and write them as output to
 	// out.requests.txt
 	RecordRequests bool
+
+	// Wait for these files to appear before starting assertion on the output.
+	// This is useful for concurrency control when tests spawn subprocesses.
+	EventuallyFiles []string
 }
 
 type ServerStub struct {

acceptance/libs/telemetry/dummy-without-server/out.upload.txt

@@ -0,0 +1 @@
+error: telemetry endpoint not found

acceptance/libs/telemetry/dummy-without-server/output.txt

@@ -0,0 +1,2 @@
+
+>>> $CLI telemetry dummy

acceptance/libs/telemetry/dummy-without-server/script

@@ -0,0 +1,5 @@
+export DATABRICKS_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt
+
+# This test ensures that the main CLI command does not error even if
+# telemetry upload fails.
+trace $CLI telemetry dummy

acceptance/libs/telemetry/dummy/out.requests.txt

@@ -0,0 +1 @@
+{"method":"POST","path":"/telemetry-ext","body":{"uploadTime":"UNIX_TIME_MILLIS","items":[],"protoLogs":["{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"$DEV_VERSION\",\"command\":\"telemetry_dummy\",\"operating_system\":\"[OS]\",\"execution_time_ms\":0,\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}","{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"$DEV_VERSION\",\"command\":\"telemetry_dummy\",\"operating_system\":\"[OS]\",\"execution_time_ms\":0,\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"]}}

acceptance/libs/telemetry/dummy/out.upload.txt

@@ -0,0 +1,3 @@
+Successfully uploaded telemetry logs
+Response: 
+{"errors":null,"numProtoSuccess":2}

acceptance/libs/telemetry/dummy/output.txt

@@ -0,0 +1,2 @@
+
+>>> $CLI telemetry dummy

acceptance/libs/telemetry/dummy/script

@@ -0,0 +1,3 @@
+export DATABRICKS_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt
+
+trace $CLI telemetry dummy

acceptance/libs/telemetry/dummy/test.toml

@@ -0,0 +1,16 @@
+RecordRequests = true
+EventuallyFiles = [
+    "out.upload.txt"
+]
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "numProtoSuccess": 2
+}
+'''
+
+[[Repls]]
+Old = '17\d{11}'
+New = '"UNIX_TIME_MILLIS"'

acceptance/libs/telemetry/upload/out.requests.txt

@@ -0,0 +1 @@
+{"method":"POST","path":"/telemetry-ext","body":{"uploadTime":1738560432192,"items":[],"protoLogs":["{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE1\"}}}}","{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"]}}

acceptance/libs/telemetry/upload/output.txt

@@ -0,0 +1,5 @@
+
+>>> $CLI telemetry upload
+Successfully uploaded telemetry logs
+Response: 
+{"errors":null,"numProtoSuccess":2}

acceptance/libs/telemetry/upload/script

@@ -0,0 +1,3 @@
+# This command / test cannot be run in inprocess / debug mode. This is because
+# it does not go through the [root.Execute] function.
+trace $CLI telemetry upload < stdin

acceptance/libs/telemetry/upload/stdin

@@ -0,0 +1,24 @@
+{
+    "logs": [
+        {
+            "frontend_log_event_id": "BB79BB52-96F6-42C5-9E44-E63EEA84888D",
+            "entry": {
+                "databricks_cli_log": {
+                    "cli_test_event": {
+                        "name": "VALUE1"
+                    }
+                }
+            }
+        },
+        {
+            "frontend_log_event_id": "A7F597B0-66D1-462D-824C-C5C706F232E8",
+            "entry": {
+                "databricks_cli_log": {
+                    "cli_test_event": {
+                        "name": "VALUE2"
+                    }
+                }
+            }
+        }
+    ]
+}

acceptance/libs/telemetry/upload/test.toml

@@ -0,0 +1,9 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "numProtoSuccess": 2
+}
+'''

cmd/cmd.go

@@ -13,6 +13,7 @@ import (
 	"github.com/databricks/cli/cmd/labs"
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/cmd/sync"
+	"github.com/databricks/cli/cmd/telemetry"
 	"github.com/databricks/cli/cmd/version"
 	"github.com/databricks/cli/cmd/workspace"
 	"github.com/spf13/cobra"
@@ -74,6 +75,6 @@ func New(ctx context.Context) *cobra.Command {
 	cli.AddCommand(labs.New(ctx))
 	cli.AddCommand(sync.New())
 	cli.AddCommand(version.New())
-
+	cli.AddCommand(telemetry.New())
 	return cli
 }

cmd/root/root.go

@@ -2,16 +2,24 @@ package root
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"os"
+	"os/exec"
+	"runtime"
 	"strings"
+	"time"
 
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/dbr"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/spf13/cobra"
 )
 
@@ -73,9 +81,6 @@ func New(ctx context.Context) *cobra.Command {
 		// get the context back
 		ctx = cmd.Context()
 
-		// Detect if the CLI is running on DBR and store this on the context.
-		ctx = dbr.DetectRuntime(ctx)
-
 		// Configure our user agent with the command that's about to be executed.
 		ctx = withCommandInUserAgent(ctx, cmd)
 		ctx = withCommandExecIdInUserAgent(ctx)
@@ -97,7 +102,9 @@ func flagErrorFunc(c *cobra.Command, err error) error {
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute(ctx context.Context, cmd *cobra.Command) error {
-	// TODO: deferred panic recovery
+	ctx = telemetry.WithNewLogger(ctx)
+	ctx = dbr.DetectRuntime(ctx)
+	start := time.Now()
 
 	// Run the command
 	cmd, err := cmd.ExecuteContextC(ctx)
@@ -121,5 +128,72 @@ func Execute(ctx context.Context, cmd *cobra.Command) error {
 		}
 	}
 
+	end := time.Now()
+	exitCode := 0
+	if err != nil {
+		exitCode = 1
+	}
+
+	uploadTelemetry(cmd.Context(), commandString(cmd), start, end, exitCode)
 	return err
 }
+
+func uploadTelemetry(ctx context.Context, cmdStr string, start, end time.Time, exitCode int) {
+	// Nothing to upload.
+	if !telemetry.HasLogs(ctx) {
+		return
+	}
+
+	telemetry.SetExecutionContext(ctx, protos.ExecutionContext{
+		CmdExecID:       cmdExecId,
+		Version:         build.GetInfo().Version,
+		Command:         cmdStr,
+		OperatingSystem: runtime.GOOS,
+		DbrVersion:      env.Get(ctx, dbr.EnvVarName),
+		ExecutionTimeMs: end.Sub(start).Milliseconds(),
+		ExitCode:        int64(exitCode),
+	})
+
+	logs := telemetry.GetLogs(ctx)
+
+	in := telemetry.UploadConfig{
+		Logs: logs,
+	}
+
+	execPath, err := os.Executable()
+	if err != nil {
+		log.Debugf(ctx, "failed to get executable path: %s", err)
+	}
+	telemetryCmd := exec.Command(execPath, "telemetry", "upload")
+	telemetryCmd.Env = os.Environ()
+	for k, v := range auth.Env(ConfigUsed(ctx)) {
+		telemetryCmd.Env = append(telemetryCmd.Env, fmt.Sprintf("%s=%s", k, v))
+	}
+
+	b, err := json.Marshal(in)
+	if err != nil {
+		log.Debugf(ctx, "failed to marshal telemetry logs: %s", err)
+		return
+	}
+
+	stdin, err := telemetryCmd.StdinPipe()
+	if err != nil {
+		log.Debugf(ctx, "failed to create stdin pipe for telemetry worker: %s", err)
+	}
+
+	err = telemetryCmd.Start()
+	if err != nil {
+		log.Debugf(ctx, "failed to start telemetry worker: %s", err)
+		return
+	}
+
+	_, err = stdin.Write(b)
+	if err != nil {
+		log.Debugf(ctx, "failed to write to telemetry worker: %s", err)
+	}
+
+	err = stdin.Close()
+	if err != nil {
+		log.Debugf(ctx, "failed to close stdin for telemetry worker: %s", err)
+	}
+}

cmd/root/user_agent_command_exec_id.go

@@ -7,8 +7,10 @@ import (
 	"github.com/google/uuid"
 )
 
+var cmdExecId = uuid.New().String()
+
 func withCommandExecIdInUserAgent(ctx context.Context) context.Context {
 	// A UUID that will allow us to correlate multiple API requests made by
 	// the same CLI invocation.
-	return useragent.InContext(ctx, "cmd-exec-id", uuid.New().String())
+	return useragent.InContext(ctx, "cmd-exec-id", cmdExecId)
 }

cmd/telemetry/dummy.go

@@ -0,0 +1,31 @@
+package telemetry
+
+import (
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/spf13/cobra"
+)
+
+func newDummyCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "dummy",
+		Short:   "log dummy telemetry events",
+		Long:    "Fire a test telemetry event against the configured Databricks workspace.",
+		Hidden:  true,
+		PreRunE: root.MustWorkspaceClient,
+	}
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		for _, v := range []string{"VALUE1", "VALUE2"} {
+			telemetry.Log(cmd.Context(), protos.DatabricksCliLog{
+				CliTestEvent: &protos.CliTestEvent{
+					Name: protos.DummyCliEnum(v),
+				},
+			})
+		}
+		return nil
+	}
+
+	return cmd
+}

cmd/telemetry/telemetry.go

@@ -0,0 +1,16 @@
+package telemetry
+
+import (
+	"github.com/spf13/cobra"
+)
+
+func New() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:    "telemetry",
+		Short:  "",
+		Hidden: true,
+	}
+
+	cmd.AddCommand(newDummyCommand())
+	return cmd
+}

libs/dbr/detect.go

@@ -11,6 +11,8 @@ import (
 // Dereference [os.Stat] to allow mocking in tests.
 var statFunc = os.Stat
 
+const EnvVarName = "DATABRICKS_RUNTIME_VERSION"
+
 // detect returns true if the current process is running on a Databricks Runtime.
 // Its return value is meant to be cached in the context.
 func detect(ctx context.Context) bool {
@@ -21,7 +23,7 @@ func detect(ctx context.Context) bool {
 	}
 
 	// Databricks Runtime always has the DATABRICKS_RUNTIME_VERSION environment variable set.
-	if value, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION"); !ok || value == "" {
+	if value, ok := env.Lookup(ctx, EnvVarName); !ok || value == "" {
 		return false
 	}
 

libs/telemetry/context.go

@@ -0,0 +1,25 @@
+package telemetry
+
+import (
+	"context"
+	"fmt"
+)
+
+// Private type to store the telemetry logger in the context
+type telemetryLogger int
+
+// Key to store the telemetry logger in the context
+var telemetryLoggerKey telemetryLogger
+
+func WithNewLogger(ctx context.Context) context.Context {
+	return context.WithValue(ctx, telemetryLoggerKey, &logger{})
+}
+
+func fromContext(ctx context.Context) *logger {
+	v := ctx.Value(telemetryLoggerKey)
+	if v == nil {
+		panic(fmt.Errorf("telemetry logger not found in the context"))
+	}
+
+	return v.(*logger)
+}