From 9ca7f8a888cfcd9a89ba9cd19e5a13d80d0a5cc1 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Fri, 3 Feb 2023 16:47:33 +0100
Subject: [PATCH] Configure user agent in root command (#195)

This configures the user agent with the bricks version and the name of
the command being executed.

Example user agent value:
```
> * User-Agent: bricks/0.0.21-devel databricks-sdk-go/0.2.0 go/1.19.4 os/darwin cmd/sync auth/pat
```

This is a follow up for #194.
---
 cmd/root/root.go            |  4 ++++
 cmd/root/user_agent.go      | 42 +++++++++++++++++++++++++++++++++++++
 cmd/root/user_agent_test.go | 29 +++++++++++++++++++++++++
 main_test.go                | 24 +++++++++++++++++++++
 4 files changed, 99 insertions(+)
 create mode 100644 cmd/root/user_agent.go
 create mode 100644 cmd/root/user_agent_test.go
 create mode 100644 main_test.go

diff --git a/cmd/root/root.go b/cmd/root/root.go
index 2722461b..6bea7806 100644
--- a/cmd/root/root.go
+++ b/cmd/root/root.go
@@ -23,6 +23,10 @@ var RootCmd = &cobra.Command{
 	SilenceUsage: true,
 
 	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		// Configure our user agent with the command that's about to be executed.
+		ctx := withCommandInUserAgent(cmd.Context(), cmd)
+		cmd.SetContext(ctx)
+
 		if Verbose {
 			logLevel = append(logLevel, "[DEBUG]")
 		}
diff --git a/cmd/root/user_agent.go b/cmd/root/user_agent.go
new file mode 100644
index 00000000..80d681aa
--- /dev/null
+++ b/cmd/root/user_agent.go
@@ -0,0 +1,42 @@
+package root
+
+import (
+	"context"
+	"strings"
+
+	"github.com/databricks/bricks/internal/build"
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/spf13/cobra"
+)
+
+// commandSeparator joins command names in a command hierachy.
+// We enforce no command name contains this character.
+// See unit test [main.TestCommandsDontUseUnderscoreInName].
+const commandSeparator = "_"
+
+// commandString walks up the command hierarchy of the specified
+// command to build a string representing this hierarchy.
+func commandString(cmd *cobra.Command) string {
+	reversed := []string{cmd.Name()}
+	cmd.VisitParents(func(p *cobra.Command) {
+		if !p.HasParent() {
+			return
+		}
+		reversed = append(reversed, p.Name())
+	})
+
+	ordered := make([]string, 0, len(reversed))
+	for i := len(reversed) - 1; i >= 0; i-- {
+		ordered = append(ordered, reversed[i])
+	}
+
+	return strings.Join(ordered, commandSeparator)
+}
+
+func withCommandInUserAgent(ctx context.Context, cmd *cobra.Command) context.Context {
+	return useragent.InContext(ctx, "cmd", commandString(cmd))
+}
+
+func init() {
+	useragent.WithProduct("bricks", build.GetInfo().Version)
+}
diff --git a/cmd/root/user_agent_test.go b/cmd/root/user_agent_test.go
new file mode 100644
index 00000000..9620bb5b
--- /dev/null
+++ b/cmd/root/user_agent_test.go
@@ -0,0 +1,29 @@
+package root
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCommandString(t *testing.T) {
+	root := &cobra.Command{
+		Use: "root",
+	}
+
+	hello := &cobra.Command{
+		Use: "hello",
+	}
+
+	world := &cobra.Command{
+		Use: "world",
+	}
+
+	root.AddCommand(hello)
+	hello.AddCommand(world)
+
+	assert.Equal(t, "root", commandString(root))
+	assert.Equal(t, "hello", commandString(hello))
+	assert.Equal(t, "hello_world", commandString(world))
+}
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 00000000..b75eb1c9
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/databricks/bricks/cmd/root"
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCommandsDontUseUnderscoreInName(t *testing.T) {
+	// We use underscore as separator between commands in logs
+	// so need to enforce that no command uses it in its name.
+	//
+	// This test lives in the main package because this is where
+	// all commands are imported.
+	//
+	queue := []*cobra.Command{root.RootCmd}
+	for len(queue) > 0 {
+		cmd := queue[0]
+		assert.NotContains(t, cmd.Name(), "_")
+		queue = append(queue[1:], cmd.Commands()...)
+	}
+}