diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 22c4f1a4c..98ef72b76 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -36,7 +36,7 @@ var ( // In order to debug CLI running under acceptance test, set this to full subtest name, e.g. "bundle/variables/empty" // Then install your breakpoints and click "debug test" near TestAccept in VSCODE. // example: var SingleTest = "bundle/variables/empty" -var SingleTest = "" +var SingleTest = "telemetry/dummy-without-server" // If enabled, instead of compiling and running CLI externally, we'll start in-process server that accepts and runs // CLI commands. The $CLI in test scripts is a helper that just forwards command-line arguments to this server (see bin/callserver.py). @@ -44,7 +44,8 @@ var SingleTest = "" var InprocessMode bool func init() { - flag.BoolVar(&InprocessMode, "inprocess", SingleTest != "", "Run CLI in the same process as test (for debugging)") + // flag.BoolVar(&InprocessMode, "inprocess", SingleTest != "", "Run CLI in the same process as test (for debugging)") + flag.BoolVar(&InprocessMode, "inprocess", false, "Run CLI in the same process as test (for debugging)") flag.BoolVar(&KeepTmp, "keeptmp", false, "Do not delete TMP directory after run") flag.BoolVar(&NoRepl, "norepl", false, "Do not apply any replacements (for debugging)") } diff --git a/acceptance/bin/wait_file b/acceptance/bin/wait_file new file mode 100755 index 000000000..5fa1ab69c --- /dev/null +++ b/acceptance/bin/wait_file @@ -0,0 +1,27 @@ +#!/bin/bash + +wait_file() { + local file_path="$1" + local max_attempts=100 + local attempt=0 + + while [ $attempt -lt $max_attempts ]; do + if [ -e "$file_path" ]; then + echo "File $file_path exists" + return 0 + fi + sleep 0.1 + attempt=$((attempt + 1)) + done + + echo "Timeout: File $file_path did not appear within 10 seconds" + return 1 +} + +if [ $# -eq 0 ]; then + echo "Usage: $0 " + exit 1 +fi + +wait_file "$1" +exit $? diff --git a/acceptance/bin/wait_pid b/acceptance/bin/wait_pid new file mode 100755 index 000000000..a71781725 --- /dev/null +++ b/acceptance/bin/wait_pid @@ -0,0 +1,41 @@ +#!/bin/bash + + +# wait in bash only works for child process. This script is more general. +wait_pid() { + local pid=$1 + local max_attempts=100 # 100 * 0.1 seconds = 10 seconds + local attempt=0 + local sleep_time=0.1 + + while [ $attempt -lt $max_attempts ]; do + if [[ "$OSTYPE" == "msys"* || "$OSTYPE" == "cygwin"* ]]; then + # Windows approach + if ! tasklist | grep -q $pid; then + echo "Process has ended" + return 0 + fi + else + # Linux/macOS approach + if ! kill -0 $pid 2>/dev/null; then + echo "Process has ended" + return 0 + fi + fi + + sleep $sleep_time + attempt=$((attempt + 1)) + done + + echo "Timeout: Process $pid did not end within 10 seconds" + return 1 +} + +# Usage +if [ $# -eq 0 ]; then + echo "Usage: $0 " + exit 1 +fi + +wait_pid $1 +exit $? diff --git a/acceptance/telemetry/dummy-without-server/out.upload.txt b/acceptance/telemetry/dummy-without-server/out.upload.txt deleted file mode 100644 index 563c5f6fd..000000000 --- a/acceptance/telemetry/dummy-without-server/out.upload.txt +++ /dev/null @@ -1 +0,0 @@ -error: telemetry endpoint not found diff --git a/acceptance/telemetry/dummy-without-server/output.txt b/acceptance/telemetry/dummy-without-server/output.txt index a7b80d140..cf0956eec 100644 --- a/acceptance/telemetry/dummy-without-server/output.txt +++ b/acceptance/telemetry/dummy-without-server/output.txt @@ -1,2 +1,4 @@ >>> $CLI telemetry dummy +File ./telemetry.pid exists +Process has ended diff --git a/acceptance/telemetry/dummy-without-server/script b/acceptance/telemetry/dummy-without-server/script index 2bc8dc082..f91cb69d1 100644 --- a/acceptance/telemetry/dummy-without-server/script +++ b/acceptance/telemetry/dummy-without-server/script @@ -1,5 +1,9 @@ -export DATABRICKS_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt +export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid # This test ensures that the main CLI command does not error even if # telemetry upload fails. trace $CLI telemetry dummy + +# Wait for the child telemetry process to finish +trace wait_file ./telemetry.pid +trace wait_pid $(cat ./telemetry.pid) diff --git a/acceptance/telemetry/dummy/out.upload.txt b/acceptance/telemetry/dummy/out.upload.txt deleted file mode 100644 index 082f3bb90..000000000 --- a/acceptance/telemetry/dummy/out.upload.txt +++ /dev/null @@ -1,3 +0,0 @@ -Successfully uploaded telemetry logs -Response: -{"errors":null,"numProtoSuccess":2} diff --git a/acceptance/telemetry/dummy/output.txt b/acceptance/telemetry/dummy/output.txt index a7b80d140..cf0956eec 100644 --- a/acceptance/telemetry/dummy/output.txt +++ b/acceptance/telemetry/dummy/output.txt @@ -1,2 +1,4 @@ >>> $CLI telemetry dummy +File ./telemetry.pid exists +Process has ended diff --git a/acceptance/telemetry/dummy/script b/acceptance/telemetry/dummy/script index d60217b23..2bcb85299 100644 --- a/acceptance/telemetry/dummy/script +++ b/acceptance/telemetry/dummy/script @@ -1,3 +1,7 @@ -export DATABRICKS_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt +export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid trace $CLI telemetry dummy + +# Wait for the child telemetry process to finish +trace wait_file ./telemetry.pid +trace wait_pid $(cat ./telemetry.pid) diff --git a/acceptance/telemetry/dummy/test.toml b/acceptance/telemetry/dummy/test.toml index 48e39414e..a7864dbaf 100644 --- a/acceptance/telemetry/dummy/test.toml +++ b/acceptance/telemetry/dummy/test.toml @@ -1,7 +1,4 @@ RecordRequests = true -EventuallyFiles = [ - "out.upload.txt" -] [[Server]] Pattern = "POST /telemetry-ext" diff --git a/cmd/root/root.go b/cmd/root/root.go index 31f13ef38..aa1b92d6f 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -10,6 +10,7 @@ import ( "os/exec" "runtime" "slices" + "strconv" "strings" "time" @@ -164,6 +165,8 @@ 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) { @@ -213,6 +216,13 @@ func uploadTelemetry(ctx context.Context, cmdStr string, start, end time.Time, e return } + if pidFilePath := env.Get(ctx, telemetryPidFileEnvVar); 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) + } + } + _, err = stdin.Write(b) if err != nil { log.Debugf(ctx, "failed to write to telemetry worker: %s", err)