Add telemetry for SDK usage from DBR

client/client_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/internal/env"
 	"github.com/databricks/databricks-sdk-go/useragent"
 	"github.com/databricks/databricks-sdk-go/version"
 	"github.com/stretchr/testify/assert"
@@ -320,6 +321,114 @@ GET /a
 	}
 }
 
+func captureUserAgent(t *testing.T) string {
+	var userAgent string
+	c, err := New(&config.Config{
+		Host:       "some",
+		Token:      "token",
+		ConfigFile: "/dev/null",
+		HTTPTransport: hc(func(r *http.Request) (*http.Response, error) {
+			// Capture the user agent via the round tripper.
+			userAgent = r.UserAgent()
+
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(strings.NewReader(`{}`)),
+				Request:    r,
+			}, nil
+		}),
+	})
+	require.NoError(t, err)
+
+	err = c.Do(context.Background(), "GET", "/a", nil, nil, nil)
+	require.NoError(t, err)
+
+	return userAgent
+}
+
+func TestUserAgentForDBR(t *testing.T) {
+	for v, sv := range map[string]string{
+		// DBR versions that don't need to be sanitized.
+		"client.0": "client.0",
+		"client.1": "client.1",
+		"15.5":     "15.5",
+		"15.5.0":   "15.5.0",
+		"13.3":     "13.3",
+
+		// DBR versions that need to be sanitized.
+		"foo🧟bar": "foo-bar",
+		"foo/bar": "foo-bar",
+		"foo bar": "foo-bar",
+	} {
+		t.Run(v, func(t *testing.T) {
+			env.CleanupEnvironment(t)
+			useragent.ClearCache()
+
+			t.Setenv("DATABRICKS_RUNTIME_VERSION", v)
+			userAgent := captureUserAgent(t)
+
+			// The user agent should contain the runtime version, with the value
+			// sanitized if necessary.
+			assert.Contains(t, userAgent, "runtime/"+sv)
+		})
+	}
+}
+
+func TestUserAgentForCiCd(t *testing.T) {
+	ciToEnv := map[string]map[string]string{
+		"github": {
+			"GITHUB_ACTIONS": "true",
+		},
+		"gitlab": {
+			"GITLAB_CI": "true",
+		},
+		"jenkins": {
+			"JENKINS_URL": "https://jenkins.example.com",
+		},
+		"azure-devops": {
+			"TF_BUILD": "True",
+		},
+		"circle": {
+			"CIRCLECI": "true",
+		},
+		"travis": {
+			"TRAVIS": "true",
+		},
+		"bitbucket": {
+			"BITBUCKET_BUILD_NUMBER": "123",
+		},
+		"google-cloud-build": {
+			"PROJECT_ID":     "",
+			"BUILD_ID":       "",
+			"PROJECT_NUMBER": "",
+			"LOCATION":       "",
+		},
+		"aws-code-build": {
+			"CODEBUILD_BUILD_ARN": "",
+		},
+		"tf-cloud": {
+			"TFC_RUN_ID": "",
+		},
+	}
+
+	for ci, envVars := range ciToEnv {
+		t.Run(ci, func(t *testing.T) {
+			env.CleanupEnvironment(t)
+			useragent.ClearCache()
+
+			for k, v := range envVars {
+				t.Setenv(k, v)
+			}
+
+			userAgent := captureUserAgent(t)
+
+			// The user agent should contain the CI/CD provider.
+			assert.Contains(t, userAgent, "cicd/"+ci)
+		})
+	}
+
+}
+
 func testNonJSONResponseIncludedInError(t *testing.T, statusCode int, status, errorMessage string) {
 	c, err := New(&config.Config{
 		Host:       "some",

config/api_client.go

@@ -49,7 +49,7 @@ func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
 				return nil
 			},
 			func(r *http.Request) error {
-				ctx := useragent.InContext(r.Context(), "auth", c.AuthType)
+				ctx := useragent.InContext(r.Context(), useragent.AuthKey, c.AuthType)
 				*r = *r.WithContext(ctx) // replace request
 				return nil
 			},
@@ -60,7 +60,18 @@ func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
 					return nil
 				}
 				// Add the detected CI/CD provider to the user agent
-				ctx := useragent.InContext(r.Context(), "cicd", provider)
+				ctx := useragent.InContext(r.Context(), useragent.CicdKey, provider)
+				*r = *r.WithContext(ctx) // replace request
+				return nil
+			},
+			func(r *http.Request) error {
+				// Detect if the SDK is being run in a Databricks Runtime.
+				v := useragent.Runtime()
+				if v == "" {
+					return nil
+				}
+				// Add the detected Databricks Runtime version to the user agent
+				ctx := useragent.InContext(r.Context(), useragent.RuntimeKey, v)
 				*r = *r.WithContext(ctx) // replace request
 				return nil
 			},

useragent/patterns.go

@@ -14,7 +14,24 @@ const (
 
 var regexpSemVer = regexp.MustCompile(`^` + semVerCore + semVerPrerelease + semVerBuildmetadata + `$`)
 
-var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_-]+$`)
+// Sanitize replaces all non-alphanumeric characters with a hyphen. Use this to
+// ensure that the user agent value is valid. This is useful when the value is not
+// ensured to be valid at compile time.
+//
+// Example: You want to avoid having '/' and ' ' in the value because it will
+// make downstream applications fail.
+//
+// Note: Semver strings are comprised of alphanumeric characters, hyphens, periods
+// and plus signs. This function will not remove these characters.
+// see:
+// 1. https://semver.org/#spec-item-9
+// 2. https://semver.org/#spec-item-10
+var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_\.\+-]+$`)
+var regexpAlphanumInverse = regexp.MustCompile(`[^0-9A-Za-z_\.\+-]`)
+
+func Sanitize(s string) string {
+	return regexpAlphanumInverse.ReplaceAllString(s, "-")
+}
 
 func matchSemVer(s string) error {
 	if regexpSemVer.MatchString(s) {

useragent/patterns_test.go

@@ -14,19 +14,65 @@ func TestMatchSemVer(t *testing.T) {
 }
 
 func TestMatchAlphanum(t *testing.T) {
-	assert.NoError(t, matchAlphanum("foo"))
-	assert.NoError(t, matchAlphanum("FOO"))
-	assert.NoError(t, matchAlphanum("FOO123"))
-	assert.NoError(t, matchAlphanum("foo_bar"))
-	assert.NoError(t, matchAlphanum("foo-bar"))
-	assert.Error(t, matchAlphanum("foo bar"))
-	assert.Error(t, matchAlphanum("foo/bar"))
+	for _, v := range []string{
+		"foo",
+		"FOO",
+		"FOO123",
+		"foo_bar",
+		"foo-bar",
+		"foo.bar",
+	} {
+		assert.NoError(t, matchAlphanum(v))
+	}
+
+	for _, v := range []string{
+		"foo bar",
+		"foo/bar",
+	} {
+		assert.Error(t, matchAlphanum(v))
+	}
 }
 
 func TestMatchAlphanumOrSemVer(t *testing.T) {
-	assert.NoError(t, matchAlphanumOrSemVer("foo"))
-	assert.NoError(t, matchAlphanumOrSemVer("1.2.3"))
-	assert.NoError(t, matchAlphanumOrSemVer("0.0.0-dev+2e014739024a"))
-	assert.Error(t, matchAlphanumOrSemVer("foo/bar"))
-	assert.Error(t, matchAlphanumOrSemVer("1/2/3"))
+	for _, v := range []string{
+		"foo",
+		"1.2.3",
+		"0.0.0-dev+2e014739024a",
+		"client.0",
+	} {
+		assert.NoError(t, matchAlphanumOrSemVer(v))
+	}
+	for _, v := range []string{
+		"foo/bar",
+		"1/2/3",
+	} {
+		assert.Error(t, matchAlphanumOrSemVer(v))
+	}
+}
+
+func TestSanitize(t *testing.T) {
+	for _, v := range []string{
+		"foo",
+		"FOO",
+		"FOO123",
+		"foo_bar",
+		"foo-bar",
+		"foo+bar",
+		"foo.bar",
+		"1.2.3",
+		"client.0",
+	} {
+		assert.Equal(t, v, Sanitize(v))
+	}
+
+	sanitizeMap := map[string]string{
+		"1@2#3?4,5/6!7 8 ": "1-2-3-4-5-6-7-8-",
+		"foo bar":          "foo-bar",
+		"foo/bar":          "foo-bar",
+		"foo:)bar":         "foo--bar",
+		"foo😊bar":          "foo-bar",
+	}
+	for k, v := range sanitizeMap {
+		assert.Equal(t, v, Sanitize(k))
+	}
 }

useragent/runtime.go

@@ -0,0 +1,30 @@
+package useragent
+
+import (
+	"os"
+	"sync"
+)
+
+// Clear cached user agent values like the DBR version or the CI/CD provider
+// being used. This is useful for testing.
+func ClearCache() {
+	// Reset the sync.Once to their default values. This will recompute the
+	// values on the next call to Runtime() or CiCdProvider().
+	runtimeOnce = sync.Once{}
+	providerOnce = sync.Once{}
+}
+
+var runtimeOnce sync.Once
+var runtimeVersion string
+
+func getRuntimeVersion() string {
+	v := os.Getenv("DATABRICKS_RUNTIME_VERSION")
+	return Sanitize(v)
+}
+
+func Runtime() string {
+	runtimeOnce.Do(func() {
+		runtimeVersion = getRuntimeVersion()
+	})
+	return runtimeVersion
+}

useragent/runtime_test.go

@@ -0,0 +1,26 @@
+package useragent
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestUserAgentRuntime(t *testing.T) {
+	for _, tc := range []string{
+		"1.2.3",
+		"0.0.0-dev+2e014739024a",
+		"client.0",
+		"foo",
+		"15.0",
+		"13.3",
+	} {
+		t.Run(tc, func(t *testing.T) {
+			t.Setenv("DATABRICKS_RUNTIME_VERSION", tc)
+
+			v := getRuntimeVersion()
+			assert.Equal(t, tc, v)
+			assert.NoError(t, matchAlphanumOrSemVer(v))
+		})
+	}
+}

useragent/user_agent.go

@@ -11,6 +11,12 @@ import (
 	"golang.org/x/mod/semver"
 )
 
+const (
+	RuntimeKey = "runtime"
+	CicdKey    = "cicd"
+	AuthKey    = "auth"
+)
+
 // WithProduct sets the product name and product version globally.
 // It should be called by developers to differentiate their application from others.
 func WithProduct(name, version string) {