Add telemetry for SDK usage from DBR

client/client_test.go

@@ -320,6 +320,37 @@ GET /a
 	}
 }
 
+func TestUserAgentForDBR(t *testing.T) {
+	useragent.DisableRuntimeCaching = true
+	for _, dbrVersion := range []string{"client.0", "client.1", "15.0", "13.3", "14.4"} {
+		t.Run(dbrVersion, func(t *testing.T) {
+			t.Setenv("DATABRICKS_RUNTIME_VERSION", dbrVersion)
+
+			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)
+			assert.Contains(t, userAgent, "runtime/"+dbrVersion)
+			require.NoError(t, err)
+		})
+	}
+}
+
 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,7 @@ const (
 
 var regexpSemVer = regexp.MustCompile(`^` + semVerCore + semVerPrerelease + semVerBuildmetadata + `$`)
 
-var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_-]+$`)
+var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_\.-]+$`)
 
 func matchSemVer(s string) error {
 	if regexpSemVer.MatchString(s) {

useragent/patterns_test.go

@@ -14,19 +14,20 @@ 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))
+	}
 }

useragent/runtime.go

@@ -0,0 +1,29 @@
+package useragent
+
+import (
+	"os"
+	"sync"
+)
+
+var runtimeOnce sync.Once
+var runtimeVersion string
+
+// DisableRuntimeCaching is a flag to disable caching of runtime version.
+// This is useful for testing.
+var DisableRuntimeCaching bool
+
+func getRuntimeVersion() string {
+	v := os.Getenv("DATABRICKS_RUNTIME_VERSION")
+	return Sanitize(v)
+}
+
+func Runtime() string {
+	if DisableRuntimeCaching {
+		return getRuntimeVersion()
+	}
+
+	runtimeOnce.Do(func() {
+		runtimeVersion = getRuntimeVersion()
+	})
+	return runtimeVersion
+}

useragent/user_agent.go

@@ -4,13 +4,20 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"regexp"
 	"runtime"
 	"strings"
 
 	"github.com/databricks/databricks-sdk-go/version"
 	"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) {
@@ -99,6 +106,22 @@ func (u info) String() string {
 
 type data []info
 
+// 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
+func Sanitize(s string) string {
+	return regexp.MustCompile(`[^0-9A-Za-z_\.\+-]`).ReplaceAllString(s, "-")
+}
+
 // With always uses the latest value for a given alphanumeric key.
 // Panics if key or value don't satisfy alphanumeric or semver format.
 func (d data) With(key, value string) data {

useragent/user_agent_test.go

@@ -74,3 +74,20 @@ func TestFromContext_Custom(t *testing.T) {
 func TestDefaultsAreValid(t *testing.T) {
 	WithProduct(productName, productVersion)
 }
+
+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))
+	}
+}