diff --git a/acceptance/telemetry/upload-succeeds/script b/acceptance/telemetry/upload-succeeds/script new file mode 100644 index 000000000..b11cb6e99 --- /dev/null +++ b/acceptance/telemetry/upload-succeeds/script @@ -0,0 +1,12 @@ +export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid +export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt + +trace $CLI selftest send-telemetry + +echo "waiting for telemetry process to finish" + +# Wait for the child telemetry process to finish +wait_pid $(cat ./telemetry.pid) + +# cleanup the pid file +rm -f ./telemetry.pid diff --git a/acceptance/telemetry/upload-timeout/out.requests.txt b/acceptance/telemetry/upload-timeout/out.requests.txt new file mode 100644 index 000000000..e69de29bb diff --git a/acceptance/telemetry/upload-timeout/out.upload_process.txt b/acceptance/telemetry/upload-timeout/out.upload_process.txt new file mode 100644 index 000000000..8cdfbf20e --- /dev/null +++ b/acceptance/telemetry/upload-timeout/out.upload_process.txt @@ -0,0 +1 @@ +error: Failed to flush telemetry log due to timeout diff --git a/acceptance/telemetry/upload-timeout/output.txt b/acceptance/telemetry/upload-timeout/output.txt new file mode 100644 index 000000000..1a1f61588 --- /dev/null +++ b/acceptance/telemetry/upload-timeout/output.txt @@ -0,0 +1,4 @@ + +>>> [CLI] selftest send-telemetry +waiting for telemetry process to finish +[wait_pid] process has ended diff --git a/acceptance/telemetry/upload-timeout/script b/acceptance/telemetry/upload-timeout/script new file mode 100644 index 000000000..d81eba279 --- /dev/null +++ b/acceptance/telemetry/upload-timeout/script @@ -0,0 +1,16 @@ +export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid +export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt + +# Configure a timeout of 0 seconds. This ensures that the timeout is respected and the +# telemetry process does not try to upload logs. +export DATABRICKS_CLI_TELEMETRY_UPLOAD_TIMEOUT="0s" + +trace $CLI selftest send-telemetry + +echo "waiting for telemetry process to finish" + +# Wait for the child telemetry process to finish +wait_pid $(cat ./telemetry.pid) + +# cleanup the pid file +rm -f ./telemetry.pid diff --git a/acceptance/telemetry/upload-timeout/test.toml b/acceptance/telemetry/upload-timeout/test.toml new file mode 100644 index 000000000..4826f42b0 --- /dev/null +++ b/acceptance/telemetry/upload-timeout/test.toml @@ -0,0 +1,11 @@ +[[Server]] +Pattern = "POST /telemetry-ext" +Response.Body = ''' +{ + "numProtoSuccess": 2 +} +''' + +[[Repls]] +Old = 'execution_time_ms\\\":\d{1,5},' +New = 'execution_time_ms\":\"SMALL_INT\",' diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 5bcadfe1b..927c6c4a6 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -186,8 +186,7 @@ func workspaceClientOrPrompt(ctx context.Context, cfg *config.Config, allowPromp return w, err } -// TODO: Make upload with oauth work. -// TODO: Move env var inheritance to the daemon library. +// TODO: Run as integration tests? func MustWorkspaceClient(cmd *cobra.Command, args []string) error { cfg := &config.Config{} diff --git a/cmd/root/root.go b/cmd/root/root.go index 8f963f381..de92bda98 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -103,6 +103,7 @@ func flagErrorFunc(c *cobra.Command, err error) error { // Execute adds all child commands to the root command and sets flags appropriately. // This is called by main.main(). It only needs to happen once to the rootCmd. func Execute(ctx context.Context, cmd *cobra.Command) error { + // TODO: deferred panic recovery ctx = telemetry.WithNewLogger(ctx) ctx = dbr.DetectRuntime(ctx) start := time.Now() diff --git a/libs/telemetry/upload.go b/libs/telemetry/upload.go index 6aadc67ab..72b68fc48 100644 --- a/libs/telemetry/upload.go +++ b/libs/telemetry/upload.go @@ -25,6 +25,9 @@ const ( // Environment variable to disable telemetry. If this is set to any value, telemetry // will be disabled. DisableEnvVar = "DATABRICKS_CLI_DISABLE_TELEMETRY" + + // Max time to try and upload the telemetry logs. Useful for testing. + UploadTimeout = "DATABRICKS_CLI_TELEMETRY_UPLOAD_TIMEOUT" ) type UploadConfig struct { @@ -68,8 +71,16 @@ func Upload() (*ResponseBody, error) { return nil, fmt.Errorf("Failed to create API client: %s\n", err) } + maxUploadTime := 30 * time.Second + if v, ok := os.LookupEnv(UploadTimeout); ok { + maxUploadTime, err = time.ParseDuration(v) + if err != nil { + return nil, fmt.Errorf("Failed to parse time limit %s: %s\n", UploadTimeout, err) + } + } + // Set a maximum total time to try telemetry uploads. - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), maxUploadTime) defer cancel() resp := &ResponseBody{}