From 97740468254ca9953fc3d5cc6f0255d5991c5c74 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 5 Mar 2025 16:27:26 +0100
Subject: [PATCH] some cleanup

---
 acceptance/telemetry/timeout/out.requests.txt | 4 ++--
 acceptance/telemetry/timeout/output.txt       | 4 ++--
 acceptance/telemetry/timeout/test.toml        | 9 ---------
 libs/telemetry/logger.go                      | 6 +++---
 4 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/acceptance/telemetry/timeout/out.requests.txt b/acceptance/telemetry/timeout/out.requests.txt
index 6fb834be7..57574b976 100644
--- a/acceptance/telemetry/timeout/out.requests.txt
+++ b/acceptance/telemetry/timeout/out.requests.txt
@@ -10,8 +10,8 @@
     "uploadTime": "UNIX_TIME_MILLIS",
     "items": [],
     "protoLogs": [
-      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"[OS]\"},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
-      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"[OS]\"},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"[OS]\",\"execution_time_ms\":1},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"[OS]\",\"execution_time_ms\":1},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
     ]
   }
 }
diff --git a/acceptance/telemetry/timeout/output.txt b/acceptance/telemetry/timeout/output.txt
index be3e332a0..8dad2627a 100644
--- a/acceptance/telemetry/timeout/output.txt
+++ b/acceptance/telemetry/timeout/output.txt
@@ -6,8 +6,8 @@ HH:MM:SS Debug: POST /telemetry-ext
 > {
 >   "items": null,
 >   "protoLogs": [
->     "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"... (210 more bytes)",
->     "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"... (210 more bytes)"
+>     "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"... (232 more bytes)",
+>     "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"... (232 more bytes)"
 >   ],
 >   "uploadTime": "UNIX_TIME_MILLIS"
 > }
diff --git a/acceptance/telemetry/timeout/test.toml b/acceptance/telemetry/timeout/test.toml
index 65da6edb0..7e1b4d5b2 100644
--- a/acceptance/telemetry/timeout/test.toml
+++ b/acceptance/telemetry/timeout/test.toml
@@ -7,12 +7,3 @@ Response.Body = '''
 }
 '''
 DelaySeconds = 5
-
-
-[[Repls]]
-Old = "(?:[01][0-9]|2[0-3]):[0-5][0-9]:[0-5][0-9]"
-New = "HH:MM:SS"
-
-[[Repls]]
-Old = "pid=[0-9]+"
-New = "pid=PID"
diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go
index a060f47cc..ae7cc561b 100644
--- a/libs/telemetry/logger.go
+++ b/libs/telemetry/logger.go
@@ -78,8 +78,8 @@ func Upload(ctx context.Context, cfg *config.Config) error {
 		return fmt.Errorf("failed to create API client: %w", err)
 	}
 
-	ctx, _ = context.WithTimeout(ctx, 3*time.Second)
-	var resp *ResponseBody
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
 
 	// Only try uploading logs for a maximum of 3 times.
 	for i := range 3 {
@@ -90,7 +90,7 @@ func Upload(ctx context.Context, cfg *config.Config) error {
 			// proceed
 		}
 
-		resp, err = attempt(ctx, apiClient, protoLogs)
+		resp, err := attempt(ctx, apiClient, protoLogs)
 
 		// All logs were uploaded successfully.
 		if err == nil && resp.NumProtoSuccess >= int64(len(protoLogs)) {