Support partners in SDK

README.md

@@ -22,6 +22,7 @@ The Databricks SDK for Go includes functionality to accelerate development with
 - [GetByName utility methods](#getbyname-utility-methods)
 - [Node type and Databricks Runtime selectors](#node-type-and-databricks-runtime-selectors)
 - [Integration with `io` interfaces for DBFS](#integration-with-io-interfaces-for-dbfs)
+- [User Agent Request Attribution](#user-agent-request-attribution)
 - [Error Handling](#error-handling)
 - [Logging](#logging)
 - [Testing](#testing)
@@ -531,6 +532,25 @@ buf, err := w.Dbfs.ReadFile(ctx, "/path/to/remote/file")
 
 Databricks SDK for Go loosely integrates with [spf13/pflag](https://github.com/spf13/pflag) by implementing [pflag.Value](https://pkg.go.dev/github.com/spf13/pflag#Value) for all enum types.
 
+## User Agent Request Attribution
+
+The Databricks SDK for Go uses the `User-Agent` header to include request metadata along with each request. By default, this includes the version of the Go SDK, the version of the Go language used by your application, and the underlying operating system. To statically add additional metadata, you can use the `WithPartner()` and `WithProduct()` functions in the `useragent` package. `WithPartner()` can be used by partners to indicate that code using the Databricks SDK for Go should be attributed to a specific partner. Multiple partners can be registered at once. Partner names can contain any number, digit, `.`, `-`, `_` or `+`.
+
+```go
+useragent.WithPartner("partner-abc")
+useragent.WithPartner("partner-xyz")
+```
+
+`WithProduct()` can be used to define the name and version of the product that is built with the Databricks SDK for Go. The product name has the same restrictions as the partner name above, and the product version must be a valid [SemVer](https://semver.org/). Subsequent calls to `WithProduct()` replace the original product with the new user-specified one.
+
+```go
+useragent.WithProduct("databricks-example-product", "1.2.0")
+```
+
+If both the `DATABRICKS_SDK_UPSTREAM` and `DATABRICKS_SDK_UPSTREAM_VERSION` environment variables are defined, these will also be included in the `User-Agent` header.
+
+If additional metadata needs to be specified that isn't already supported by the above interfaces, you can use the `WithUserAgentExtra()` function to register arbitrary key-value pairs to include in the user agent. Multiple values associated with the same key are allowed. Keys have the same restrictions as the partner name above. Values must be either as described above or SemVer strings.
+
 ## Error handling
 
 The Databricks SDK for Go converts error responses from the Databricks API into the [`apierr.APIError`](https://pkg.go.dev/github.com/databricks/databricks-sdk-go/apierr#APIError) type. This allows you to inspect the error code, message, and details by asserting the error as `apierr.APIError`:

examples/useragent/main.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/useragent"
+)
+
+// Run this example by running `go run ./examples/useragent` from the root of this repository.
+//
+// To run this example, ensure that you have a DEFAULT profile configured in your `.databrickscfg`.
+func main() {
+	useragent.WithProduct("databricks-sdk-example", "0.0.1")
+	useragent.WithPartner("databricks-sdk-example-abc")
+	useragent.WithPartner("databricks-sdk-example-def")
+	useragent.WithUserAgentExtra("test-key", "test-value")
+	w := databricks.Must(databricks.NewWorkspaceClient())
+	_, err := w.CurrentUser.Me(context.Background())
+	if err != nil {
+		panic(err)
+	}
+}

useragent/user_agent.go

@@ -52,6 +52,10 @@ func WithUserAgentExtra(key, value string) {
 	extra = extra.With(key, value)
 }
 
+func WithPartner(partner string) {
+	WithUserAgentExtra("partner", partner)
+}
+
 func getUpstreamUserAgentInfo() []info {
 	product := os.Getenv("DATABRICKS_SDK_UPSTREAM")
 	version := os.Getenv("DATABRICKS_SDK_UPSTREAM_VERSION")
@@ -114,19 +118,7 @@ func (d data) With(key, value string) data {
 	if err := matchAlphanumOrSemVer(value); err != nil {
 		panic(err)
 	}
-	var c data
-	var found bool
-	for _, i := range d {
-		if i.Key == key {
-			i.Value = value
-			found = true
-		}
-		c = append(c, i)
-	}
-	if !found {
-		c = append(c, info{key, value})
-	}
-	return c
+	return append(d, info{key, value})
 }
 
 func (d data) String() string {

useragent/user_agent_test.go

@@ -49,11 +49,11 @@ func TestFromContext_Custom(t *testing.T) {
 	userAgent2 := FromContext(ctx2)
 
 	// tests may be developed and run on different versions of different things
-	expectedFormat := "unit-tests/0.0.1 databricks-sdk-go/%s go/%s os/%s pulumi/3.8.4 terraform/1.3.6 a/e"
+	expectedFormat := "unit-tests/0.0.1 databricks-sdk-go/%s go/%s os/%s pulumi/3.8.4 terraform/1.3.4 terraform/1.3.5 terraform/1.3.6 a/b a/d a/e"
 	expected := fmt.Sprintf(expectedFormat, version.Version, goVersion(), runtime.GOOS)
 	assert.Equal(t, expected, userAgent)
 
-	expectedFormat2 := "unit-tests/0.0.1 databricks-sdk-go/%s go/%s os/%s pulumi/3.8.4 terraform/1.3.6 a/f foo/bar"
+	expectedFormat2 := "unit-tests/0.0.1 databricks-sdk-go/%s go/%s os/%s pulumi/3.8.4 terraform/1.3.4 terraform/1.3.5 terraform/1.3.6 a/b a/d a/e a/f foo/bar"
 	expected2 := fmt.Sprintf(expectedFormat2, version.Version, goVersion(), runtime.GOOS)
 	assert.Equal(t, expected2, userAgent2)
 
@@ -74,3 +74,11 @@ func TestFromContext_Custom(t *testing.T) {
 func TestDefaultsAreValid(t *testing.T) {
 	WithProduct(productName, productVersion)
 }
+
+func TestMultiplePartners(t *testing.T) {
+	WithPartner("partner1")
+	WithPartner("partner2")
+	userAgent := FromContext(context.Background())
+	assert.Contains(t, userAgent, "partner/partner1")
+	assert.Contains(t, userAgent, "partner/partner2")
+}