mirror of https://github.com/databricks/cli.git
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:
parent
9ca7f8a888
commit
2e01473902
|
@ -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 {
|
||||||
|
|
|
@ -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)
|
||||||
|
}
|
|
@ -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/")
|
||||||
|
}
|
Loading…
Reference in New Issue