Add telemetry for SDK usage from DBR

client/client_test.go

@@ -320,6 +320,36 @@ GET /a
 	}
 }
 
+func TestUserAgentForDBR(t *testing.T) {
+	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

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"os"
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
@@ -49,7 +50,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 +61,19 @@ 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, ok := os.LookupEnv("DATABRICKS_RUNTIME_VERSION")
+				if !ok {
+					return nil
+				}
+				v = useragent.Sanitize(v)
+				// 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,25 +14,28 @@ const (
 
 var regexpSemVer = regexp.MustCompile(`^` + semVerCore + semVerPrerelease + semVerBuildmetadata + `$`)
 
-var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_-]+$`)
-
 func matchSemVer(s string) error {
 	if regexpSemVer.MatchString(s) {
 		return nil
 	}
 	return fmt.Errorf("invalid semver string: %s", s)
 }
 
-func matchAlphanum(s string) error {
-	if regexpAlphanum.MatchString(s) {
-		return nil
-	}
-	return fmt.Errorf("invalid alphanumeric string: %s", s)
+// Alphanumeric characters, hyphen, underscore, and period. This is the subset of
+// characters that we allow in user agent keys and values. This is to ensure that
+// downstream applications can correctly parse the full user agent header.
+//
+// NOTE: HTTP headers in general only work well with ASCII characters. see:
+// https://stackoverflow.com/questions/4400678/what-character-encoding-should-i-use-for-a-http-header
+var validChars = `0-9A-Za-z_\-\.`
+
+func isValid(s string) bool {
+	return regexp.MustCompile(`^[` + validChars + `]+$`).MatchString(s)
 }
 
-func matchAlphanumOrSemVer(s string) error {
-	if regexpAlphanum.MatchString(s) || regexpSemVer.MatchString(s) {
+func matchValidChars(s string) error {
+	if isValid(s) {
 		return nil
 	}
-	return fmt.Errorf("invalid alphanumeric or semver string: %s", s)
+	return fmt.Errorf("invalid alphanumeric string: %s", s)
 }

useragent/patterns_test.go

@@ -13,20 +13,28 @@ func TestMatchSemVer(t *testing.T) {
 	assert.Error(t, matchSemVer("1.2"))
 }
 
-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"))
-}
+func TestMatchValidChars(t *testing.T) {
+	assert.NoError(t, matchValidChars("foo"))
+	assert.True(t, isValid("foo"))
+
+	assert.NoError(t, matchValidChars("FOO"))
+	assert.True(t, isValid("FOO"))
+
+	assert.NoError(t, matchValidChars("FOO123"))
+	assert.True(t, isValid("FOO123"))
+
+	assert.NoError(t, matchValidChars("foo_bar"))
+	assert.True(t, isValid("foo_bar"))
+
+	assert.NoError(t, matchValidChars("foo-bar"))
+	assert.True(t, isValid("foo-bar"))
+
+	assert.NoError(t, matchValidChars("foo.bar"))
+	assert.True(t, isValid("foo.bar"))
+
+	assert.Error(t, matchValidChars("foo bar"))
+	assert.False(t, isValid("foo bar"))
 
-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"))
+	assert.Error(t, matchValidChars("foo/bar"))
+	assert.False(t, isValid("foo/bar"))
 }

useragent/user_agent.go

@@ -4,17 +4,24 @@ 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) {
-	if err := matchAlphanum(name); err != nil {
+	if err := matchValidChars(name); err != nil {
 		panic(err)
 	}
 	if err := matchSemVer(version); err != nil {
@@ -99,13 +106,29 @@ func (u info) String() string {
 
 type data []info
 
+// Validate the key value pair being set in the user agent. Error if invalid.
+func validate(key, value string) error {
+	if !isValid(key) {
+		return fmt.Errorf("expected user agent key to be alphanumeric: %q", key)
+	}
+	if !isValid(value) {
+		return fmt.Errorf("expected user agent value for key %q to be alphanumeric: %q", key, value)
+	}
+	return nil
+}
+
+// Sanitize the user agent value. This is useful when the value is not ensured to be
+// to be valid at compile time. Having this sanitization then ensures downstream
+// applications can correctly parse the full user agent header, by making sure
+// characters like '/' and ' ' are not present in the value.
+func Sanitize(s string) string {
+	return regexp.MustCompile(`[^`+validChars+`]`).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 {
-	if err := matchAlphanum(key); err != nil {
-		panic(err)
-	}
-	if err := matchAlphanumOrSemVer(value); err != nil {
+	if err := validate(key, value); err != nil {
 		panic(err)
 	}
 	var c data

useragent/user_agent_test.go

@@ -74,3 +74,56 @@ func TestFromContext_Custom(t *testing.T) {
 func TestDefaultsAreValid(t *testing.T) {
 	WithProduct(productName, productVersion)
 }
+
+func TestUserAgentValidate(t *testing.T) {
+	assert.EqualError(t, validate("foobar!", ""), "expected user agent key to be alphanumeric: \"foobar!\"")
+	assert.EqualError(t, validate("foo", "invalid!"), "expected user agent value for key \"foo\" to be alphanumeric: \"invalid!\"")
+	assert.EqualError(t, validate("foo", "whatever#!@"), "expected user agent value for key \"foo\" to be alphanumeric: \"whatever#!@\"")
+
+	assert.NoError(t, validate("foo", "7.3"))
+	assert.NoError(t, validate("foo", "client.7"))
+	assert.NoError(t, validate("foo", "1.1.1"))
+}
+
+func TestSanitize(t *testing.T) {
+	// Valid values
+	assert.True(t, isValid("foo"))
+	assert.Equal(t, "foo", Sanitize("foo"))
+
+	assert.True(t, isValid("FOO"))
+	assert.Equal(t, "FOO", Sanitize("FOO"))
+
+	assert.True(t, isValid("FOO123"))
+	assert.Equal(t, "FOO123", Sanitize("FOO123"))
+
+	assert.True(t, isValid("foo_bar"))
+	assert.Equal(t, "foo_bar", Sanitize("foo_bar"))
+
+	assert.True(t, isValid("foo-bar"))
+	assert.Equal(t, "foo-bar", Sanitize("foo-bar"))
+
+	assert.True(t, isValid("foo.bar"))
+	assert.Equal(t, "foo.bar", Sanitize("foo.bar"))
+
+	assert.True(t, isValid("1.2.3"))
+	assert.Equal(t, "1.2.3", Sanitize("1.2.3"))
+
+	assert.True(t, isValid("client.0"))
+	assert.Equal(t, "client.0", Sanitize("client.0"))
+
+	// Invalid Values, being sanitized correctly.
+	assert.False(t, isValid("1@2#3?4,5/6!7 8 "))
+	assert.Equal(t, "1-2-3-4-5-6-7-8-", Sanitize("1@2#3?4,5/6!7 8 "))
+
+	assert.False(t, isValid("foo bar"))
+	assert.Equal(t, "foo-bar", Sanitize("foo bar"))
+
+	assert.False(t, isValid("foo/bar"))
+	assert.Equal(t, "foo-bar", Sanitize("foo/bar"))
+
+	assert.False(t, isValid("foo:)bar"))
+	assert.Equal(t, "foo--bar", Sanitize("foo:)bar"))
+
+	assert.False(t, isValid("foo😊bar"))
+	assert.Equal(t, "foo-bar", Sanitize("foo😊bar"))
+}