mirror of https://github.com/databricks/cli.git
Add `cmd-exec-id` to user agent (#1808)
## Changes This PR adds the `cmd-exec-id` field to the user agent. This allows us to correlate multiple HTTP requests made from the CLI. ### Why Not Use HTTP traceparent? We considered using the traceparent header in HTTP as an alternative, but it's not a good fit for our use case. Here's why: 1. Purpose of traceparent: It's designed to trace a single HTTP request across a distributed system as it moves through subsystems and proxies. 2. Our requirement: We need to trace multiple HTTP requests made during a single command execution in the CLI. For more details about how traceparent itself works and how it's used in the Go SDK, see https://github.com/databricks/databricks-sdk-go/pull/914. ## Tests Unit test
This commit is contained in:
parent
2bbdd04217
commit
f3bf33da27
|
@ -75,6 +75,7 @@ func New(ctx context.Context) *cobra.Command {
|
||||||
|
|
||||||
// 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(ctx, cmd)
|
ctx = withCommandInUserAgent(ctx, cmd)
|
||||||
|
ctx = withCommandExecIdInUserAgent(ctx)
|
||||||
ctx = withUpstreamInUserAgent(ctx)
|
ctx = withUpstreamInUserAgent(ctx)
|
||||||
cmd.SetContext(ctx)
|
cmd.SetContext(ctx)
|
||||||
return nil
|
return nil
|
||||||
|
|
|
@ -0,0 +1,14 @@
|
||||||
|
package root
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
"github.com/databricks/databricks-sdk-go/useragent"
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
func withCommandExecIdInUserAgent(ctx context.Context) context.Context {
|
||||||
|
// A UUID that will allow us to correlate multiple API requests made by
|
||||||
|
// the same CLI invocation.
|
||||||
|
return useragent.InContext(ctx, "cmd-exec-id", uuid.New().String())
|
||||||
|
}
|
|
@ -0,0 +1,26 @@
|
||||||
|
package root
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"regexp"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/databricks/databricks-sdk-go/useragent"
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestWithCommandExecIdInUserAgent(t *testing.T) {
|
||||||
|
ctx := withCommandExecIdInUserAgent(context.Background())
|
||||||
|
|
||||||
|
// Check that the command exec ID is in the user agent string.
|
||||||
|
ua := useragent.FromContext(ctx)
|
||||||
|
re := regexp.MustCompile(`cmd-exec-id/([a-f0-9-]+)`)
|
||||||
|
matches := re.FindAllStringSubmatch(ua, -1)
|
||||||
|
|
||||||
|
// Assert that we have exactly one match and that it's a valid UUID.
|
||||||
|
require.Len(t, matches, 1)
|
||||||
|
_, err := uuid.Parse(matches[0][1])
|
||||||
|
assert.NoError(t, err)
|
||||||
|
}
|
|
@ -1,13 +1,15 @@
|
||||||
package root
|
package root
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/databricks/databricks-sdk-go/useragent"
|
||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestCommandString(t *testing.T) {
|
func TestWithCommandInUserAgent(t *testing.T) {
|
||||||
root := &cobra.Command{
|
root := &cobra.Command{
|
||||||
Use: "root",
|
Use: "root",
|
||||||
}
|
}
|
||||||
|
@ -26,4 +28,9 @@ func TestCommandString(t *testing.T) {
|
||||||
assert.Equal(t, "root", commandString(root))
|
assert.Equal(t, "root", commandString(root))
|
||||||
assert.Equal(t, "hello", commandString(hello))
|
assert.Equal(t, "hello", commandString(hello))
|
||||||
assert.Equal(t, "hello_world", commandString(world))
|
assert.Equal(t, "hello_world", commandString(world))
|
||||||
|
|
||||||
|
ctx := withCommandInUserAgent(context.Background(), world)
|
||||||
|
|
||||||
|
ua := useragent.FromContext(ctx)
|
||||||
|
assert.Contains(t, ua, "cmd/hello_world")
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue