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"
@@ -64,6 +65,17 @@ func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
 				*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
+				}
+				// Add the detected Databricks Runtime version to the user agent
+				ctx := useragent.InContext(r.Context(), "runtime", v)
+				*r = *r.WithContext(ctx) // replace request
+				return nil
+			},
 		},
 		TransientErrors: []string{
 			"REQUEST_LIMIT_EXCEEDED", // This is temporary workaround for SCIM API returning 500.  Remove when it's fixed

useragent/patterns.go

@@ -36,3 +36,11 @@ func matchAlphanumOrSemVer(s string) error {
 	}
 	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

@@ -85,13 +85,27 @@ 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 {
+	// 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 == "runtime" {
+		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,15 @@ 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"))
+}