From d5e03f08d56c1b6fc71516435d4be42143b50edf Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Tue, 4 Feb 2025 16:25:59 +0100 Subject: [PATCH] - --- .../dummy-without-server/out.upload.txt | 26 +++++++++++ .../telemetry/dummy-without-server/script | 1 + .../telemetry/dummy-without-server/test.toml | 5 -- acceptance/telemetry/dummy/out.upload.txt | 3 ++ acceptance/telemetry/dummy/script | 1 + acceptance/telemetry/upload/out.upload.txt | 3 ++ acceptance/telemetry/upload/output.txt | 3 -- acceptance/telemetry/upload/script | 2 + cmd/root/root.go | 4 +- libs/telemetry/upload.go | 46 +++++++------------ libs/telemetry/upload_test.go | 4 +- main.go | 20 ++++++-- 12 files changed, 71 insertions(+), 47 deletions(-) create mode 100644 acceptance/telemetry/dummy-without-server/out.upload.txt delete mode 100644 acceptance/telemetry/dummy-without-server/test.toml create mode 100644 acceptance/telemetry/dummy/out.upload.txt create mode 100644 acceptance/telemetry/upload/out.upload.txt diff --git a/acceptance/telemetry/dummy-without-server/out.upload.txt b/acceptance/telemetry/dummy-without-server/out.upload.txt new file mode 100644 index 000000000..2d2b20cb4 --- /dev/null +++ b/acceptance/telemetry/dummy-without-server/out.upload.txt @@ -0,0 +1,26 @@ +error: Failed to upload telemetry logs: unable to parse response. This is likely a bug in the Databricks SDK for Go or the underlying REST API. Please report this issue with the following debugging information to the SDK issue tracker at https://github.com/databricks/databricks-sdk-go/issues. Request log: +``` +POST /telemetry-ext +> * Host: +> * Accept: application/json +> * Authorization: REDACTED +> * Content-Type: application/json +> * Traceparent: 00-1997351cf21b2042c96c8bf415bbcfe8-7542779df43bc559-01 +> * User-Agent: cli/$DEV_VERSION databricks-sdk-go/0.56.1 go/1.23.4 os/darwin auth/pat +> { +> "items": null, +> "protoLogs": [ +> "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"$DEV_VERSION\",\"command\":\"telemetry_dummy\",\"operating_system\":\"darwin\",\"execution_time_ms\":0,\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}", +> "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"$DEV_VERSION\",\"command\":\"telemetry_dummy\",\"operating_system\":\"darwin\",\"execution_time_ms\":0,\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}" +> ], +> "uploadTime": 1738682716822 +> } +< HTTP/1.1 404 Not Found +< * Content-Length: 19 +< * Content-Type: text/plain; charset=utf-8 +< * Date: Tue, 04 Feb 2025 15:25:16 GMT +< * X-Content-Type-Options: nosniff +< 404 page not found +< +``` + diff --git a/acceptance/telemetry/dummy-without-server/script b/acceptance/telemetry/dummy-without-server/script index 4a7b3151f..b4be7f986 100644 --- a/acceptance/telemetry/dummy-without-server/script +++ b/acceptance/telemetry/dummy-without-server/script @@ -1,4 +1,5 @@ export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid +export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt # This test ensures that the main CLI command does not error even if # telemetry upload fails. diff --git a/acceptance/telemetry/dummy-without-server/test.toml b/acceptance/telemetry/dummy-without-server/test.toml deleted file mode 100644 index fdaf052da..000000000 --- a/acceptance/telemetry/dummy-without-server/test.toml +++ /dev/null @@ -1,5 +0,0 @@ -# TODO: We do not run this on windows because we don't have proper stubbing -# for API errors yet. Simply not defining a server stub does not return a 404 on windows. -# Run this on windows as well once we can stub API errors. -[GOOS] -windows = false diff --git a/acceptance/telemetry/dummy/out.upload.txt b/acceptance/telemetry/dummy/out.upload.txt new file mode 100644 index 000000000..7ee6e3b6a --- /dev/null +++ b/acceptance/telemetry/dummy/out.upload.txt @@ -0,0 +1,3 @@ +Telemetry logs uploaded successfully +Response: +{"errors":null,"numProtoSuccess":2} diff --git a/acceptance/telemetry/dummy/script b/acceptance/telemetry/dummy/script index cc291e941..af8c69071 100644 --- a/acceptance/telemetry/dummy/script +++ b/acceptance/telemetry/dummy/script @@ -1,4 +1,5 @@ export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid +export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt trace $CLI telemetry dummy diff --git a/acceptance/telemetry/upload/out.upload.txt b/acceptance/telemetry/upload/out.upload.txt new file mode 100644 index 000000000..7ee6e3b6a --- /dev/null +++ b/acceptance/telemetry/upload/out.upload.txt @@ -0,0 +1,3 @@ +Telemetry logs uploaded successfully +Response: +{"errors":null,"numProtoSuccess":2} diff --git a/acceptance/telemetry/upload/output.txt b/acceptance/telemetry/upload/output.txt index c54d49324..0254788a7 100644 --- a/acceptance/telemetry/upload/output.txt +++ b/acceptance/telemetry/upload/output.txt @@ -1,5 +1,2 @@ >>> $CLI telemetry upload -Successfully uploaded telemetry logs -Response: -{"errors":null,"numProtoSuccess":2} diff --git a/acceptance/telemetry/upload/script b/acceptance/telemetry/upload/script index 595a5303b..7312d8359 100644 --- a/acceptance/telemetry/upload/script +++ b/acceptance/telemetry/upload/script @@ -1,3 +1,5 @@ +export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt + # This command / test cannot be run in inprocess / debug mode. This is because # it does not go through the [root.Execute] function. trace $CLI telemetry upload < stdin diff --git a/cmd/root/root.go b/cmd/root/root.go index aa1b92d6f..ee2959f72 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -165,8 +165,6 @@ func inheritEnvVars() []string { return out } -const telemetryPidFileEnvVar = "DATABRICKS_CLI_TELEMETRY_PID_FILE" - func uploadTelemetry(ctx context.Context, cmdStr string, start, end time.Time, exitCode int) { // Nothing to upload. if !telemetry.HasLogs(ctx) { @@ -216,7 +214,7 @@ func uploadTelemetry(ctx context.Context, cmdStr string, start, end time.Time, e return } - if pidFilePath := env.Get(ctx, telemetryPidFileEnvVar); pidFilePath != "" { + if pidFilePath := env.Get(ctx, telemetry.PidFileEnvVar); pidFilePath != "" { err = os.WriteFile(pidFilePath, []byte(strconv.Itoa(telemetryCmd.Process.Pid)), 0o644) if err != nil { log.Debugf(ctx, "failed to write telemetry worker PID file: %s", err) diff --git a/libs/telemetry/upload.go b/libs/telemetry/upload.go index 885d4da74..5db9a5e9f 100644 --- a/libs/telemetry/upload.go +++ b/libs/telemetry/upload.go @@ -11,13 +11,17 @@ import ( "time" "github.com/databricks/cli/libs/telemetry/protos" - "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" ) // File containing debug logs from the upload process. -const UploadLogsFileEnvVar = "DATABRICKS_TELEMETRY_UPLOAD_LOGS_FILE" +const ( + UploadLogsFileEnvVar = "DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE" + + // File containing the PID of the telemetry upload process. + PidFileEnvVar = "DATABRICKS_CLI_TELEMETRY_PID_FILE" +) type UploadConfig struct { Logs []protos.FrontendLog `json:"logs"` @@ -26,30 +30,29 @@ type UploadConfig struct { // Upload reads telemetry logs from stdin and uploads them to the telemetry endpoint. // This function is always expected to be called in a separate child process from // the main CLI process. -func Upload() error { +func Upload() (*ResponseBody, error) { var err error b, err := io.ReadAll(os.Stdin) if err != nil { - return fmt.Errorf("failed to read from stdin: %s\n", err) + return nil, fmt.Errorf("failed to read from stdin: %s\n", err) } in := UploadConfig{} err = json.Unmarshal(b, &in) if err != nil { - printJson(in) - return fmt.Errorf("failed to unmarshal input: %s\n", err) + return nil, fmt.Errorf("failed to unmarshal input: %s\n", err) } if len(in.Logs) == 0 { - return fmt.Errorf("No logs to upload: %s\n", err) + return nil, fmt.Errorf("No logs to upload: %s\n", err) } protoLogs := make([]string, len(in.Logs)) for i, log := range in.Logs { b, err := json.Marshal(log) if err != nil { - return fmt.Errorf("failed to marshal log: %s\n", err) + return nil, fmt.Errorf("failed to marshal log: %s\n", err) } protoLogs[i] = string(b) } @@ -58,7 +61,7 @@ func Upload() error { // configure authentication. apiClient, err := client.New(&config.Config{}) if err != nil { - return fmt.Errorf("Failed to create API client: %s\n", err) + return nil, fmt.Errorf("Failed to create API client: %s\n", err) } // Set a maximum total time to try telemetry uploads. @@ -69,7 +72,7 @@ func Upload() error { for { select { case <-ctx.Done(): - return errors.New("Failed to flush telemetry log due to timeout") + return nil, errors.New("Failed to flush telemetry log due to timeout") default: // Proceed @@ -81,33 +84,16 @@ func Upload() error { Items: []string{}, ProtoLogs: protoLogs, }, resp) - var apierr *apierr.APIError - if errors.As(err, &apierr) && apierr.StatusCode == http.StatusNotFound { - return errors.New("telemetry endpoint not found") - } if err != nil { - return fmt.Errorf("Failed to upload telemetry logs: %s\n", err) + return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", err) } if len(resp.Errors) > 0 { - return fmt.Errorf("Failed to upload telemetry logs: %s\n", resp.Errors) + return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", resp.Errors) } if resp.NumProtoSuccess == int64(len(in.Logs)) { - fmt.Println("Successfully uploaded telemetry logs") - fmt.Println("Response: ") - printJson(resp) - break + return resp, nil } } - - return nil -} - -func printJson(v any) { - b, err := json.Marshal(v) - if err != nil { - fmt.Printf("failed to marshal JSON: %s\n", err) - } - fmt.Println(string(b)) } diff --git a/libs/telemetry/upload_test.go b/libs/telemetry/upload_test.go index 9db975fe6..01ea79bcc 100644 --- a/libs/telemetry/upload_test.go +++ b/libs/telemetry/upload_test.go @@ -79,8 +79,8 @@ func TestTelemetryUpload(t *testing.T) { os.Stdin = old }) - err = Upload() + resp, err := Upload() require.NoError(t, err) - + assert.Equal(t, int64(2), resp.NumProtoSuccess) assert.Equal(t, 2, count) } diff --git a/main.go b/main.go index 88c627fbf..3150e79ad 100644 --- a/main.go +++ b/main.go @@ -2,7 +2,9 @@ package main import ( "context" + "encoding/json" "fmt" + "io" "os" "github.com/databricks/cli/cmd" @@ -22,18 +24,28 @@ func main() { if len(os.Args) == 3 && os.Args[1] == "telemetry" && os.Args[2] == "upload" { var err error + // By default, this command should not write anything to stdout or stderr. + outW := io.Discard + errW := io.Discard + // If the environment variable is set, redirect stdout to the file. // This is useful for testing. if v := os.Getenv(telemetry.UploadLogsFileEnvVar); v != "" { - os.Stdout, _ = os.OpenFile(v, os.O_CREATE|os.O_WRONLY, 0o644) - os.Stderr = os.Stdout + outW, _ = os.OpenFile(v, os.O_CREATE|os.O_WRONLY, 0o644) + errW = outW } - err = telemetry.Upload() + resp, err := telemetry.Upload() if err != nil { - fmt.Fprintf(os.Stderr, "error: %s\n", err) + fmt.Fprintf(errW, "error: %s\n", err) os.Exit(1) } + fmt.Fprintf(outW, "Telemetry logs uploaded successfully\n") + fmt.Fprintln(outW, "Response:") + b, err := json.Marshal(resp) + if err == nil { + fmt.Fprintln(outW, string(b)) + } os.Exit(0) }