squash all commits

cmd/bundle/init.go

@@ -6,7 +6,10 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/template"
+	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/spf13/cobra"
 )
 
@@ -36,12 +39,28 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 	cmd.Flags().StringVar(&branch, "tag", "", "Git tag to use for template initialization")
 	cmd.Flags().StringVar(&tag, "branch", "", "Git branch to use for template initialization")
 
-	cmd.PreRunE = root.MustWorkspaceClient
-	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		if tag != "" && branch != "" {
-			return errors.New("only one of --tag or --branch can be specified")
+	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
+		// Configure the logger to send telemetry to Databricks.
+		ctx := telemetry.WithDefaultLogger(cmd.Context())
+		cmd.SetContext(ctx)
+
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
+	cmd.PostRun = func(cmd *cobra.Command, args []string) {
+		ctx := cmd.Context()
+		w := root.WorkspaceClient(ctx)
+		apiClient, err := client.New(w.Config)
+		if err != nil {
+			// Uploading telemetry is best effort. Do not error.
+			log.Debugf(ctx, "Could not create API client to send telemetry using: %v", err)
+			return
 		}
 
+		telemetry.Flush(cmd.Context(), apiClient)
+	}
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		var templatePathOrUrl string
 		if len(args) > 0 {
 			templatePathOrUrl = args[0]
@@ -67,7 +86,14 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 		}
 		defer tmpl.Reader.Cleanup(ctx)
 
-		return tmpl.Writer.Materialize(ctx, tmpl.Reader)
+		err = tmpl.Writer.Materialize(ctx, tmpl.Reader)
+		if err != nil {
+			return err
+		}
+		defer tmpl.Reader.Cleanup(ctx)
+
+		tmpl.Writer.LogTelemetry(ctx)
+		return nil
 	}
 	return cmd
 }

integration/bundle/helpers_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/folders"
+	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/template"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/stretchr/testify/require"
@@ -38,6 +39,7 @@ func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, te
 	ctx = root.SetWorkspaceClient(ctx, nil)
 	cmd := cmdio.NewIO(ctx, flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "bundles")
 	ctx = cmdio.InContext(ctx, cmd)
+	ctx = telemetry.WithMockLogger(ctx)
 
 	r := template.Resolver{
 		TemplatePathOrUrl: templateRoot,

integration/bundle/init_test.go

@@ -15,6 +15,8 @@ import (
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/iamutil"
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/template"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -42,6 +44,9 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
+	// Use mock logger to introspect the telemetry payload.
+	ctx = telemetry.WithMockLogger(ctx)
+
 	tmpDir1 := t.TempDir()
 	tmpDir2 := t.TempDir()
 
@@ -64,6 +69,29 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 	assert.NoFileExists(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
 	testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", "mlops-stacks", "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json"))
 
+	// Assert the telemetry payload is correctly logged.
+	telemetryEvents := telemetry.Introspect(ctx)
+	require.Len(t, telemetry.Introspect(ctx), 1)
+	event := telemetryEvents[0].BundleInitEvent
+	assert.Equal(t, "mlops-stacks", event.TemplateName)
+
+	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, "no", get("input_include_models_in_unity_catalog"))
+	assert.Equal(t, strings.ToLower(env), get("input_cloud"))
+
+	// Freeform strings should not be present in the telemetry payload.
+	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"))
 	assert.Contains(t, contents, "# "+projectName)
@@ -99,6 +127,156 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 	assert.Contains(t, job.Settings.Name, fmt.Sprintf("dev-%s-batch-inference-job", projectName))
 }
 
+func TestBundleInitTelemetryForDefaultTemplates(t *testing.T) {
+	projectName := testutil.RandomName("name_")
+
+	tcases := []struct {
+		name         string
+		args         map[string]string
+		expectedArgs map[string]string
+	}{
+		{
+			name: "dbt-sql",
+			args: map[string]string{
+				"project_name":     fmt.Sprintf("dbt-sql-%s", projectName),
+				"http_path":        "/sql/1.0/warehouses/id",
+				"default_catalog":  "abcd",
+				"personal_schemas": "yes, use a schema based on the current user name during development",
+			},
+			expectedArgs: map[string]string{
+				"personal_schemas": "yes, use a schema based on the current user name during development",
+			},
+		},
+		{
+			name: "default-python",
+			args: map[string]string{
+				"project_name":     fmt.Sprintf("default_python_%s", projectName),
+				"include_notebook": "yes",
+				"include_dlt":      "yes",
+				"include_python":   "no",
+			},
+			expectedArgs: map[string]string{
+				"include_notebook": "yes",
+				"include_dlt":      "yes",
+				"include_python":   "no",
+			},
+		},
+		{
+			name: "default-sql",
+			args: map[string]string{
+				"project_name":     fmt.Sprintf("sql_project_%s", projectName),
+				"http_path":        "/sql/1.0/warehouses/id",
+				"default_catalog":  "abcd",
+				"personal_schemas": "yes, automatically use a schema based on the current user name during development",
+			},
+			expectedArgs: map[string]string{
+				"personal_schemas": "yes, automatically use a schema based on the current user name during development",
+			},
+		},
+	}
+
+	for _, tc := range tcases {
+		ctx, _ := acc.WorkspaceTest(t)
+
+		// Use mock logger to introspect the telemetry payload.
+		ctx = telemetry.WithMockLogger(ctx)
+
+		tmpDir1 := t.TempDir()
+		tmpDir2 := t.TempDir()
+
+		// Create a config file with the project name and root dir
+		initConfig := tc.args
+		b, err := json.Marshal(initConfig)
+		require.NoError(t, err)
+		err = os.WriteFile(filepath.Join(tmpDir1, "config.json"), b, 0o644)
+		require.NoError(t, err)
+
+		// Run bundle init
+		assert.NoDirExists(t, filepath.Join(tmpDir2, tc.args["project_name"]))
+		testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", tc.name, "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json"))
+		assert.DirExists(t, filepath.Join(tmpDir2, tc.args["project_name"]))
+
+		// Assert the telemetry payload is correctly logged.
+		logs := telemetry.Introspect(ctx)
+		require.Len(t, logs, 1)
+		event := logs[0].BundleInitEvent
+		assert.Equal(t, event.TemplateName, tc.name)
+
+		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)
+		}
+	}
+}
+
+func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
+	ctx, _ := acc.WorkspaceTest(t)
+
+	tmpDir1 := t.TempDir()
+	tmpDir2 := t.TempDir()
+	tmpDir3 := t.TempDir()
+
+	err := os.Mkdir(filepath.Join(tmpDir1, "template"), 0o755)
+	require.NoError(t, err)
+	err = os.WriteFile(filepath.Join(tmpDir1, "template", "foo.txt.tmpl"), []byte("{{bundle_uuid}}"), 0o644)
+	require.NoError(t, err)
+	err = os.WriteFile(filepath.Join(tmpDir1, "databricks_template_schema.json"), []byte(`
+	{
+		"properties": {
+			"a": {
+				"description": "whatever",
+				"type": "string"
+			},
+			"b": {
+				"description": "whatever",
+				"type": "string",
+				"enum": ["yes", "no"]
+			}
+		}
+	}
+	`), 0o644)
+	require.NoError(t, err)
+
+	// Create a config file with the project name and root dir
+	initConfig := map[string]string{
+		"a": "v1",
+		"b": "yes",
+	}
+	b, err := json.Marshal(initConfig)
+	require.NoError(t, err)
+	err = os.WriteFile(filepath.Join(tmpDir3, "config.json"), b, 0o644)
+	require.NoError(t, err)
+
+	// Use mock logger to introspect the telemetry payload.
+	ctx = telemetry.WithMockLogger(ctx)
+
+	// Run bundle init.
+	testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", tmpDir1, "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir3, "config.json"))
+
+	// Assert the telemetry payload is correctly logged. For custom templates we should
+	// never set template_enum_args.
+	telemetryEvents := telemetry.Introspect(ctx)
+	require.Len(t, telemetryEvents, 1)
+	event := telemetryEvents[0].BundleInitEvent
+	assert.Equal(t, string(template.Custom), event.TemplateName)
+	assert.Empty(t, event.TemplateEnumArgs)
+
+	// Ensure that the UUID returned by the `bundle_uuid` helper is the same UUID
+	// that's logged in the telemetry event.
+	fileC := testutil.ReadFile(t, filepath.Join(tmpDir2, "foo.txt"))
+	assert.Equal(t, event.Uuid, fileC)
+}
+
 func TestBundleInitHelpers(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W

libs/telemetry/protos/bundle_init.go

@@ -6,7 +6,7 @@ type BundleInitEvent struct {
 	// UUID associated with the DAB itself. This is serialized into the DAB
 	// when a user runs `databricks bundle init` and all subsequent deployments of
 	// that DAB can then be associated with this init event.
-	Uuid string `json:"bundle_uuid,omitempty"`
+	Uuid string `json:"uuid,omitempty"`
 
 	// 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

libs/template/config.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"slices"
 
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/jsonschema"
@@ -273,3 +274,25 @@ 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 {
+		// For now, we only send over string enum values since only those are
+		// being used in first party Databricks templates.
+		if p.Type != jsonschema.StringType {
+			continue
+		}
+		if p.Enum == nil {
+			continue
+		}
+		v := c.values[k]
+
+		if slices.Contains(p.Enum, v) {
+			res[k] = v.(string)
+		}
+	}
+	return res
+}

libs/template/config_test.go

@@ -564,3 +564,42 @@ func TestPromptIsSkippedAnyOf(t *testing.T) {
 	assert.True(t, skip)
 	assert.Equal(t, "hello-world", c.values["xyz"])
 }
+
+func TestConfigEnumValues(t *testing.T) {
+	c := &config{
+		schema: &jsonschema.Schema{
+			Properties: map[string]*jsonschema.Schema{
+				"a": {
+					Type: jsonschema.StringType,
+				},
+				"b": {
+					Type: jsonschema.BooleanType,
+				},
+				"c": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v1", "v2"},
+				},
+				"d": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v3", "v4"},
+				},
+				"e": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v5", "v6"},
+				},
+			},
+		},
+		values: map[string]any{
+			"a": "w1",
+			"b": false,
+			"c": "v1",
+			"d": "v3",
+			"e": "v7",
+		},
+	}
+
+	assert.Equal(t, map[string]string{
+		"c": "v1",
+		"d": "v3",
+	}, c.enumValues())
+}

libs/template/resolver.go

@@ -10,7 +10,6 @@ import (
 
 var gitUrlPrefixes = []string{
 	"https://",
-	"git@",
 }
 
 func isRepoUrl(url string) bool {
@@ -27,7 +26,6 @@ func isRepoUrl(url string) bool {
 type Resolver struct {
 	// One of the following three:
 	// 1. Path to a local template directory.
-	// 2. URL to a Git repository containing a template.
 	// 3. Name of a built-in template.
 	TemplatePathOrUrl string