Merge f98369d9e9 into 24ac8d8d59

## Changes

This PR adds acceptance tests for native Databricks auth methods: basic,
oauth, and pat.

In the future we could compare this with auth credentials used by
downstream tools like TF or the telemetry process to ensure consistent
auth credentials are picked up and used.

Note: 
We do not add acceptance tests for other auth methods like Azure because
they communicate with external endpoints. To test them locally, we would
need to set up a reverse proxy server, which is out of scope for this
change.

## Tests
N/A

NOTICE

@@ -114,3 +114,7 @@ dario.cat/mergo
 Copyright (c) 2013 Dario Castañé. All rights reserved.
 Copyright (c) 2012 The Go Authors. All rights reserved.
 https://github.com/darccio/mergo/blob/master/LICENSE
+
+https://github.com/gorilla/mux
+Copyright (c) 2023 The Gorilla Authors. All rights reserved.
+https://github.com/gorilla/mux/blob/main/LICENSE

acceptance/acceptance_test.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"slices"
 	"sort"
@@ -26,6 +27,7 @@ import (
 	"github.com/databricks/cli/libs/testdiff"
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/stretchr/testify/require"
 )
 
@@ -72,7 +74,8 @@ func TestInprocessMode(t *testing.T) {
 	if InprocessMode {
 		t.Skip("Already tested by TestAccept")
 	}
-	require.Equal(t, 1, testAccept(t, true, "selftest"))
+	require.Equal(t, 1, testAccept(t, true, "selftest/basic"))
+	require.Equal(t, 1, testAccept(t, true, "selftest/server"))
 }
 
 func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
@@ -118,14 +121,12 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	uvCache := getUVDefaultCacheDir(t)
 	t.Setenv("UV_CACHE_DIR", uvCache)
 
-	ctx := context.Background()
 	cloudEnv := os.Getenv("CLOUD_ENV")
 
 	if cloudEnv == "" {
 		defaultServer := testserver.New(t)
 		AddHandlers(defaultServer)
-		// Redirect API access to local server:
+		t.Setenv("DATABRICKS_DEFAULT_HOST", defaultServer.URL)
-		t.Setenv("DATABRICKS_HOST", defaultServer.URL)
 
 		homeDir := t.TempDir()
 		// Do not read user's ~/.databrickscfg
@@ -148,27 +149,12 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	// do it last so that full paths match first:
 	repls.SetPath(buildDir, "[BUILD_DIR]")
 
-	var config databricks.Config
-	if cloudEnv == "" {
-		// use fake token for local tests
-		config = databricks.Config{Token: "dbapi1234"}
-	} else {
-		// non-local tests rely on environment variables
-		config = databricks.Config{}
-	}
-	workspaceClient, err := databricks.NewWorkspaceClient(&config)
-	require.NoError(t, err)
-
-	user, err := workspaceClient.CurrentUser.Me(ctx)
-	require.NoError(t, err)
-	require.NotNil(t, user)
-	testdiff.PrepareReplacementsUser(t, &repls, *user)
-	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
-	testdiff.PrepareReplacementsUUID(t, &repls)
 	testdiff.PrepareReplacementsDevVersion(t, &repls)
 	testdiff.PrepareReplacementSdkVersion(t, &repls)
 	testdiff.PrepareReplacementsGoVersion(t, &repls)
 
+	repls.Repls = append(repls.Repls, testdiff.Replacement{Old: regexp.MustCompile("dbapi[0-9a-f]+"), New: "[DATABRICKS_TOKEN]"})
+
 	testDirs := getTests(t)
 	require.NotEmpty(t, testDirs)
 
@@ -180,8 +166,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	}
 
 	for _, dir := range testDirs {
-		testName := strings.ReplaceAll(dir, "\\", "/")
+		t.Run(dir, func(t *testing.T) {
-		t.Run(testName, func(t *testing.T) {
 			if !InprocessMode {
 				t.Parallel()
 			}
@@ -203,7 +188,8 @@ func getTests(t *testing.T) []string {
 		name := filepath.Base(path)
 		if name == EntryPointScript {
 			// Presence of 'script' marks a test case in this directory
-			testDirs = append(testDirs, filepath.Dir(path))
+			testName := filepath.ToSlash(filepath.Dir(path))
+			testDirs = append(testDirs, testName)
 		}
 		return nil
 	})
@@ -239,7 +225,6 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	}
 
 	repls.SetPathWithParents(tmpDir, "[TMPDIR]")
-	repls.Repls = append(repls.Repls, config.Repls...)
 
 	scriptContents := readMergedScriptContents(t, dir)
 	testutil.WriteFile(t, filepath.Join(tmpDir, EntryPointScript), scriptContents)
@@ -253,38 +238,79 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Env = os.Environ()
 
+	var workspaceClient *databricks.WorkspaceClient
+	var user iam.User
+
 	// Start a new server with a custom configuration if the acceptance test
 	// specifies a custom server stubs.
 	var server *testserver.Server
 
-	// Start a new server for this test if either:
+	if cloudEnv == "" {
-	// 1. A custom server spec is defined in the test configuration.
+		// Start a new server for this test if either:
-	// 2. The test is configured to record requests and assert on them. We need
+		// 1. A custom server spec is defined in the test configuration.
-	//    a duplicate of the default server to record requests because the default
+		// 2. The test is configured to record requests and assert on them. We need
-	//    server otherwise is a shared resource.
+		//    a duplicate of the default server to record requests because the default
-	if len(config.Server) > 0 || config.RecordRequests {
+		//    server otherwise is a shared resource.
-		server = testserver.New(t)
-		server.RecordRequests = config.RecordRequests
-		server.IncludeRequestHeaders = config.IncludeRequestHeaders
 
-		// If no custom server stubs are defined, add the default handlers.
+		databricksLocalHost := os.Getenv("DATABRICKS_DEFAULT_HOST")
-		if len(config.Server) == 0 {
+
+		if len(config.Server) > 0 || config.RecordRequests {
+			server = testserver.New(t)
+			server.RecordRequests = config.RecordRequests
+			server.IncludeRequestHeaders = config.IncludeRequestHeaders
+
+			for _, stub := range config.Server {
+				require.NotEmpty(t, stub.Pattern)
+				items := strings.Split(stub.Pattern, " ")
+				require.Len(t, items, 2)
+				server.Handle(items[0], items[1], func(fakeWorkspace *testserver.FakeWorkspace, req *http.Request) (any, int) {
+					statusCode := http.StatusOK
+					if stub.Response.StatusCode != 0 {
+						statusCode = stub.Response.StatusCode
+					}
+					return stub.Response.Body, statusCode
+				})
+			}
+
+			// The earliest handlers take precedence, add default handlers last
 			AddHandlers(server)
+			databricksLocalHost = server.URL
 		}
 
-		for _, stub := range config.Server {
+		// Each local test should use a new token that will result into a new fake workspace,
-			require.NotEmpty(t, stub.Pattern)
+		// so that test don't interfere with each other.
-			server.Handle(stub.Pattern, func(fakeWorkspace *testserver.FakeWorkspace, req *http.Request) (any, int) {
+		tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
-				statusCode := http.StatusOK
+		config := databricks.Config{
-				if stub.Response.StatusCode != 0 {
+			Host:  databricksLocalHost,
-					statusCode = stub.Response.StatusCode
+			Token: "dbapi" + tokenSuffix,
-				}
-				return stub.Response.Body, statusCode
-			})
 		}
-		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+server.URL)
+		workspaceClient, err = databricks.NewWorkspaceClient(&config)
+		require.NoError(t, err)
+
+		cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+config.Host)
+		cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+config.Token)
+
+		// For the purposes of replacements, use testUser.
+		// Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement:
+		user = testUser
+	} else {
+		// Use whatever authentication mechanism is configured by the test runner.
+		workspaceClient, err = databricks.NewWorkspaceClient(&databricks.Config{})
+		require.NoError(t, err)
+		pUser, err := workspaceClient.CurrentUser.Me(context.Background())
+		require.NoError(t, err, "Failed to get current user")
+		user = *pUser
 	}
 
+	testdiff.PrepareReplacementsUser(t, &repls, user)
+	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
+
+	// Must be added PrepareReplacementsUser, otherwise conflicts with [USERNAME]
+	testdiff.PrepareReplacementsUUID(t, &repls)
+
+	// User replacements come last:
+	repls.Repls = append(repls.Repls, config.Repls...)
+
 	if coverDir != "" {
 		// Creating individual coverage directory for each test, because writing to the same one
 		// results in sporadic failures like this one (only if tests are running in parallel):
@@ -295,15 +321,6 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 		cmd.Env = append(cmd.Env, "GOCOVERDIR="+coverDir)
 	}
 
-	// Each local test should use a new token that will result into a new fake workspace,
-	// so that test don't interfere with each other.
-	if cloudEnv == "" {
-		tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
-		token := "dbapi" + tokenSuffix
-		cmd.Env = append(cmd.Env, "DATABRICKS_TOKEN="+token)
-		repls.Set(token, "[DATABRICKS_TOKEN]")
-	}
-
 	// Write combined output to a file
 	out, err := os.Create(filepath.Join(tmpDir, "output.txt"))
 	require.NoError(t, err)
@@ -320,7 +337,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 
 		for _, req := range server.Requests {
 			reqJson, err := json.MarshalIndent(req, "", "  ")
-			require.NoError(t, err)
+			require.NoErrorf(t, err, "Failed to indent: %#v", req)
 
 			reqJsonWithRepls := repls.Replace(string(reqJson))
 			_, err = f.WriteString(reqJsonWithRepls + "\n")

acceptance/auth/bundle_and_profile/output.txt

@@ -13,13 +13,13 @@
 
 === Inside the bundle, profile flag not matching bundle host. Badness: should use profile from flag instead and not fail
 >>> errcode [CLI] current-user me -p profile_name
-Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_URL]
+Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_TARGET]
 
 Exit code: 1
 
 === Inside the bundle, target and not matching profile
 >>> errcode [CLI] current-user me -t dev -p profile_name
-Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_URL]
+Error: cannot resolve bundle auth configuration: config host mismatch: profile uses host https://non-existing-subdomain.databricks.com, but CLI configured to use [DATABRICKS_TARGET]
 
 Exit code: 1
 

acceptance/auth/bundle_and_profile/test.toml

@@ -5,4 +5,8 @@ Badness = "When -p flag is used inside the bundle folder for any CLI commands, C
 # This is a workaround to replace DATABRICKS_URL with DATABRICKS_HOST
 [[Repls]]
 Old='DATABRICKS_HOST'
-New='DATABRICKS_URL'
+New='DATABRICKS_TARGET'
+
+[[Repls]]
+Old='DATABRICKS_URL'
+New='DATABRICKS_TARGET'

acceptance/auth/credentials/basic/out.requests.txt

@@ -0,0 +1,12 @@
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ],
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/basic"
+    ]
+  },
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}

acceptance/auth/credentials/basic/output.txt

@@ -0,0 +1,4 @@
+{
+  "id":"[USERID]",
+  "userName":"[USERNAME]"
+}

acceptance/auth/credentials/basic/script

@@ -0,0 +1,8 @@
+# Unset the token which is configured by default
+# in acceptance tests
+export DATABRICKS_TOKEN=""
+
+export DATABRICKS_USERNAME=username
+export DATABRICKS_PASSWORD=password
+
+$CLI current-user me

acceptance/auth/credentials/basic/test.toml

@@ -0,0 +1,4 @@
+# "username:password" in base64 is dXNlcm5hbWU6cGFzc3dvcmQ=, expect to see this in Authorization header
+[[Repls]]
+Old = "dXNlcm5hbWU6cGFzc3dvcmQ="
+New = "[ENCODED_AUTH]"

acceptance/auth/credentials/oauth/out.requests.txt

@@ -0,0 +1,34 @@
+{
+  "headers": {
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS]"
+    ]
+  },
+  "method": "GET",
+  "path": "/oidc/.well-known/oauth-authorization-server"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ],
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS]"
+    ]
+  },
+  "method": "POST",
+  "path": "/oidc/v1/token",
+  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Bearer oauth-token"
+    ],
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/oauth-m2m"
+    ]
+  },
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}

acceptance/auth/credentials/oauth/output.txt

@@ -0,0 +1,4 @@
+{
+  "id":"[USERID]",
+  "userName":"[USERNAME]"
+}

acceptance/auth/credentials/oauth/script

@@ -0,0 +1,8 @@
+# Unset the token which is configured by default
+# in acceptance tests
+export DATABRICKS_TOKEN=""
+
+export DATABRICKS_CLIENT_ID=client_id
+export DATABRICKS_CLIENT_SECRET=client_secret
+
+$CLI current-user me

acceptance/auth/credentials/oauth/test.toml

@@ -0,0 +1,5 @@
+# "client_id:client_secret" in base64 is Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ=, expect to
+# see this in Authorization header
+[[Repls]]
+Old = "Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ="
+New = "[ENCODED_AUTH]"

acceptance/auth/credentials/pat/out.requests.txt

@@ -0,0 +1,12 @@
+{
+  "headers": {
+    "Authorization": [
+      "Bearer dapi1234"
+    ],
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/pat"
+    ]
+  },
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}

acceptance/auth/credentials/pat/output.txt

@@ -0,0 +1,4 @@
+{
+  "id":"[USERID]",
+  "userName":"[USERNAME]"
+}

acceptance/auth/credentials/pat/script

@@ -0,0 +1,3 @@
+export DATABRICKS_TOKEN=dapi1234
+
+$CLI current-user me

acceptance/auth/credentials/test.toml

@@ -0,0 +1,20 @@
+LocalOnly = true
+
+RecordRequests = true
+IncludeRequestHeaders = ["Authorization", "User-Agent"]
+
+[[Repls]]
+Old = '(linux|darwin|windows)'
+New = '[OS]'
+
+[[Repls]]
+Old = " upstream/[A-Za-z0-9.-]+"
+New = ""
+
+[[Repls]]
+Old = " upstream-version/[A-Za-z0-9.-]+"
+New = ""
+
+[[Repls]]
+Old = " cicd/[A-Za-z0-9.-]+"
+New = ""

acceptance/bin/wait_file

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+wait_file() {
+    local file_path="$1"
+    local max_attempts=100
+    local attempt=0
+
+    while [ $attempt -lt $max_attempts ]; do
+        if [ -e "$file_path" ]; then
+            echo "File $file_path exists"
+            return 0
+        fi
+        sleep 0.1
+        attempt=$((attempt + 1))
+    done
+
+    echo "Timeout: File $file_path did not appear within 10 seconds"
+    return 1
+}
+
+if [ $# -eq 0 ]; then
+    echo "Usage: $0 <file_path>"
+    exit 1
+fi
+
+wait_file "$1"
+exit $?

acceptance/bin/wait_pid

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+
+# wait <pid> in bash only works for processes that are direct children to the calling
+# shell. This script is more general purpose.
+wait_pid() {
+    local pid=$1
+    local max_attempts=100  # 100 * 0.1 seconds = 10 seconds
+    local attempt=0
+    local sleep_time=0.1
+
+    while [ $attempt -lt $max_attempts ]; do
+        if [[ "$OSTYPE" == "msys"* || "$OSTYPE" == "cygwin"* ]]; then
+            # Windows approach
+            if ! tasklist | grep -q $pid; then
+                echo "Process has ended"
+                return 0
+            fi
+        else
+            # Linux/macOS approach
+            if ! kill -0 $pid 2>/dev/null; then
+                echo "Process has ended"
+                return 0
+            fi
+        fi
+
+        sleep $sleep_time
+        attempt=$((attempt + 1))
+    done
+
+    echo "Timeout: Process $pid did not end within 10 seconds"
+    return 1
+}
+
+# Usage
+if [ $# -eq 0 ]; then
+    echo "Usage: $0 <PID>"
+    exit 1
+fi
+
+wait_pid $1
+exit $?

acceptance/cmd_server_test.go

@@ -14,11 +14,7 @@ import (
 
 func StartCmdServer(t *testing.T) *testserver.Server {
 	server := testserver.New(t)
-
+	server.Handle("GET", "/", func(_ *testserver.FakeWorkspace, r *http.Request) (any, int) {
-	// {$} is a wildcard that only matches the end of the URL. We explicitly use
-	// /{$} to disambiguate it from the generic handler for '/' which is used to
-	// identify unhandled API endpoints in the test server.
-	server.Handle("/{$}", func(w *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		q := r.URL.Query()
 		args := strings.Split(q.Get("args"), " ")
 

acceptance/selftest/server/out.requests.txt

@@ -0,0 +1,8 @@
+{
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}
+{
+  "method": "GET",
+  "path": "/custom/endpoint"
+}

acceptance/selftest/server/output.txt

@@ -0,0 +1,15 @@
+
+>>> curl -s [DATABRICKS_URL]/api/2.0/preview/scim/v2/Me
+{
+    "id": "[USERID]",
+    "userName": "[USERNAME]"
+}
+>>> curl -sD - [DATABRICKS_URL]/custom/endpoint?query=param
+HTTP/1.1 201 Created
+Content-Type: application/json
+Date: (redacted)
+Content-Length: (redacted)
+
+custom
+---
+response

acceptance/selftest/server/script

@@ -0,0 +1,2 @@
+trace curl -s $DATABRICKS_HOST/api/2.0/preview/scim/v2/Me
+trace curl -sD - $DATABRICKS_HOST/custom/endpoint?query=param

acceptance/selftest/server/test.toml

@@ -0,0 +1,18 @@
+LocalOnly = true
+RecordRequests = true
+
+[[Server]]
+Pattern = "GET /custom/endpoint"
+Response.Body = '''custom
+---
+response
+'''
+Response.StatusCode = 201
+
+[[Repls]]
+Old = 'Date: .*'
+New = 'Date: (redacted)'
+
+[[Repls]]
+Old = 'Content-Length: [0-9]*'
+New = 'Content-Length: (redacted)'

acceptance/server_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/gorilla/mux"
 
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -16,8 +17,13 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
+var testUser = iam.User{
+	Id:       "1000012345",
+	UserName: "tester@databricks.com",
+}
+
 func AddHandlers(server *testserver.Server) {
-	server.Handle("GET /api/2.0/policies/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/policies/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{
 				{
@@ -32,7 +38,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/instance-pools/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/instance-pools/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListInstancePools{
 			InstancePools: []compute.InstancePoolAndStats{
 				{
@@ -43,7 +49,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.1/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.1/clusters/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return compute.ListClustersResponse{
 			Clusters: []compute.ClusterDetails{
 				{
@@ -58,20 +64,17 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		return iam.User{
+		return testUser, http.StatusOK
-			Id:       "1000012345",
-			UserName: "tester@databricks.com",
-		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/workspace/get-status", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/workspace/get-status", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 
 		return fakeWorkspace.WorkspaceGetStatus(path)
 	})
 
-	server.Handle("POST /api/2.0/workspace/mkdirs", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST", "/api/2.0/workspace/mkdirs", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		request := workspace.Mkdirs{}
 		decoder := json.NewDecoder(r.Body)
 
@@ -83,13 +86,13 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceMkdirs(request)
 	})
 
-	server.Handle("GET /api/2.0/workspace/export", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/workspace/export", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 
 		return fakeWorkspace.WorkspaceExport(path)
 	})
 
-	server.Handle("POST /api/2.0/workspace/delete", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST", "/api/2.0/workspace/delete", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		path := r.URL.Query().Get("path")
 		recursiveStr := r.URL.Query().Get("recursive")
 		var recursive bool
@@ -103,8 +106,9 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceDelete(path, recursive)
 	})
 
-	server.Handle("POST /api/2.0/workspace-files/import-file/{path}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST", "/api/2.0/workspace-files/import-file/{path:.*}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		path := r.PathValue("path")
+		vars := mux.Vars(r)
+		path := vars["path"]
 
 		body := new(bytes.Buffer)
 		_, err := body.ReadFrom(r.Body)
@@ -115,14 +119,15 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.WorkspaceFilesImportFile(path, body.Bytes())
 	})
 
-	server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.1/unity-catalog/current-metastore-assignment", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return catalog.MetastoreAssignment{
 			DefaultCatalogName: "main",
 		}, http.StatusOK
 	})
 
-	server.Handle("GET /api/2.0/permissions/directories/{objectId}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.0/permissions/directories/{objectId}", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
-		objectId := r.PathValue("objectId")
+		vars := mux.Vars(r)
+		objectId := vars["objectId"]
 
 		return workspace.WorkspaceObjectPermissions{
 			ObjectId:   objectId,
@@ -140,7 +145,7 @@ func AddHandlers(server *testserver.Server) {
 		}, http.StatusOK
 	})
 
-	server.Handle("POST /api/2.1/jobs/create", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("POST", "/api/2.1/jobs/create", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		request := jobs.CreateJob{}
 		decoder := json.NewDecoder(r.Body)
 
@@ -152,15 +157,31 @@ func AddHandlers(server *testserver.Server) {
 		return fakeWorkspace.JobsCreate(request)
 	})
 
-	server.Handle("GET /api/2.1/jobs/get", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.1/jobs/get", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		jobId := r.URL.Query().Get("job_id")
 
 		return fakeWorkspace.JobsGet(jobId)
 	})
 
-	server.Handle("GET /api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+	server.Handle("GET", "/api/2.1/jobs/list", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
 		return fakeWorkspace.JobsList()
 	})
+
+	server.Handle("GET", "/oidc/.well-known/oauth-authorization-server", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+		return map[string]string{
+			"authorization_endpoint": server.URL + "oidc/v1/authorize",
+			"token_endpoint":         server.URL + "/oidc/v1/token",
+		}, http.StatusOK
+	})
+
+	server.Handle("POST", "/oidc/v1/token", func(fakeWorkspace *testserver.FakeWorkspace, r *http.Request) (any, int) {
+		return map[string]string{
+			"access_token": "oauth-token",
+			"expires_in":   "3600",
+			"scope":        "all-apis",
+			"token_type":   "Bearer",
+		}, http.StatusOK
+	})
 }
 
 func internalError(err error) (any, int) {

acceptance/telemetry/dummy-error/out.requests.txt

@@ -0,0 +1,12 @@
+{
+  "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\":\"SMALL_INT\",\"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\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/dummy-error/out.upload.txt

@@ -0,0 +1,2 @@
+error: Failed to upload telemetry logs: Endpoint not implemented.
+

acceptance/telemetry/dummy-error/output.txt

@@ -0,0 +1,5 @@
+
+>>> [CLI] telemetry dummy
+waiting for telemetry process to finish
+File ./telemetry.pid exists
+Process has ended

acceptance/telemetry/dummy-error/script

@@ -0,0 +1,15 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_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
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_file ./telemetry.pid
+wait_pid $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/dummy-error/test.toml

@@ -0,0 +1,13 @@
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "error_code": "ERROR_CODE",
+    "message": "Endpoint not implemented."
+}
+'''
+Response.StatusCode = 501
+
+[[Repls]]
+Old = 'execution_time_ms\\\":\d{1,5},'
+New = 'execution_time_ms\":\"SMALL_INT\",'

acceptance/telemetry/dummy/out.requests.txt

@@ -0,0 +1,12 @@
+{
+  "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\":\"SMALL_INT\",\"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\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/dummy/out.upload.txt

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

acceptance/telemetry/dummy/output.txt

@@ -0,0 +1,5 @@
+
+>>> [CLI] telemetry dummy
+waiting for telemetry process to finish
+File ./telemetry.pid exists
+Process has ended

acceptance/telemetry/dummy/script

@@ -0,0 +1,13 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt
+
+trace $CLI telemetry dummy
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_file ./telemetry.pid
+wait_pid $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/dummy/test.toml

@@ -0,0 +1,11 @@
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "numProtoSuccess": 2
+}
+'''
+
+[[Repls]]
+Old = 'execution_time_ms\\\":\d{1,5},'
+New = 'execution_time_ms\":\"SMALL_INT\",'

acceptance/telemetry/test.toml

@@ -0,0 +1,10 @@
+LocalOnly = true
+RecordRequests = true
+
+[[Repls]]
+Old = '17\d{11}'
+New = '"UNIX_TIME_MILLIS"'
+
+[[Repls]]
+Old = 'darwin|linux|windows'
+New = 'OS'

acceptance/telemetry/upload/out.requests.txt

@@ -0,0 +1,12 @@
+{
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "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/telemetry/upload/out.upload.txt

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

acceptance/telemetry/upload/output.txt

@@ -0,0 +1,2 @@
+
+>>> [CLI] telemetry upload

acceptance/telemetry/upload/script

@@ -0,0 +1,5 @@
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt
+
+# 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/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/telemetry/upload/test.toml

@@ -0,0 +1,7 @@
+[[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,26 @@ package root
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"os"
+	"os/exec"
+	"runtime"
+	"slices"
+	"strconv"
 	"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 +83,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 +104,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)
@@ -126,5 +135,110 @@ 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
 }
+
+// We want child telemetry processes to inherit environment variables like $HOME or $HTTPS_PROXY
+// because they influence auth resolution.
+func inheritEnvVars() []string {
+	base := os.Environ()
+	out := []string{}
+	authEnvVars := auth.EnvVars()
+
+	// Remove any existing auth environment variables. This is done because
+	// the CLI offers multiple modalities of configuring authentication like
+	// `--profile` or `DATABRICKS_CONFIG_PROFILE` or `profile: <profile>` in the
+	// bundle config file.
+	//
+	// Each of these modalities have different priorities and thus we don't want
+	// any auth configuration to piggyback into the child process environment.
+	//
+	// This is a precaution to avoid conflicting auth configurations being passed
+	// to the child telemetry process.
+	for _, v := range base {
+		k, _, found := strings.Cut(v, "=")
+		if !found {
+			continue
+		}
+		if slices.Contains(authEnvVars, k) {
+			continue
+		}
+		out = append(out, v)
+	}
+
+	return out
+}
+
+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 = inheritEnvVars()
+	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
+	}
+
+	if pidFilePath := env.Get(ctx, telemetry.PidFileEnvVar); pidFilePath != "" {
+		err = os.WriteFile(pidFilePath, []byte(strconv.Itoa(telemetryCmd.Process.Pid)), 0o644)
+		if err != nil {
+			log.Debugf(ctx, "failed to write telemetry worker PID file: %s", err)
+		}
+	}
+
+	_, 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
+}

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/databricks/databricks-sdk-go v0.57.0 // Apache 2.0
 	github.com/fatih/color v1.18.0 // MIT
 	github.com/google/uuid v1.6.0 // BSD-3-Clause
+	github.com/gorilla/mux v1.8.1 // BSD 3-Clause
 	github.com/hashicorp/go-version v1.7.0 // MPL 2.0
 	github.com/hashicorp/hc-install v0.9.1 // MPL 2.0
 	github.com/hashicorp/terraform-exec v0.22.0 // MPL 2.0

go.sum

@@ -97,6 +97,8 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfF
 github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
 github.com/googleapis/gax-go/v2 v2.12.4 h1:9gWcmF85Wvq4ryPFvGFaOgPIs1AQX0d0bcbGw4Z96qg=
 github.com/googleapis/gax-go/v2 v2.12.4/go.mod h1:KYEYLorsnIGDi/rPC8b5TdlB9kbKoFubselGIoBMCwI=
+github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
+github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
 github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
 github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB11/k=

libs/auth/env.go

@@ -24,3 +24,17 @@ func Env(cfg *config.Config) map[string]string {
 
 	return out
 }
+
+func EnvVars() []string {
+	out := []string{}
+
+	for _, attr := range config.ConfigAttributes {
+		if len(attr.EnvVars) == 0 {
+			continue
+		}
+
+		out = append(out, attr.EnvVars[0])
+	}
+
+	return out
+}

libs/auth/env_test.go

@@ -40,3 +40,40 @@ func TestAuthEnv(t *testing.T) {
 	out := Env(in)
 	assert.Equal(t, expected, out)
 }
+
+func TestAuthEnvVars(t *testing.T) {
+	expected := []string{
+		"DATABRICKS_HOST",
+		"DATABRICKS_CLUSTER_ID",
+		"DATABRICKS_WAREHOUSE_ID",
+		"DATABRICKS_SERVERLESS_COMPUTE_ID",
+		"DATABRICKS_METADATA_SERVICE_URL",
+		"DATABRICKS_ACCOUNT_ID",
+		"DATABRICKS_TOKEN",
+		"DATABRICKS_USERNAME",
+		"DATABRICKS_PASSWORD",
+		"DATABRICKS_CONFIG_PROFILE",
+		"DATABRICKS_CONFIG_FILE",
+		"DATABRICKS_GOOGLE_SERVICE_ACCOUNT",
+		"GOOGLE_CREDENTIALS",
+		"DATABRICKS_AZURE_RESOURCE_ID",
+		"ARM_USE_MSI",
+		"ARM_CLIENT_SECRET",
+		"ARM_CLIENT_ID",
+		"ARM_TENANT_ID",
+		"ACTIONS_ID_TOKEN_REQUEST_URL",
+		"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
+		"ARM_ENVIRONMENT",
+		"DATABRICKS_AZURE_LOGIN_APP_ID",
+		"DATABRICKS_CLIENT_ID",
+		"DATABRICKS_CLIENT_SECRET",
+		"DATABRICKS_CLI_PATH",
+		"DATABRICKS_AUTH_TYPE",
+		"DATABRICKS_DEBUG_TRUNCATE_BYTES",
+		"DATABRICKS_DEBUG_HEADERS",
+		"DATABRICKS_RATE_LIMIT",
+	}
+
+	out := EnvVars()
+	assert.Equal(t, expected, out)
+}

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"
+	"errors"
+)
+
+// 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(errors.New("telemetry logger not found in the context"))
+	}
+
+	return v.(*logger)
+}

libs/telemetry/logger.go

@@ -0,0 +1,50 @@
+package telemetry
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/google/uuid"
+)
+
+func Log(ctx context.Context, event protos.DatabricksCliLog) {
+	fromContext(ctx).log(event)
+}
+
+func GetLogs(ctx context.Context) []protos.FrontendLog {
+	return fromContext(ctx).getLogs()
+}
+
+func HasLogs(ctx context.Context) bool {
+	return len(fromContext(ctx).getLogs()) > 0
+}
+
+func SetExecutionContext(ctx context.Context, ec protos.ExecutionContext) {
+	fromContext(ctx).setExecutionContext(ec)
+}
+
+type logger struct {
+	logs []protos.FrontendLog
+}
+
+func (l *logger) log(event protos.DatabricksCliLog) {
+	if l.logs == nil {
+		l.logs = make([]protos.FrontendLog, 0)
+	}
+	l.logs = append(l.logs, protos.FrontendLog{
+		FrontendLogEventID: uuid.New().String(),
+		Entry: protos.FrontendLogEntry{
+			DatabricksCliLog: event,
+		},
+	})
+}
+
+func (l *logger) getLogs() []protos.FrontendLog {
+	return l.logs
+}
+
+func (l *logger) setExecutionContext(ec protos.ExecutionContext) {
+	for i := range l.logs {
+		l.logs[i].Entry.DatabricksCliLog.ExecutionContext = &ec
+	}
+}

libs/telemetry/protos/databricks_cli_log.go

@@ -24,10 +24,12 @@ type ExecutionContext struct {
 	FromWebTerminal bool `json:"from_web_terminal,omitempty"`
 
 	// Time taken for the CLI command to execute.
-	ExecutionTimeMs int64 `json:"execution_time_ms,omitempty"`
+	// We want to serialize the zero value as well so the omitempty tag is not set.
+	ExecutionTimeMs int64 `json:"execution_time_ms"`
 
 	// Exit code of the CLI command.
-	ExitCode int64 `json:"exit_code,omitempty"`
+	// We want to serialize the zero value as well so the omitempty tag is not set.
+	ExitCode int64 `json:"exit_code"`
 }
 
 type CliTestEvent struct {

libs/telemetry/upload.go

@@ -0,0 +1,99 @@
+package telemetry
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/databricks/databricks-sdk-go/client"
+	"github.com/databricks/databricks-sdk-go/config"
+)
+
+const (
+	// File containing debug logs from the upload process.
+	UploadLogsFileEnvVar = "DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE"
+
+	// File containing the PID of the telemetry upload process.
+	PidFileEnvVar = "DATABRICKS_CLI_TELEMETRY_PID_FILE"
+)
+
+type UploadConfig struct {
+	Logs []protos.FrontendLog `json:"logs"`
+}
+
+// Upload reads telemetry logs from stdin and uploads them to the telemetry endpoint.
+// This function is always expected to be called in a separate child process from
+// the main CLI process.
+func Upload() (*ResponseBody, error) {
+	var err error
+
+	b, err := io.ReadAll(os.Stdin)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read from stdin: %s\n", err)
+	}
+
+	in := UploadConfig{}
+	err = json.Unmarshal(b, &in)
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal input: %s\n", err)
+	}
+
+	if len(in.Logs) == 0 {
+		return nil, fmt.Errorf("No logs to upload: %s\n", err)
+	}
+
+	protoLogs := make([]string, len(in.Logs))
+	for i, log := range in.Logs {
+		b, err := json.Marshal(log)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal log: %s\n", err)
+		}
+		protoLogs[i] = string(b)
+	}
+
+	// Parent process is responsible for setting environment variables to
+	// configure authentication.
+	apiClient, err := client.New(&config.Config{})
+	if err != nil {
+		return nil, fmt.Errorf("Failed to create API client: %s\n", err)
+	}
+
+	// Set a maximum total time to try telemetry uploads.
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	resp := &ResponseBody{}
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, errors.New("Failed to flush telemetry log due to timeout")
+
+		default:
+			// Proceed
+		}
+
+		// Log the CLI telemetry events.
+		err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", nil, nil, RequestBody{
+			UploadTime: time.Now().UnixMilli(),
+			Items:      []string{},
+			ProtoLogs:  protoLogs,
+		}, resp)
+		if err != nil {
+			return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", err)
+		}
+
+		if len(resp.Errors) > 0 {
+			return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", resp.Errors)
+		}
+
+		if resp.NumProtoSuccess == int64(len(in.Logs)) {
+			return resp, nil
+		}
+	}
+}

libs/telemetry/upload_test.go

@@ -0,0 +1,85 @@
+package telemetry
+
+import (
+	"encoding/json"
+	"net/http"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/databricks/cli/libs/testserver"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTelemetryUpload(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	count := 0
+	server.Handle("POST", "/telemetry-ext", func(_ *testserver.FakeWorkspace, req *http.Request) (resp any, statusCode int) {
+		count++
+		if count == 1 {
+			return ResponseBody{
+				NumProtoSuccess: 1,
+			}, http.StatusOK
+		}
+		if count == 2 {
+			return ResponseBody{
+				NumProtoSuccess: 2,
+			}, http.StatusOK
+		}
+		return nil, http.StatusInternalServerError
+	})
+
+	t.Setenv("DATABRICKS_HOST", server.URL)
+	t.Setenv("DATABRICKS_TOKEN", "token")
+
+	logs := []protos.FrontendLog{
+		{
+			FrontendLogEventID: uuid.New().String(),
+			Entry: protos.FrontendLogEntry{
+				DatabricksCliLog: protos.DatabricksCliLog{
+					CliTestEvent: &protos.CliTestEvent{Name: protos.DummyCliEnumValue1},
+				},
+			},
+		},
+		{
+			FrontendLogEventID: uuid.New().String(),
+			Entry: protos.FrontendLogEntry{
+				DatabricksCliLog: protos.DatabricksCliLog{
+					CliTestEvent: &protos.CliTestEvent{Name: protos.DummyCliEnumValue2},
+				},
+			},
+		},
+	}
+
+	processIn := UploadConfig{
+		Logs: logs,
+	}
+
+	b, err := json.Marshal(processIn)
+	require.NoError(t, err)
+
+	tmpDir := t.TempDir()
+	testutil.WriteFile(t, filepath.Join(tmpDir, "stdin"), string(b))
+
+	fd, err := os.OpenFile(filepath.Join(tmpDir, "stdin"), os.O_RDONLY, 0o644)
+	require.NoError(t, err)
+
+	// Redirect stdin to the file containing the telemetry logs.
+	old := os.Stdin
+	os.Stdin = fd
+	t.Cleanup(func() {
+		fd.Close()
+		os.Stdin = old
+	})
+
+	resp, err := Upload()
+	require.NoError(t, err)
+	assert.Equal(t, int64(2), resp.NumProtoSuccess)
+	assert.Equal(t, 2, count)
+}

libs/testserver/server.go

@@ -9,6 +9,8 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/gorilla/mux"
+
 	"github.com/stretchr/testify/assert"
 
 	"github.com/databricks/cli/internal/testutil"
@@ -17,7 +19,7 @@ import (
 
 type Server struct {
 	*httptest.Server
-	Mux *http.ServeMux
+	Router *mux.Router
 
 	t testutil.TestingT
 
@@ -34,26 +36,25 @@ type Request struct {
 	Headers http.Header `json:"headers,omitempty"`
 	Method  string      `json:"method"`
 	Path    string      `json:"path"`
-	Body    any         `json:"body"`
+	Body    any         `json:"body,omitempty"`
+	RawBody string      `json:"raw_body,omitempty"`
 }
 
 func New(t testutil.TestingT) *Server {
-	mux := http.NewServeMux()
+	router := mux.NewRouter()
-	server := httptest.NewServer(mux)
+	server := httptest.NewServer(router)
 	t.Cleanup(server.Close)
 
 	s := &Server{
 		Server:         server,
-		Mux:            mux,
+		Router:         router,
 		t:              t,
 		mu:             &sync.Mutex{},
 		fakeWorkspaces: map[string]*FakeWorkspace{},
 	}
 
-	// The server resolves conflicting handlers by using the one with higher
+	// Set up the not found handler as fallback
-	// specificity. This handler is the least specific, so it will be used as a
+	router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-	// fallback when no other handlers match.
-	s.Handle("/", func(fakeWorkspace *FakeWorkspace, r *http.Request) (any, int) {
 		pattern := r.Method + " " + r.URL.Path
 
 		t.Errorf(`
@@ -74,9 +75,22 @@ Response.StatusCode = <response status-code here>
 
 `, pattern, pattern)
 
-		return apierr.APIError{
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusNotImplemented)
+
+		resp := apierr.APIError{
 			Message: "No stub found for pattern: " + pattern,
-		}, http.StatusNotImplemented
+		}
+
+		respBytes, err := json.Marshal(resp)
+		if err != nil {
+			t.Errorf("JSON encoding error: %s", err)
+			respBytes = []byte("{\"message\": \"JSON encoding error\"}")
+		}
+
+		if _, err := w.Write(respBytes); err != nil {
+			t.Errorf("Response write error: %s", err)
+		}
 	})
 
 	return s
@@ -84,8 +98,8 @@ Response.StatusCode = <response status-code here>
 
 type HandlerFunc func(fakeWorkspace *FakeWorkspace, req *http.Request) (resp any, statusCode int)
 
-func (s *Server) Handle(pattern string, handler HandlerFunc) {
+func (s *Server) Handle(method, path string, handler HandlerFunc) {
-	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
+	s.Router.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
 		// For simplicity we process requests sequentially. It's fast enough because
 		// we don't do any IO except reading and writing request/response bodies.
 		s.mu.Lock()
@@ -119,13 +133,19 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) {
 				}
 			}
 
-			s.Requests = append(s.Requests, Request{
+			req := Request{
 				Headers: headers,
 				Method:  r.Method,
 				Path:    r.URL.Path,
-				Body:    json.RawMessage(body),
+			}
-			})
 
+			if json.Valid(body) {
+				req.Body = json.RawMessage(body)
+			} else {
+				req.RawBody = string(body)
+			}
+
+			s.Requests = append(s.Requests, req)
 		}
 
 		w.Header().Set("Content-Type", "application/json")
@@ -149,7 +169,7 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
-	})
+	}).Methods(method)
 }
 
 func getToken(r *http.Request) string {

main.go

@@ -2,14 +2,53 @@ package main
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
+	"io"
 	"os"
 
 	"github.com/databricks/cli/cmd"
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/telemetry"
 )
 
 func main() {
 	ctx := context.Background()
+
+	// Uploading telemetry data spawns a new process. We handle this separately
+	// from the rest of the CLI commands.
+	// This is done because [root.Execute] spawns a new process to run the
+	// "telemetry upload" command if there are logs to be uploaded. Having this outside
+	// of [root.Execute] ensures that the telemetry upload process is not spawned
+	// infinitely in a recursive manner.
+	if len(os.Args) == 3 && os.Args[1] == "telemetry" && os.Args[2] == "upload" {
+		var err error
+
+		// By default, this command should not write anything to stdout or stderr.
+		outW := io.Discard
+		errW := io.Discard
+
+		// If the environment variable is set, redirect stdout to the file.
+		// This is useful for testing.
+		if v := os.Getenv(telemetry.UploadLogsFileEnvVar); v != "" {
+			outW, _ = os.OpenFile(v, os.O_CREATE|os.O_WRONLY, 0o644)
+			errW = outW
+		}
+
+		resp, err := telemetry.Upload()
+		if err != nil {
+			fmt.Fprintf(errW, "error: %s\n", err)
+			os.Exit(1)
+		}
+		fmt.Fprintf(outW, "Telemetry logs uploaded successfully\n")
+		fmt.Fprintln(outW, "Response:")
+		b, err := json.Marshal(resp)
+		if err == nil {
+			fmt.Fprintln(outW, string(b))
+		}
+		os.Exit(0)
+	}
+
 	err := root.Execute(ctx, cmd.New(ctx))
 	if err != nil {
 		os.Exit(1)