patch template args logging

integration/bundle/init_test.go

@@ -74,12 +74,22 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 	require.Equal(t, 1, len(logs))
 	event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
 	assert.Equal(t, event.TemplateName, "mlops-stacks")
+
+	get := func(key string) string {
+		for _, v := range event.TemplateEnumArgs {
+			if v.Key == key {
+				return v.Value
+			}
+		}
+		return ""
+	}
+
 	// Enum values should be present in the telemetry payload.
-	assert.Equal(t, event.TemplateEnumArgs["input_include_models_in_unity_catalog"], "no")
-	assert.Equal(t, event.TemplateEnumArgs["input_cloud"], strings.ToLower(env))
+	assert.Equal(t, get("input_include_models_in_unity_catalog"), "no")
+	assert.Equal(t, get("input_cloud"), strings.ToLower(env))
 	// Freeform strings should not be present in the telemetry payload.
-	assert.NotContains(t, event.TemplateEnumArgs, "input_project_name")
-	assert.NotContains(t, event.TemplateEnumArgs, "input_root_dir")
+	assert.Equal(t, get("input_project_name"), "")
+	assert.Equal(t, get("input_root_dir"), "")
 
 	// Assert that the README.md file was created
 	contents := testutil.ReadFile(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
@@ -191,7 +201,21 @@ func TestBundleInitTelemetryForDefaultTemplates(t *testing.T) {
 		require.Equal(t, 1, len(logs))
 		event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
 		assert.Equal(t, event.TemplateName, tc.name)
-		assert.Equal(t, event.TemplateEnumArgs, tc.expectedArgs)
+
+		get := func(key string) string {
+			for _, v := range event.TemplateEnumArgs {
+				if v.Key == key {
+					return v.Value
+				}
+			}
+			return ""
+		}
+
+		// Assert the template enum args are correctly logged.
+		assert.Len(t, event.TemplateEnumArgs, len(tc.expectedArgs))
+		for k, v := range tc.expectedArgs {
+			assert.Equal(t, get(k), v)
+		}
 	}
 }
 

integration/internal/main.go

@@ -1,8 +1,6 @@
 package internal
 
 import (
-	"fmt"
-	"os"
 	"testing"
 )
 
@@ -10,11 +8,11 @@ import (
 // We use this for all integration tests defined in this subtree to ensure
 // they are not inadvertently executed when calling `go test ./...`.
 func Main(m *testing.M) {
-	value := os.Getenv("CLOUD_ENV")
-	if value == "" {
-		fmt.Println("CLOUD_ENV is not set, skipping integration tests")
-		return
-	}
+	// value := os.Getenv("CLOUD_ENV")
+	// if value == "" {
+	// 	fmt.Println("CLOUD_ENV is not set, skipping integration tests")
+	// 	return
+	// }
 
 	m.Run()
 }

integration/libs/telemetry/telemetry_test.go

@@ -43,9 +43,15 @@ func TestTelemetryLogger(t *testing.T) {
 			BundleInitEvent: &events.BundleInitEvent{
 				Uuid:         uuid.New().String(),
 				TemplateName: "abc",
-				TemplateEnumArgs: map[string]string{
-					"a": "b",
-					"c": "d",
+				TemplateEnumArgs: []events.BundleInitTemplateEnumArg{
+					{
+						Key:   "a",
+						Value: "b",
+					},
+					{
+						Key:   "c",
+						Value: "d",
+					},
 				},
 			},
 		},

libs/telemetry/events/bundle_init.go

@@ -19,5 +19,21 @@ type BundleInitEvent struct {
 	// We use a generic map representation here because a bundle template's args are
 	// managed in the template itself and maintaining a copy typed schema for it here
 	// will be untenable in the long term.
-	TemplateEnumArgs map[string]string `json:"template_enum_args,omitempty"`
+	TemplateEnumArgs []BundleInitTemplateEnumArg `json:"template_enum_args,omitempty"`
+}
+
+type BundleInitTemplateEnumArg struct {
+	// Valid key values for the template. These correspond to the keys specified in
+	// the "properties" section of the `databricks_template_schema.json` file.
+	//
+	// Note: `databricks_template_schema.json` contains a JSON schema type specification
+	// for the arguments that the template accepts.
+	Key string `json:"key"`
+
+	// Value that the user set for the field. This is only populated for properties
+	// that have the "enum" field specified in the JSON schema type specification.
+	//
+	// The Databricks CLI ensures that the value here is one of the "enum" values from
+	// the template specification.
+	Value string `json:"value"`
 }

libs/template/materialize.go

@@ -109,9 +109,14 @@ func (t *Template) printSuccessMessage(ctx context.Context) error {
 func (t *Template) logTelemetry(ctx context.Context) error {
 	// Only log telemetry input for Databricks owned templates. This is to prevent
 	// accidentally collecting PII from custom user templates.
-	templateEnumArgs := map[string]string{}
+	templateEnumArgs := []events.BundleInitTemplateEnumArg{}
 	if t.IsDatabricksOwned {
-		templateEnumArgs = t.config.enumValues()
+		for k, v := range t.config.enumValues() {
+			templateEnumArgs = append(templateEnumArgs, events.BundleInitTemplateEnumArg{
+				Key:   k,
+				Value: v,
+			})
+		}
 	} else {
 		t.Name = "custom"
 	}