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.
|
||||
ctx = withCommandInUserAgent(ctx, cmd)
|
||||
ctx = withCommandExecIdInUserAgent(ctx)
|
||||
ctx = withUpstreamInUserAgent(ctx)
|
||||
cmd.SetContext(ctx)
|
||||
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
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/databricks/databricks-sdk-go/useragent"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestCommandString(t *testing.T) {
|
||||
func TestWithCommandInUserAgent(t *testing.T) {
|
||||
root := &cobra.Command{
|
||||
Use: "root",
|
||||
}
|
||||
|
@ -26,4 +28,9 @@ func TestCommandString(t *testing.T) {
|
|||
assert.Equal(t, "root", commandString(root))
|
||||
assert.Equal(t, "hello", commandString(hello))
|
||||
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