From 941958cad9ea268d36c7df3f033f7ac2ead1f8ff Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 31 Oct 2024 11:16:14 +0100
Subject: [PATCH] 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 eda873d12..703bd80e7 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 000000000..0a2c7cfd0
--- /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 000000000..1dc993e17
--- /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 9620bb5b8..a3f5bbcb1 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 acc226038..fc6ea0c75 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)
-}