Let caller set BRICKS_UPSTREAM for user agent (#196)

Example when called from vscode (and everything is hooked up):

```
> * User-Agent: bricks/0.0.21-devel databricks-sdk-go/0.2.0 go/1.19.4 os/darwin upstream/databricks-vscode
```
This commit is contained in:
Pieter Noordhuis 2023-02-03 17:05:58 +01:00 committed by GitHub
parent 9ca7f8a888
commit 2e01473902
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 82 additions and 1 deletions

View File

@ -23,8 +23,11 @@ var RootCmd = &cobra.Command{
SilenceUsage: true, SilenceUsage: true,
PersistentPreRun: func(cmd *cobra.Command, args []string) { PersistentPreRun: func(cmd *cobra.Command, args []string) {
ctx := cmd.Context()
// Configure our user agent with the command that's about to be executed. // Configure our user agent with the command that's about to be executed.
ctx := withCommandInUserAgent(cmd.Context(), cmd) ctx = withCommandInUserAgent(ctx, cmd)
ctx = withUpstreamInUserAgent(ctx)
cmd.SetContext(ctx) cmd.SetContext(ctx)
if Verbose { if Verbose {

View File

@ -0,0 +1,33 @@
package root
import (
"context"
"os"
"github.com/databricks/databricks-sdk-go/useragent"
)
// Environment variables that caller can set to convey what is upstream to bricks.
const upstreamEnvVar = "BRICKS_UPSTREAM"
const upstreamVersionEnvVar = "BRICKS_UPSTREAM_VERSION"
// Keys in the user agent.
const upstreamKey = "upstream"
const upstreamVersionKey = "upstream-version"
func withUpstreamInUserAgent(ctx context.Context) context.Context {
value := os.Getenv(upstreamEnvVar)
if value == "" {
return ctx
}
ctx = useragent.InContext(ctx, upstreamKey, value)
// Include upstream version as well, if set.
value = os.Getenv(upstreamVersionEnvVar)
if value == "" {
return ctx
}
return useragent.InContext(ctx, upstreamVersionKey, value)
}

View File

@ -0,0 +1,45 @@
package root
import (
"context"
"testing"
"github.com/databricks/databricks-sdk-go/useragent"
"github.com/stretchr/testify/assert"
)
func TestUpstreamSet(t *testing.T) {
t.Setenv(upstreamEnvVar, "foobar")
ctx := withUpstreamInUserAgent(context.Background())
assert.Contains(t, useragent.FromContext(ctx), "upstream/foobar")
}
func TestUpstreamSetEmpty(t *testing.T) {
t.Setenv(upstreamEnvVar, "")
ctx := withUpstreamInUserAgent(context.Background())
assert.NotContains(t, useragent.FromContext(ctx), "upstream/")
}
func TestUpstreamVersionSet(t *testing.T) {
t.Setenv(upstreamEnvVar, "foobar")
t.Setenv(upstreamVersionEnvVar, "0.0.1")
ctx := withUpstreamInUserAgent(context.Background())
assert.Contains(t, useragent.FromContext(ctx), "upstream/foobar")
assert.Contains(t, useragent.FromContext(ctx), "upstream-version/0.0.1")
}
func TestUpstreamVersionSetEmpty(t *testing.T) {
t.Setenv(upstreamEnvVar, "foobar")
t.Setenv(upstreamVersionEnvVar, "")
ctx := withUpstreamInUserAgent(context.Background())
assert.Contains(t, useragent.FromContext(ctx), "upstream/foobar")
assert.NotContains(t, useragent.FromContext(ctx), "upstream-version/")
}
func TestUpstreamVersionSetUpstreamNotSet(t *testing.T) {
t.Setenv(upstreamEnvVar, "")
t.Setenv(upstreamVersionEnvVar, "0.0.1")
ctx := withUpstreamInUserAgent(context.Background())
assert.NotContains(t, useragent.FromContext(ctx), "upstream/")
assert.NotContains(t, useragent.FromContext(ctx), "upstream-version/")
}