Add telemetry for SDK usage from DBR

client/client_test.go

@@ -320,6 +320,64 @@ GET /a
 	}
 }
 
+func TestUserAgentForDBR(t *testing.T) {
+	goVersion := strings.TrimPrefix(runtime.Version(), "go")
+	cicdHeader := ""
+	if useragent.CiCdProvider() != "" {
+		cicdHeader = fmt.Sprintf(" cicd/%s", useragent.CiCdProvider())
+	}
+	t.Setenv("DATABRICKS_RUNTIME_VERSION", "15.1")
+	expectedUserAgent := `unknown/0.0.0 databricks-sdk-go/` + version.Version + ` go/` + goVersion + ` os/` + runtime.GOOS + ` auth/pat` + cicdHeader + ` runtime/15.1`
+
+	var userAgent string
+	c, err := New(&config.Config{
+		Host:       "some",
+		Token:      "token",
+		ConfigFile: "/dev/null",
+		HTTPTransport: hc(func(r *http.Request) (*http.Response, error) {
+			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.Equal(t, userAgent, expectedUserAgent)
+	require.NoError(t, err)
+}
+
+func TestUserAgentForServerless(t *testing.T) {
+	goVersion := strings.TrimPrefix(runtime.Version(), "go")
+	cicdHeader := ""
+	if useragent.CiCdProvider() != "" {
+		cicdHeader = fmt.Sprintf(" cicd/%s", useragent.CiCdProvider())
+	}
+	t.Setenv("DATABRICKS_RUNTIME_VERSION", "client.1")
+	expectedUserAgent := `unknown/0.0.0 databricks-sdk-go/` + version.Version + ` go/` + goVersion + ` os/` + runtime.GOOS + ` auth/pat` + cicdHeader + ` runtime/client.1`
+
+	var userAgent string
+	c, err := New(&config.Config{
+		Host:       "some",
+		Token:      "token",
+		ConfigFile: "/dev/null",
+		HTTPTransport: hc(func(r *http.Request) (*http.Response, error) {
+			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.Equal(t, userAgent, expectedUserAgent)
+	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

@@ -17,22 +17,30 @@ var regexpSemVer = regexp.MustCompile(`^` + semVerCore + semVerPrerelease + semV
 var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_-]+$`)
 
 func matchSemVer(s string) error {
-	if regexpSemVer.MatchString(s) {
+	if isSemVer(s) {
 		return nil
 	}
 	return fmt.Errorf("invalid semver string: %s", s)
 }
 
 func matchAlphanum(s string) error {
-	if regexpAlphanum.MatchString(s) {
+	if isAlphanum(s) {
 		return nil
 	}
 	return fmt.Errorf("invalid alphanumeric string: %s", s)
 }
 
 func matchAlphanumOrSemVer(s string) error {
-	if regexpAlphanum.MatchString(s) || regexpSemVer.MatchString(s) {
+	if isAlphanum(s) || isSemVer(s) {
 		return nil
 	}
 	return fmt.Errorf("invalid alphanumeric or semver string: %s", s)
 }
+
+func isAlphanum(s string) bool {
+	return regexpAlphanum.MatchString(s)
+}
+
+func isSemVer(s string) bool {
+	return regexpSemVer.MatchString(s)
+}

useragent/patterns_test.go

@@ -8,25 +8,54 @@ import (
 
 func TestMatchSemVer(t *testing.T) {
 	assert.NoError(t, matchSemVer("1.2.3"))
+	assert.True(t, isSemVer("1.2.3"))
+
 	assert.NoError(t, matchSemVer("0.0.0-dev+2e014739024a"))
+	assert.True(t, isSemVer("0.0.0-dev+2e014739024a"))
+
 	assert.Error(t, matchSemVer("1.2.3.4"))
+	assert.False(t, isSemVer("1.2.3.4"))
+
 	assert.Error(t, matchSemVer("1.2"))
+	assert.False(t, isSemVer("1.2"))
 }
 
 func TestMatchAlphanum(t *testing.T) {
 	assert.NoError(t, matchAlphanum("foo"))
+	assert.True(t, isAlphanum("foo"))
+
 	assert.NoError(t, matchAlphanum("FOO"))
+	assert.True(t, isAlphanum("FOO"))
+
 	assert.NoError(t, matchAlphanum("FOO123"))
+	assert.True(t, isAlphanum("FOO123"))
+
 	assert.NoError(t, matchAlphanum("foo_bar"))
+	assert.True(t, isAlphanum("foo_bar"))
+
 	assert.NoError(t, matchAlphanum("foo-bar"))
+	assert.True(t, isAlphanum("foo-bar"))
+
 	assert.Error(t, matchAlphanum("foo bar"))
+	assert.False(t, isAlphanum("foo bar"))
+
 	assert.Error(t, matchAlphanum("foo/bar"))
+	assert.False(t, isAlphanum("foo/bar"))
 }
 
 func TestMatchAlphanumOrSemVer(t *testing.T) {
 	assert.NoError(t, matchAlphanumOrSemVer("foo"))
+	assert.True(t, isAlphanum("foo") || isSemVer("foo"))
+
 	assert.NoError(t, matchAlphanumOrSemVer("1.2.3"))
+	assert.True(t, isAlphanum("1.2.3") || isSemVer("1.2.3"))
+
 	assert.NoError(t, matchAlphanumOrSemVer("0.0.0-dev+2e014739024a"))
+	assert.True(t, isAlphanum("0.0.0-dev+2e014739024a") || isSemVer("0.0.0-dev+2e014739024a"))
+
 	assert.Error(t, matchAlphanumOrSemVer("foo/bar"))
+	assert.False(t, isAlphanum("foo/bar") || isSemVer("foo/bar"))
+
 	assert.Error(t, matchAlphanumOrSemVer("1/2/3"))
+	assert.False(t, isAlphanum("1/2/3") || isSemVer("1/2/3"))
 }

useragent/user_agent.go

@@ -5,11 +5,18 @@ import (
 	"fmt"
 	"runtime"
 	"strings"
+	"unicode"
 
 	"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) {
@@ -79,19 +86,48 @@ type info struct {
 	Value string
 }
 
+// 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 {
+	allowList := func(r rune) rune {
+		if r == '.' || unicode.IsLetter(r) || unicode.IsDigit(r) {
+			return r
+		}
+		return '-'
+	}
+
+	return strings.Map(allowList, s)
+}
+
 func (u info) String() string {
 	return fmt.Sprintf("%s/%s", u.Key, u.Value)
 }
 
 type data []info
 
+// Validate the key value pair being set in the user agent. Error if invalid.
+func validate(key, value string) error {
+	// DBR versions as set in the `DATABRICKS_RUNTIME_VERSION` environment variable
+	// are not valid semver strings. Eg: 15.1 or client.0. Thus we allow arbitrary
+	// values for the runtime key.
+	if key == RuntimeKey {
+		return nil
+	}
+	if !isAlphanum(key) {
+		return fmt.Errorf("expected user agent key to be alphanumeric: %s", key)
+	}
+	if !isAlphanum(value) && !isSemVer(value) {
+		return fmt.Errorf("expected user agent value for key %q to be alphanumeric or semver: %s", key, value)
+	}
+	return nil
+}
+
 // 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

@@ -60,3 +60,22 @@ func TestFromContext_Custom(t *testing.T) {
 func TestDefaultsAreValid(t *testing.T) {
 	WithProduct(productName, productVersion)
 }
+
+func TestUserAgentValidate(t *testing.T) {
+	assert.EqualError(t, validate("non-alphanumic!", "abc"), "expected user agent key to be alphanumeric: non-alphanumic!")
+	assert.EqualError(t, validate("abc", "non-alphanumeric!"), "expected user agent value for key \"abc\" to be alphanumeric or semver: non-alphanumeric!")
+	assert.EqualError(t, validate("abc", "1.1.invalid"), "expected user agent value for key \"abc\" to be alphanumeric or semver: 1.1.invalid")
+
+	assert.NoError(t, validate("runtime", "7.3"))
+	assert.NoError(t, validate("runtime", "client.7"))
+	assert.NoError(t, validate("runtime", "whatever#!@"))
+	assert.NoError(t, validate("abc", "123"))
+	assert.NoError(t, validate("abc", "1.1.1"))
+}
+
+func TestUserAgentNormalizeString(t *testing.T) {
+	assert.Equal(t, "abc123", Sanitize("abc123"))
+	assert.Equal(t, "1-2-3-4-5-6-7-8-", Sanitize("1@2#3?4,5/6!7 8 "))
+	assert.Equal(t, "1.2.3", Sanitize("1.2.3"))
+	assert.Equal(t, "client.0", Sanitize("client.0"))
+}