From c92f7bdf39daa387d4769440a069abf47074047c Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 4 Oct 2024 14:03:32 +0200 Subject: [PATCH 1/7] Add trace id to user agent to correlate API calls for the same command execution --- cmd/root/user_agent_command.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cmd/root/user_agent_command.go b/cmd/root/user_agent_command.go index 306f2d7bfa..baed6ab900 100644 --- a/cmd/root/user_agent_command.go +++ b/cmd/root/user_agent_command.go @@ -5,6 +5,7 @@ import ( "strings" "github.com/databricks/databricks-sdk-go/useragent" + "github.com/google/uuid" "github.com/spf13/cobra" ) @@ -33,5 +34,9 @@ func commandString(cmd *cobra.Command) string { } func withCommandInUserAgent(ctx context.Context, cmd *cobra.Command) context.Context { - return useragent.InContext(ctx, "cmd", commandString(cmd)) + // We add a uuid to the command context to trace multiple API requests made + // by the same command invocation. + newCtx := useragent.InContext(ctx, "cmd-trace-id", uuid.New().String()) + + return useragent.InContext(newCtx, "cmd", commandString(cmd)) } From 821627b08bd9704f95ebfcbf94e67dbc4a8d8ed4 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 30 Oct 2024 15:31:13 +0100 Subject: [PATCH 2/7] - --- cmd/root/user_agent_command.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cmd/root/user_agent_command.go b/cmd/root/user_agent_command.go index baed6ab900..574f6bb348 100644 --- a/cmd/root/user_agent_command.go +++ b/cmd/root/user_agent_command.go @@ -34,9 +34,11 @@ func commandString(cmd *cobra.Command) string { } func withCommandInUserAgent(ctx context.Context, cmd *cobra.Command) context.Context { - // We add a uuid to the command context to trace multiple API requests made - // by the same command invocation. - newCtx := useragent.InContext(ctx, "cmd-trace-id", uuid.New().String()) + // A UUID that'll will allow use to correlate multiple API requests made by + // the same command invocation. + // When we add telemetry to the CLI, this trace ID will allow allow us to + // correlate logs in HTTP access logs with logs in Frontend logs. + newCtx := useragent.InContext(ctx, "command-trace-id", uuid.New().String()) return useragent.InContext(newCtx, "cmd", commandString(cmd)) } From ce17497f3c9f94a7cd6cccefbc570e2ed1870720 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 30 Oct 2024 15:52:47 +0100 Subject: [PATCH 3/7] add unit test --- cmd/root/user_agent_upstream_test.go | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/cmd/root/user_agent_upstream_test.go b/cmd/root/user_agent_upstream_test.go index fc6ea0c75d..acc2260382 100644 --- a/cmd/root/user_agent_upstream_test.go +++ b/cmd/root/user_agent_upstream_test.go @@ -2,10 +2,14 @@ package root import ( "context" + "regexp" "testing" "github.com/databricks/databricks-sdk-go/useragent" + "github.com/google/uuid" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestUpstreamSet(t *testing.T) { @@ -43,3 +47,20 @@ func TestUpstreamVersionSetUpstreamNotSet(t *testing.T) { assert.NotContains(t, useragent.FromContext(ctx), "upstream/") assert.NotContains(t, useragent.FromContext(ctx), "upstream-version/") } + +func TestWithCommandInUserAgent(t *testing.T) { + ctx := withCommandInUserAgent(context.Background(), &cobra.Command{Use: "foo"}) + + // Check that the command name is in the user agent string. + ua := useragent.FromContext(ctx) + assert.Contains(t, ua, "cmd/foo") + + // Check that the command trace ID is in the user agent string. + re := regexp.MustCompile(`command-trace-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) +} From 941958cad9ea268d36c7df3f033f7ac2ead1f8ff Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 31 Oct 2024 11:16:14 +0100 Subject: [PATCH 4/7] address comments --- cmd/root/root.go | 1 + cmd/root/user_agent_command_exec_id.go | 16 +++++++++++++ cmd/root/user_agent_command_exec_id_test.go | 26 +++++++++++++++++++++ cmd/root/user_agent_command_test.go | 9 ++++++- cmd/root/user_agent_upstream_test.go | 21 ----------------- 5 files changed, 51 insertions(+), 22 deletions(-) create mode 100644 cmd/root/user_agent_command_exec_id.go create mode 100644 cmd/root/user_agent_command_exec_id_test.go diff --git a/cmd/root/root.go b/cmd/root/root.go index eda873d122..703bd80e7a 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -76,6 +76,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 = withUpstreamInUserAgent(ctx) + ctx = withCommandExecIdInUserAgent(ctx) cmd.SetContext(ctx) return nil } diff --git a/cmd/root/user_agent_command_exec_id.go b/cmd/root/user_agent_command_exec_id.go new file mode 100644 index 0000000000..0a2c7cfd03 --- /dev/null +++ b/cmd/root/user_agent_command_exec_id.go @@ -0,0 +1,16 @@ +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'll will allow use to correlate multiple API requests made by + // the same command invocation. + // When we add telemetry to the CLI, this exec ID will allow allow us to + // correlate logs in HTTP access logs with logs in Frontend logs. + return useragent.InContext(ctx, "cmd-exec-id", uuid.New().String()) +} diff --git a/cmd/root/user_agent_command_exec_id_test.go b/cmd/root/user_agent_command_exec_id_test.go new file mode 100644 index 0000000000..1dc993e173 --- /dev/null +++ b/cmd/root/user_agent_command_exec_id_test.go @@ -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 trace 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) +} diff --git a/cmd/root/user_agent_command_test.go b/cmd/root/user_agent_command_test.go index 9620bb5b8d..a3f5bbcb1c 100644 --- a/cmd/root/user_agent_command_test.go +++ b/cmd/root/user_agent_command_test.go @@ -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") } diff --git a/cmd/root/user_agent_upstream_test.go b/cmd/root/user_agent_upstream_test.go index acc2260382..fc6ea0c75d 100644 --- a/cmd/root/user_agent_upstream_test.go +++ b/cmd/root/user_agent_upstream_test.go @@ -2,14 +2,10 @@ package root import ( "context" - "regexp" "testing" "github.com/databricks/databricks-sdk-go/useragent" - "github.com/google/uuid" - "github.com/spf13/cobra" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" ) func TestUpstreamSet(t *testing.T) { @@ -47,20 +43,3 @@ func TestUpstreamVersionSetUpstreamNotSet(t *testing.T) { assert.NotContains(t, useragent.FromContext(ctx), "upstream/") assert.NotContains(t, useragent.FromContext(ctx), "upstream-version/") } - -func TestWithCommandInUserAgent(t *testing.T) { - ctx := withCommandInUserAgent(context.Background(), &cobra.Command{Use: "foo"}) - - // Check that the command name is in the user agent string. - ua := useragent.FromContext(ctx) - assert.Contains(t, ua, "cmd/foo") - - // Check that the command trace ID is in the user agent string. - re := regexp.MustCompile(`command-trace-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) -} From 9302520ce5f7806bc5e15485951b5445c3f85e06 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 31 Oct 2024 11:17:20 +0100 Subject: [PATCH 5/7] - --- cmd/root/user_agent_command.go | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/cmd/root/user_agent_command.go b/cmd/root/user_agent_command.go index 574f6bb348..306f2d7bfa 100644 --- a/cmd/root/user_agent_command.go +++ b/cmd/root/user_agent_command.go @@ -5,7 +5,6 @@ import ( "strings" "github.com/databricks/databricks-sdk-go/useragent" - "github.com/google/uuid" "github.com/spf13/cobra" ) @@ -34,11 +33,5 @@ func commandString(cmd *cobra.Command) string { } func withCommandInUserAgent(ctx context.Context, cmd *cobra.Command) context.Context { - // A UUID that'll will allow use to correlate multiple API requests made by - // the same command invocation. - // When we add telemetry to the CLI, this trace ID will allow allow us to - // correlate logs in HTTP access logs with logs in Frontend logs. - newCtx := useragent.InContext(ctx, "command-trace-id", uuid.New().String()) - - return useragent.InContext(newCtx, "cmd", commandString(cmd)) + return useragent.InContext(ctx, "cmd", commandString(cmd)) } From bbc26324745953d9eda5acb8e2b596e95232cc0e Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 31 Oct 2024 11:18:39 +0100 Subject: [PATCH 6/7] - --- cmd/root/user_agent_command_exec_id_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/root/user_agent_command_exec_id_test.go b/cmd/root/user_agent_command_exec_id_test.go index 1dc993e173..5c4365107f 100644 --- a/cmd/root/user_agent_command_exec_id_test.go +++ b/cmd/root/user_agent_command_exec_id_test.go @@ -14,7 +14,7 @@ import ( func TestWithCommandExecIdInUserAgent(t *testing.T) { ctx := withCommandExecIdInUserAgent(context.Background()) - // Check that the command trace ID is in the user agent string. + // 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) From a7e54a22d57ff09178a736cec2ce511888e41f57 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 31 Oct 2024 14:48:30 +0100 Subject: [PATCH 7/7] address comments --- cmd/root/root.go | 2 +- cmd/root/user_agent_command_exec_id.go | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/cmd/root/root.go b/cmd/root/root.go index 703bd80e7a..7059586f30 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -75,8 +75,8 @@ 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 = withUpstreamInUserAgent(ctx) ctx = withCommandExecIdInUserAgent(ctx) + ctx = withUpstreamInUserAgent(ctx) cmd.SetContext(ctx) return nil } diff --git a/cmd/root/user_agent_command_exec_id.go b/cmd/root/user_agent_command_exec_id.go index 0a2c7cfd03..3bf32b703f 100644 --- a/cmd/root/user_agent_command_exec_id.go +++ b/cmd/root/user_agent_command_exec_id.go @@ -8,9 +8,7 @@ import ( ) func withCommandExecIdInUserAgent(ctx context.Context) context.Context { - // A UUID that'll will allow use to correlate multiple API requests made by - // the same command invocation. - // When we add telemetry to the CLI, this exec ID will allow allow us to - // correlate logs in HTTP access logs with logs in Frontend logs. + // 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()) }