From 2e35cdf3b7e992c608e6973d1a7820e147f89d81 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Wed, 26 Jun 2024 17:10:22 +0200
Subject: [PATCH 1/2] PythonMutator: register product in user agent extra

---
 bundle/deploy/terraform/init.go      | 23 +++++++++++++++++++++++
 bundle/deploy/terraform/init_test.go | 21 +++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/bundle/deploy/terraform/init.go b/bundle/deploy/terraform/init.go
index d1847cf2443..cfd7e526ae2 100644
--- a/bundle/deploy/terraform/init.go
+++ b/bundle/deploy/terraform/init.go
@@ -218,6 +218,24 @@ func setProxyEnvVars(ctx context.Context, environ map[string]string, b *bundle.B
 	return nil
 }
 
+func setUserAgentExtraEnvVar(environ map[string]string, b *bundle.Bundle) error {
+	var products []string
+
+	if experimental := b.Config.Experimental; experimental != nil {
+		if experimental.PyDABs.Enabled {
+			products = append(products, "databricks-pydabs/0.0.0")
+		}
+	}
+
+	userAgentExtra := strings.Join(products, " ")
+
+	if userAgentExtra != "" {
+		environ["DATABRICKS_USER_AGENT_EXTRA"] = userAgentExtra
+	}
+
+	return nil
+}
+
 func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 	tfConfig := b.Config.Bundle.Terraform
 	if tfConfig == nil {
@@ -262,6 +280,11 @@ func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnosti
 		return diag.FromErr(err)
 	}
 
+	err = setUserAgentExtraEnvVar(environ, b)
+	if err != nil {
+		return diag.FromErr(err)
+	}
+
 	// Configure environment variables for auth for Terraform to use.
 	log.Debugf(ctx, "Environment variables for Terraform: %s", strings.Join(maps.Keys(environ), ", "))
 	err = tf.SetEnv(environ)
diff --git a/bundle/deploy/terraform/init_test.go b/bundle/deploy/terraform/init_test.go
index 421e9be3fbf..aa9b2f77fdd 100644
--- a/bundle/deploy/terraform/init_test.go
+++ b/bundle/deploy/terraform/init_test.go
@@ -248,6 +248,27 @@ func TestSetProxyEnvVars(t *testing.T) {
 	assert.ElementsMatch(t, []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY"}, maps.Keys(env))
 }
 
+func TestSetUserAgentExtraEnvVar(t *testing.T) {
+	b := &bundle.Bundle{
+		RootPath: t.TempDir(),
+		Config: config.Root{
+			Experimental: &config.Experimental{
+				PyDABs: config.PyDABs{
+					Enabled: true,
+				},
+			},
+		},
+	}
+
+	env := make(map[string]string, 0)
+	err := setUserAgentExtraEnvVar(env, b)
+
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"DATABRICKS_USER_AGENT_EXTRA": "databricks-pydabs/0.0.0",
+	}, env)
+}
+
 func TestInheritEnvVars(t *testing.T) {
 	env := map[string]string{}
 

From b84b66979140b82fe077df266a3aa93dab6c9358 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Thu, 27 Jun 2024 15:28:20 +0200
Subject: [PATCH 2/2] Fix new line

---
 bundle/deploy/terraform/init.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/bundle/deploy/terraform/init.go b/bundle/deploy/terraform/init.go
index cfd7e526ae2..d480242ce99 100644
--- a/bundle/deploy/terraform/init.go
+++ b/bundle/deploy/terraform/init.go
@@ -228,7 +228,6 @@ func setUserAgentExtraEnvVar(environ map[string]string, b *bundle.Bundle) error
 	}
 
 	userAgentExtra := strings.Join(products, " ")
-
 	if userAgentExtra != "" {
 		environ["DATABRICKS_USER_AGENT_EXTRA"] = userAgentExtra
 	}