diff --git a/integration/bundle/init_test.go b/integration/bundle/init_test.go index 3a83848ad..4fe61af75 100644 --- a/integration/bundle/init_test.go +++ b/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) + } } } diff --git a/integration/internal/main.go b/integration/internal/main.go index 6d69dcf70..01eba1fd8 100644 --- a/integration/internal/main.go +++ b/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() } diff --git a/integration/libs/telemetry/telemetry_test.go b/integration/libs/telemetry/telemetry_test.go index 80918ebba..5fdceb2cd 100644 --- a/integration/libs/telemetry/telemetry_test.go +++ b/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", + }, }, }, }, diff --git a/internal/main.go b/internal/main.go new file mode 100644 index 000000000..e69de29bb diff --git a/libs/telemetry/events/bundle_init.go b/libs/telemetry/events/bundle_init.go index 38be47f60..73e681509 100644 --- a/libs/telemetry/events/bundle_init.go +++ b/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"` } diff --git a/libs/template/materialize.go b/libs/template/materialize.go index 32478c859..5b79b96cb 100644 --- a/libs/template/materialize.go +++ b/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" }