From f8f804fe17ea49650068c521d8ac4cd8501ef22b Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <kanterov@users.noreply.github.com>
Date: Mon, 13 Jan 2025 10:16:29 +0100
Subject: [PATCH] PythonMutator: update instrumentation (#2124)

## Changes
Update instrumentation for PythonMutator to handle `experimental/python`
config.

## Tests
Unit tests
---
 bundle/deploy/terraform/init.go      |  6 +++++-
 bundle/deploy/terraform/init_test.go | 22 +++++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/bundle/deploy/terraform/init.go b/bundle/deploy/terraform/init.go
index e69f0bf0f..d982354e1 100644
--- a/bundle/deploy/terraform/init.go
+++ b/bundle/deploy/terraform/init.go
@@ -232,7 +232,11 @@ func setUserAgentExtraEnvVar(environ map[string]string, b *bundle.Bundle) error
 	// Terraform provider to the CLI.
 	products := []string{"cli/" + build.GetInfo().Version}
 	if experimental := b.Config.Experimental; experimental != nil {
-		if experimental.PyDABs.Enabled {
+		hasPython := experimental.Python.Resources != nil || experimental.Python.Mutators != nil
+
+		if hasPython {
+			products = append(products, "databricks-pydabs/0.7.0")
+		} else if experimental.PyDABs.Enabled {
 			products = append(products, "databricks-pydabs/0.0.0")
 		}
 	}
diff --git a/bundle/deploy/terraform/init_test.go b/bundle/deploy/terraform/init_test.go
index 30ac9e301..c7a4ffe4a 100644
--- a/bundle/deploy/terraform/init_test.go
+++ b/bundle/deploy/terraform/init_test.go
@@ -248,7 +248,7 @@ func TestSetProxyEnvVars(t *testing.T) {
 	assert.ElementsMatch(t, []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY"}, maps.Keys(env))
 }
 
-func TestSetUserAgentExtraEnvVar(t *testing.T) {
+func TestSetUserAgentExtraEnvVar_PyDABs(t *testing.T) {
 	b := &bundle.Bundle{
 		BundleRootPath: t.TempDir(),
 		Config: config.Root{
@@ -268,6 +268,26 @@ func TestSetUserAgentExtraEnvVar(t *testing.T) {
 	}, env)
 }
 
+func TestSetUserAgentExtraEnvVar_Python(t *testing.T) {
+	b := &bundle.Bundle{
+		BundleRootPath: t.TempDir(),
+		Config: config.Root{
+			Experimental: &config.Experimental{
+				Python: config.Python{
+					Resources: []string{"my_project.resources:load_resources"},
+				},
+			},
+		},
+	}
+
+	env := make(map[string]string, 0)
+	err := setUserAgentExtraEnvVar(env, b)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"DATABRICKS_USER_AGENT_EXTRA": "cli/0.0.0-dev databricks-pydabs/0.7.0",
+	}, env)
+}
+
 func TestInheritEnvVars(t *testing.T) {
 	t.Setenv("HOME", "/home/testuser")
 	t.Setenv("PATH", "/foo:/bar")