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:
shreyas-goenka 2024-11-01 19:38:09 +05:30 committed by GitHub
parent 2bbdd04217
commit f3bf33da27
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 49 additions and 1 deletions

View File

@ -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

View File

@ -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())
}

View File

@ -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)
}

View File

@ -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")
} }