cmd/bundle/init.go

@@ -15,7 +15,6 @@ import (
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/git"
-	"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"
@@ -213,7 +212,6 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 		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
 		}
 

integration/bundle/init_test.go

@@ -69,7 +69,7 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 	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.
-	logs := telemetry.GetLogs(ctx)
+	logs, err := telemetry.GetLogs(ctx)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(logs))
 	event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
@@ -186,7 +186,7 @@ func TestBundleInitTelemetryForDefaultTemplates(t *testing.T) {
 		assert.DirExists(t, filepath.Join(tmpDir2, tc.args["project_name"]))
 
 		// Assert the telemetry payload is correctly logged.
-		logs := telemetry.GetLogs(ctx)
+		logs, err := telemetry.GetLogs(ctx)
 		require.NoError(t, err)
 		require.Equal(t, 1, len(logs))
 		event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
@@ -241,7 +241,7 @@ func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
 
 	// Assert the telemetry payload is correctly logged. For custom templates we should
 	// never set template_enum_args.
-	logs := telemetry.GetLogs(ctx)
+	logs, err := telemetry.GetLogs(ctx)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(logs))
 	event := logs[0].Entry.DatabricksCliLog.BundleInitEvent

integration/libs/telemetry/telemetry_test.go

@@ -33,19 +33,23 @@ func (wrapper *apiClientWrapper) Do(ctx context.Context, method, path string,
 }
 
 func TestTelemetryLogger(t *testing.T) {
-	events := []telemetry.DatabricksCliLog{
+	events := []telemetry.FrontendLogEntry{
 		{
-			CliTestEvent: &events.CliTestEvent{
-				Name: events.DummyCliEnumValue1,
+			DatabricksCliLog: telemetry.DatabricksCliLog{
+				CliTestEvent: &events.CliTestEvent{
+					Name: events.DummyCliEnumValue1,
+				},
 			},
 		},
 		{
-			BundleInitEvent: &events.BundleInitEvent{
-				Uuid:         uuid.New().String(),
-				TemplateName: "abc",
-				TemplateEnumArgs: map[string]string{
-					"a": "b",
-					"c": "d",
+			DatabricksCliLog: telemetry.DatabricksCliLog{
+				BundleInitEvent: &events.BundleInitEvent{
+					Uuid:         uuid.New().String(),
+					TemplateName: "abc",
+					TemplateEnumArgs: map[string]string{
+						"a": "b",
+						"c": "d",
+					},
 				},
 			},
 		},
@@ -64,7 +68,8 @@ func TestTelemetryLogger(t *testing.T) {
 	})
 
 	for _, event := range events {
-		telemetry.Log(ctx, event)
+		err := telemetry.Log(ctx, event)
+		require.NoError(t, err)
 	}
 
 	apiClient, err := client.New(w.W.Config)

libs/telemetry/context.go

@@ -18,7 +18,7 @@ func ContextWithLogger(ctx context.Context) context.Context {
 		return ctx
 	}
 
-	return context.WithValue(ctx, telemetryLoggerKey, &logger{logs: []FrontendLog{}})
+	return context.WithValue(ctx, telemetryLoggerKey, &logger{protoLogs: []string{}})
 }
 
 func fromContext(ctx context.Context) *logger {

libs/telemetry/logger.go

@@ -3,6 +3,7 @@ package telemetry
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"time"
 
@@ -17,34 +18,50 @@ type DatabricksApiClient interface {
 		visitors ...func(*http.Request) error) error
 }
 
-func Log(ctx context.Context, event DatabricksCliLog) {
+func Log(ctx context.Context, event FrontendLogEntry) error {
 	l := fromContext(ctx)
 
-	l.logs = append(l.logs, FrontendLog{
+	FrontendLog := FrontendLog{
 		// The telemetry endpoint deduplicates logs based on the FrontendLogEventID.
 		// This it's important to generate a unique ID for each log event.
 		FrontendLogEventID: uuid.New().String(),
-		Entry: FrontendLogEntry{
-			DatabricksCliLog: event,
-		},
-	})
+		Entry:              event,
+	}
+
+	protoLog, err := json.Marshal(FrontendLog)
+	if err != nil {
+		return fmt.Errorf("error marshalling the telemetry event: %v", err)
+	}
+
+	l.protoLogs = append(l.protoLogs, string(protoLog))
+	return nil
 }
 
 type logger struct {
-	logs []FrontendLog
+	protoLogs []string
 }
 
 // This function is meant to be only to be used in tests to introspect the telemetry logs
 // that have been logged so far.
-func GetLogs(ctx context.Context) []FrontendLog {
+func GetLogs(ctx context.Context) ([]FrontendLog, error) {
 	l := fromContext(ctx)
-	return l.logs
+	res := []FrontendLog{}
+
+	for _, log := range l.protoLogs {
+		frontendLog := FrontendLog{}
+		err := json.Unmarshal([]byte(log), &frontendLog)
+		if err != nil {
+			return nil, fmt.Errorf("error unmarshalling the telemetry event: %v", err)
+		}
+		res = append(res, frontendLog)
+	}
+	return res, nil
 }
 
 // Maximum additional time to wait for the telemetry event to flush. We expect the flush
 // method to be called when the CLI command is about to exist, so this caps the maximum
 // additional time the user will experience because of us logging CLI telemetry.
-var MaxAdditionalWaitTime = 5 * time.Second
+var MaxAdditionalWaitTime = 2 * time.Second
 
 // We make the API call to the /telemetry-ext endpoint to log the CLI telemetry events
 // right about as the CLI command is about to exit. The API endpoint can handle
@@ -56,22 +73,11 @@ func Flush(ctx context.Context, apiClient DatabricksApiClient) {
 	defer cancel()
 	l := fromContext(ctx)
 
-	if len(l.logs) == 0 {
+	if len(l.protoLogs) == 0 {
 		log.Debugf(ctx, "No telemetry events to flush")
 		return
 	}
 
-	var protoLogs []string
-	for _, event := range l.logs {
-		s, err := json.Marshal(event)
-		if err != nil {
-			log.Debugf(ctx, "Error marshalling the telemetry event %v: %v", event, err)
-			continue
-		}
-
-		protoLogs = append(protoLogs, string(s))
-	}
-
 	resp := &ResponseBody{}
 	for {
 		select {
@@ -85,7 +91,7 @@ func Flush(ctx context.Context, apiClient DatabricksApiClient) {
 		// Log the CLI telemetry events.
 		err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", nil, RequestBody{
 			UploadTime: time.Now().Unix(),
-			ProtoLogs:  protoLogs,
+			ProtoLogs:  l.protoLogs,
 
 			// A bug in the telemetry API requires us to send an empty items array.
 			// Otherwise we get an opaque 500 internal server error.
@@ -102,7 +108,7 @@ func Flush(ctx context.Context, apiClient DatabricksApiClient) {
 		//
 		// Note: This will result in server side duplications but that's fine since
 		// we can always deduplicate in the data pipeline itself.
-		if len(l.logs) > int(resp.NumProtoSuccess) {
+		if len(l.protoLogs) > int(resp.NumProtoSuccess) {
 			log.Debugf(ctx, "Not all logs were successfully sent. Retrying...")
 			continue
 		}

libs/telemetry/logger_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/databricks/cli/libs/telemetry/events"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type mockDatabricksClient struct {
@@ -67,9 +68,10 @@ func TestTelemetryLoggerFlushesEvents(t *testing.T) {
 	ctx := ContextWithLogger(context.Background())
 
 	for _, v := range []events.DummyCliEnum{events.DummyCliEnumValue1, events.DummyCliEnumValue2, events.DummyCliEnumValue2, events.DummyCliEnumValue3} {
-		Log(ctx, DatabricksCliLog{
+		err := Log(ctx, FrontendLogEntry{DatabricksCliLog: DatabricksCliLog{
 			CliTestEvent: &events.CliTestEvent{Name: v},
-		})
+		}})
+		require.NoError(t, err)
 	}
 
 	// Flush the events.
@@ -100,9 +102,10 @@ func TestTelemetryLoggerFlushExitsOnTimeout(t *testing.T) {
 	ctx := ContextWithLogger(context.Background())
 
 	for _, v := range []events.DummyCliEnum{events.DummyCliEnumValue1, events.DummyCliEnumValue2, events.DummyCliEnumValue2, events.DummyCliEnumValue3} {
-		Log(ctx, DatabricksCliLog{
+		err := Log(ctx, FrontendLogEntry{DatabricksCliLog: DatabricksCliLog{
 			CliTestEvent: &events.CliTestEvent{Name: v},
-		})
+		}})
+		require.NoError(t, err)
 	}
 
 	// Flush the events.

libs/template/materialize.go

@@ -124,8 +124,9 @@ func (t *Template) logTelemetry(ctx context.Context) error {
 		},
 	}
 
-	telemetry.Log(ctx, event)
-	return nil
+	return telemetry.Log(ctx, telemetry.FrontendLogEntry{
+		DatabricksCliLog: event,
+	})
 }
 
 // This function materializes the input templates as a project, using user defined