add event proper

cmd/bundle/init.go

@@ -101,16 +101,16 @@ func getNativeTemplateByDescription(description string) string {
 	return ""
 }
 
-func getUrlForNativeTemplate(name string) string {
+func getNativeTemplateByName(name string) *nativeTemplate {
 	for _, template := range nativeTemplates {
 		if template.name == name {
-			return template.gitUrl
+			return &template
 		}
 		if slices.Contains(template.aliases, name) {
-			return template.gitUrl
+			return &template
 		}
 	}
-	return ""
+	return nil
 }
 
 func getFsForNativeTemplate(name string) (fs.FS, error) {
@@ -235,10 +235,21 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			return nil
 		}
 
-		// Expand templatePath to a git URL if it's an alias for a known native template
+		// If templatePath refers to a native template, store it's name in a new
-		// and we know it's git URL.
+		// variable.
-		if gitUrl := getUrlForNativeTemplate(templatePath); gitUrl != "" {
+		nt := getNativeTemplateByName(templatePath)
-			templatePath = gitUrl
+		templateName := "custom"
+		isTemplateDatabricksOwned := false
+		if nt != nil {
+			templateName = templatePath
+
+			// if we have a Git URL for the native template, expand templatePath
+			// to the full URL.
+			if nt.gitUrl != "" {
+				templatePath = nt.gitUrl
+			}
+
+			isTemplateDatabricksOwned = true
 		}
 
 		if !isRepoUrl(templatePath) {
@@ -256,9 +267,19 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 				templateFS = os.DirFS(templatePath)
 			}
 
+			t := template.Template{
+				TemplateOpts: template.TemplateOpts{
+					ConfigFilePath:    configFile,
+					TemplateFS:        templateFS,
+					OutputFiler:       outputFiler,
+					IsDatabricksOwned: isTemplateDatabricksOwned,
+					Name:              templateName,
+				},
+			}
+
 			// skip downloading the repo because input arg is not a URL. We assume
 			// it's a path on the local file system in that case
-			return template.Materialize(ctx, configFile, templateFS, outputFiler)
+			return t.Materialize(ctx)
 		}
 
 		// Create a temporary directory with the name of the repository.  The '*'
@@ -283,7 +304,16 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 		// Clean up downloaded repository once the template is materialized.
 		defer os.RemoveAll(repoDir)
 		templateFS := os.DirFS(filepath.Join(repoDir, templateDir))
-		return template.Materialize(ctx, configFile, templateFS, outputFiler)
+		t := template.Template{
+			TemplateOpts: template.TemplateOpts{
+				ConfigFilePath:    configFile,
+				TemplateFS:        templateFS,
+				OutputFiler:       outputFiler,
+				IsDatabricksOwned: isTemplateDatabricksOwned,
+				Name:              templateName,
+			},
+		}
+		return t.Materialize(ctx)
 	}
 	return cmd
 }

cmd/bundle/init_test.go

@@ -46,9 +46,9 @@ func TestNativeTemplateHelpDescriptions(t *testing.T) {
 	assert.Equal(t, expected, nativeTemplateHelpDescriptions())
 }
 
-func TestGetUrlForNativeTemplate(t *testing.T) {
+func TestGetNativeTemplateByName(t *testing.T) {
-	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getUrlForNativeTemplate("mlops-stacks"))
+	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getNativeTemplateByName("mlops-stacks").gitUrl)
-	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getUrlForNativeTemplate("mlops-stack"))
+	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getNativeTemplateByName("mlops-stack").gitUrl)
-	assert.Equal(t, "", getUrlForNativeTemplate("default-python"))
+	assert.Equal(t, "", getNativeTemplateByName("default-python").gitUrl)
-	assert.Equal(t, "", getUrlForNativeTemplate("invalid"))
+	assert.Nil(t, nil, getNativeTemplateByName("invalid"))
 }

cmd/root/root.go

@@ -12,6 +12,8 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/spf13/cobra"
 )
 
@@ -52,6 +54,9 @@ func New(ctx context.Context) *cobra.Command {
 			return err
 		}
 
+		// Configure the logger to send telemetry to Databricks.
+		ctx = telemetry.NewContext(ctx)
+
 		logger := log.GetLogger(ctx)
 		logger.Info("start",
 			slog.String("version", build.GetInfo().Version),
@@ -84,6 +89,18 @@ func New(ctx context.Context) *cobra.Command {
 		return nil
 	}
 
+	cmd.PersistentPostRun = func(cmd *cobra.Command, args []string) {
+		ctx := cmd.Context()
+		w := WorkspaceClient(ctx)
+		apiClient, err := client.New(w.Config)
+		if err != nil {
+			// Uploading telemetry is best effort. Do not error.
+			return
+		}
+
+		telemetry.Flush(cmd.Context(), apiClient)
+	}
+
 	cmd.SetFlagErrorFunc(flagErrorFunc)
 	cmd.SetVersionTemplate("Databricks CLI v{{.Version}}\n")
 	return cmd

integration/bundle/helpers_test.go

@@ -42,7 +42,15 @@ func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, te
 
 	out, err := filer.NewLocalClient(bundleRoot)
 	require.NoError(t, err)
-	err = template.Materialize(ctx, configFilePath, os.DirFS(templateRoot), out)
+	tmpl := template.Template{
+		TemplateOpts: template.TemplateOpts{
+			ConfigFilePath: configFilePath,
+			TemplateFS:     os.DirFS(templateRoot),
+			OutputFiler:    out,
+		},
+	}
+
+	err = tmpl.Materialize(ctx)
 	require.NoError(t, err)
 	return bundleRoot
 }

libs/telemetry/events/bundle_init.go

@@ -11,7 +11,7 @@ type BundleInitEvent struct {
 	// Name of the template initialized when the user ran `databricks bundle init`
 	// This is only populated when the template is a first party template like
 	// mlops-stacks or default-python.
-	TemplateName BundleTemplate `json:"template_name,omitempty"`
+	TemplateName string `json:"template_name,omitempty"`
 
 	// Arguments used by the user to initialize the template. Only enum
 	// values will be set here by the Databricks CLI.
@@ -21,13 +21,3 @@ type BundleInitEvent struct {
 	// will be untenable in the long term.
 	TemplateEnumArgs map[string]string `json:"template_enum_args,omitempty"`
 }
-
-type BundleTemplate string
-
-const (
-	BundleTemplateMlopsStacks   BundleTemplate = "mlops-stacks"
-	BundleTemplateDefaultPython BundleTemplate = "default-python"
-	BundleTemplateDefaultSql    BundleTemplate = "default-sql"
-	BundleTemplateDbtSql        BundleTemplate = "dbt-sql"
-	BundleTemplateCustom        BundleTemplate = "custom"
-)

libs/telemetry/logger.go

@@ -7,9 +7,7 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/log"
-	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/google/uuid"
 )
 
@@ -58,17 +56,23 @@ func Flush(ctx context.Context, apiClient DatabricksApiClient) {
 	defer cancel()
 	l := fromContext(ctx)
 
-	// We pass the API client as an arg to mock it in unit tests.
+	if len(l.protoLogs) == 0 {
-	if apiClient == nil {
+		log.Debugf(ctx, "No telemetry events to flush")
-		var err error
-
-		// Create API client to make the the telemetry API call.
-		apiClient, err = client.New(root.WorkspaceClient(ctx).Config)
-		if err != nil {
-			log.Debugf(ctx, "error creating API client for telemetry: %v", err)
 		return
 	}
-	}
+
+	// We pass the API client as an arg to mock it in unit tests.
+	// TODO: Cleanup and remove this section.
+	// if apiClient == nil {
+	// 	var err error
+
+	// 	// Create API client to make the the telemetry API call.
+	// 	apiClient, err = client.New(root.WorkspaceClient(ctx).Config)
+	// 	if err != nil {
+	// 		log.Debugf(ctx, "error creating API client for telemetry: %v", err)
+	// 		return
+	// 	}
+	// }
 
 	resp := &ResponseBody{}
 	for {

libs/template/config.go

@@ -273,3 +273,16 @@ func (c *config) validate() error {
 	}
 	return nil
 }
+
+// Return enum values selected by the user during template initialization. These
+// values are safe to send over in telemetry events due to their limited cardinality.
+func (c *config) enumValues() map[string]string {
+	res := map[string]string{}
+	for k, p := range c.schema.Properties {
+		if p.Enum == nil {
+			continue
+		}
+		res[k] = c.values[k].(string)
+	}
+	return res
+}

libs/template/materialize.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/telemetry/events"
 )
 
 const (
@@ -16,42 +18,65 @@ const (
 	schemaFileName  = "databricks_template_schema.json"
 )
 
-// This function materializes the input templates as a project, using user defined
+type TemplateOpts struct {
-// configurations.
+	// file path containing user defined config values
-// Parameters:
+	ConfigFilePath string
-//
+	// root of the template definition
-//	ctx: 			context containing a cmdio object. This is used to prompt the user
+	TemplateFS fs.FS
-//	configFilePath: file path containing user defined config values
+	// filer to use for writing the initialized template
-//	templateFS: 	root of the template definition
+	OutputFiler filer.Filer
-//	outputFiler: 	filer to use for writing the initialized template
+	// If true, we'll include the enum template args in the telemetry payload..
-func Materialize(ctx context.Context, configFilePath string, templateFS fs.FS, outputFiler filer.Filer) error {
+	IsDatabricksOwned bool
-	if _, err := fs.Stat(templateFS, schemaFileName); errors.Is(err, fs.ErrNotExist) {
+	// Name of the template. For non-Databricks owned templates, this is set to
+	// custom.
+	// TODO: move this enum to the template package.
+	Name string
+}
+
+type Template struct {
+	TemplateOpts
+
+	// internal object used to prompt user for config values and store them.
+	config *config
+
+	// internal object user to render the template.
+	renderer *renderer
+}
+
+// This function resolves input to use to materialize the template in two steps.
+//  1. First, this function loads any user specified input configuration if the user
+//     has provided a config file path.
+//  2. For any values that are required by the template but not provided in the config
+//     file, this function prompts the user for them.
+func (t *Template) resolveTemplateInput(ctx context.Context) error {
+	if _, err := fs.Stat(t.TemplateFS, schemaFileName); errors.Is(err, fs.ErrNotExist) {
 		return fmt.Errorf("not a bundle template: expected to find a template schema file at %s", schemaFileName)
 	}
 
-	config, err := newConfig(ctx, templateFS, schemaFileName)
+	var err error
+	t.config, err = newConfig(ctx, t.TemplateFS, schemaFileName)
 	if err != nil {
 		return err
 	}
 
 	// Read and assign config values from file
-	if configFilePath != "" {
+	if t.ConfigFilePath != "" {
-		err = config.assignValuesFromFile(configFilePath)
+		err = t.config.assignValuesFromFile(t.ConfigFilePath)
 		if err != nil {
 			return err
 		}
 	}
 
 	helpers := loadHelpers(ctx)
-	r, err := newRenderer(ctx, config.values, helpers, templateFS, templateDirName, libraryDirName)
+	t.renderer, err = newRenderer(ctx, t.config.values, helpers, t.TemplateFS, templateDirName, libraryDirName)
 	if err != nil {
 		return err
 	}
 
 	// Print welcome message
-	welcome := config.schema.WelcomeMessage
+	welcome := t.config.schema.WelcomeMessage
 	if welcome != "" {
-		welcome, err = r.executeTemplate(welcome)
+		welcome, err = t.renderer.executeTemplate(welcome)
 		if err != nil {
 			return err
 		}
@@ -60,35 +85,74 @@ func Materialize(ctx context.Context, configFilePath string, templateFS fs.FS, o
 
 	// Prompt user for any missing config values. Assign default values if
 	// terminal is not TTY
-	err = config.promptOrAssignDefaultValues(r)
+	err = t.config.promptOrAssignDefaultValues(t.renderer)
 	if err != nil {
 		return err
 	}
-	err = config.validate()
+	return t.config.validate()
-	if err != nil {
-		return err
 }
 
-	// Walk and render the template, since input configuration is complete
+func (t *Template) printSuccessMessage(ctx context.Context) error {
-	err = r.walk()
+	success := t.config.schema.SuccessMessage
-	if err != nil {
-		return err
-	}
-
-	err = r.persistToDisk(ctx, outputFiler)
-	if err != nil {
-		return err
-	}
-
-	success := config.schema.SuccessMessage
 	if success == "" {
 		cmdio.LogString(ctx, "✨ Successfully initialized template")
-	} else {
+		return nil
-		success, err = r.executeTemplate(success)
+	}
+
+	success, err := t.renderer.executeTemplate(success)
 	if err != nil {
 		return err
 	}
 	cmdio.LogString(ctx, success)
-	}
 	return nil
 }
+
+func (t *Template) logTelemetry(ctx context.Context) error {
+	// Only log telemetry input for Databricks owned templates. This is to prevent
+	// accidentally collecting PUII from custom user templates.
+	templateEnumArgs := map[string]string{}
+	if !t.IsDatabricksOwned {
+		templateEnumArgs = t.config.enumValues()
+	}
+
+	event := telemetry.DatabricksCliLog{
+		BundleInitEvent: &events.BundleInitEvent{
+			Uuid:             bundleUuid,
+			TemplateName:     t.Name,
+			TemplateEnumArgs: templateEnumArgs,
+		},
+	}
+
+	return telemetry.Log(ctx, telemetry.FrontendLogEntry{
+		DatabricksCliLog: event,
+	})
+}
+
+// This function materializes the input templates as a project, using user defined
+// configurations.
+func (t *Template) Materialize(ctx context.Context) error {
+	err := t.resolveTemplateInput(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Walk the template file tree and compute in-memory representations of the
+	// output files.
+	err = t.renderer.walk()
+	if err != nil {
+		return err
+	}
+
+	// Flush the output files to disk.
+	err = t.renderer.persistToDisk(ctx, t.OutputFiler)
+	if err != nil {
+		return err
+	}
+
+	err = t.printSuccessMessage(ctx)
+	if err != nil {
+		return err
+	}
+
+	return t.logTelemetry(ctx)
+}

libs/template/materialize_test.go

@@ -18,7 +18,15 @@ func TestMaterializeForNonTemplateDirectory(t *testing.T) {
 	require.NoError(t, err)
 	ctx := root.SetWorkspaceClient(context.Background(), w)
 
+	tmpl := Template{
+		TemplateOpts: TemplateOpts{
+			ConfigFilePath: "",
+			TemplateFS:     os.DirFS(tmpDir),
+			OutputFiler:    nil,
+		},
+	}
+
 	// Try to materialize a non-template directory.
-	err = Materialize(ctx, "", os.DirFS(tmpDir), nil)
+	err = tmpl.Materialize(ctx)
 	assert.EqualError(t, err, fmt.Sprintf("not a bundle template: expected to find a template schema file at %s", schemaFileName))
 }