diff --git a/acceptance/internal/config.go b/acceptance/internal/config.go index 2129fa5573..3b5a3aadc2 100644 --- a/acceptance/internal/config.go +++ b/acceptance/internal/config.go @@ -156,6 +156,10 @@ type ServerStub struct { // Configure as "1ms", "2s", "3m", etc. // See [time.ParseDuration] for details. Delay time.Duration + + // If true, kill the caller process instead of returning a response. + // Requires DATABRICKS_CLI_TEST_PID=1 to be set in the test environment. + KillCaller bool } // FindConfigs finds all the config relevant for this test, diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index a401a1cac8..200532e7e4 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -209,6 +209,44 @@ func startLocalServer(t *testing.T, } } + if stub.KillCaller { + pid := testserver.ExtractPidFromHeaders(req.Headers) + if pid == 0 { + t.Errorf("KillCaller configured but test-pid not found in User-Agent") + return testserver.Response{ + StatusCode: http.StatusBadRequest, + Body: "test-pid not found in User-Agent (set DATABRICKS_CLI_TEST_PID=1)", + } + } + + t.Logf("KillCaller: killing PID %d (pattern: %s)", pid, stub.Pattern) + + process, err := os.FindProcess(pid) + if err != nil { + t.Errorf("Failed to find process %d: %s", pid, err) + return testserver.Response{ + StatusCode: http.StatusInternalServerError, + Body: fmt.Sprintf("failed to find process: %s", err), + } + } + + // Use process.Kill() for cross-platform compatibility. + // On Unix, this sends SIGKILL. On Windows, this calls TerminateProcess. + if err := process.Kill(); err != nil { + t.Errorf("Failed to kill process %d: %s", pid, err) + return testserver.Response{ + StatusCode: http.StatusInternalServerError, + Body: fmt.Sprintf("failed to kill process: %s", err), + } + } + + // Return a response (the CLI will likely be killed before it receives this) + return testserver.Response{ + StatusCode: http.StatusOK, + Body: fmt.Sprintf("killed PID %d", pid), + } + } + return stub.Response }) } diff --git a/acceptance/selftest/kill_caller/mid_request/out.test.toml b/acceptance/selftest/kill_caller/mid_request/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/selftest/kill_caller/mid_request/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/selftest/kill_caller/mid_request/output.txt b/acceptance/selftest/kill_caller/mid_request/output.txt new file mode 100644 index 0000000000..fac04c326a --- /dev/null +++ b/acceptance/selftest/kill_caller/mid_request/output.txt @@ -0,0 +1,5 @@ + +>>> errcode [CLI] workspace list / +[PROCESS_KILLED] + +Exit code: [KILLED] diff --git a/acceptance/selftest/kill_caller/mid_request/script b/acceptance/selftest/kill_caller/mid_request/script new file mode 100644 index 0000000000..817a90cc76 --- /dev/null +++ b/acceptance/selftest/kill_caller/mid_request/script @@ -0,0 +1 @@ +trace errcode $CLI workspace list / diff --git a/acceptance/selftest/kill_caller/mid_request/test.toml b/acceptance/selftest/kill_caller/mid_request/test.toml new file mode 100644 index 0000000000..dedb9fd7a3 --- /dev/null +++ b/acceptance/selftest/kill_caller/mid_request/test.toml @@ -0,0 +1,30 @@ +Local = true +Env.DATABRICKS_CLI_TEST_PID = "1" + +[[Server]] +Pattern = "GET /api/2.0/workspace/list" +KillCaller = true + +[[Repls]] +# macOS bash shows "Killed: 9" (with signal number), Linux shows "Killed" +# Normalize the whole killed line to a placeholder +Old = 'script: line \d+:\s+\d+ Killed(: 9)?\s+"\$@"' +New = '[PROCESS_KILLED]' + +[[Repls]] +# On Windows, there's no "Killed" message - just empty line before Exit code +# Insert [PROCESS_KILLED] placeholder for consistency +Old = '(\n>>> errcode [^\n]+\n)\nExit code:' +New = """${1}[PROCESS_KILLED] + +Exit code:""" + +[[Repls]] +# Normalize exit code: 137 on Unix (128 + SIGKILL), 1 on Windows +Old = 'Exit code: (137|1)' +New = 'Exit code: [KILLED]' + +[[Repls]] +# Normalize Windows line endings (CRLF -> LF) - must be LAST +Old = "\r" +New = '' diff --git a/acceptance/selftest/kill_caller/out.test.toml b/acceptance/selftest/kill_caller/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/selftest/kill_caller/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/selftest/kill_caller/output.txt b/acceptance/selftest/kill_caller/output.txt new file mode 100644 index 0000000000..2e413420b8 --- /dev/null +++ b/acceptance/selftest/kill_caller/output.txt @@ -0,0 +1,5 @@ + +>>> errcode [CLI] current-user me +[PROCESS_KILLED] + +Exit code: [KILLED] diff --git a/acceptance/selftest/kill_caller/script b/acceptance/selftest/kill_caller/script new file mode 100644 index 0000000000..b6e8c070f6 --- /dev/null +++ b/acceptance/selftest/kill_caller/script @@ -0,0 +1 @@ +trace errcode $CLI current-user me diff --git a/acceptance/selftest/kill_caller/test.toml b/acceptance/selftest/kill_caller/test.toml new file mode 100644 index 0000000000..0fd483f97a --- /dev/null +++ b/acceptance/selftest/kill_caller/test.toml @@ -0,0 +1,31 @@ +Local = true +Env.DATABRICKS_CLI_TEST_PID = "1" + +# Kill the CLI when it calls /Me endpoint +[[Server]] +Pattern = "GET /api/2.0/preview/scim/v2/Me" +KillCaller = true + +[[Repls]] +# macOS bash shows "Killed: 9" (with signal number), Linux shows "Killed" +# Normalize the whole killed line to a placeholder +Old = 'script: line \d+:\s+\d+ Killed(: 9)?\s+"\$@"' +New = '[PROCESS_KILLED]' + +[[Repls]] +# On Windows, there's no "Killed" message - just empty line before Exit code +# Insert [PROCESS_KILLED] placeholder for consistency +Old = '(\n>>> errcode [^\n]+\n)\nExit code:' +New = """${1}[PROCESS_KILLED] + +Exit code:""" + +[[Repls]] +# Normalize exit code: 137 on Unix (128 + SIGKILL), 1 on Windows +Old = 'Exit code: (137|1)' +New = 'Exit code: [KILLED]' + +[[Repls]] +# Normalize Windows line endings (CRLF -> LF) - must be LAST +Old = "\r" +New = '' diff --git a/cmd/pipelines/root/root.go b/cmd/pipelines/root/root.go index 72d0b52664..1eb874067c 100644 --- a/cmd/pipelines/root/root.go +++ b/cmd/pipelines/root/root.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/internal/build" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/testserver" "github.com/spf13/cobra" ) @@ -72,6 +73,7 @@ func New(ctx context.Context) *cobra.Command { ctx = withCommandInUserAgent(ctx, cmd) ctx = withCommandExecIdInUserAgent(ctx) ctx = withUpstreamInUserAgent(ctx) + ctx = testserver.InjectPidToUserAgent(ctx) cmd.SetContext(ctx) return nil } diff --git a/cmd/root/root.go b/cmd/root/root.go index 96fde20846..20bf96691d 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/telemetry" "github.com/databricks/cli/libs/telemetry/protos" + "github.com/databricks/cli/libs/testserver" "github.com/spf13/cobra" ) @@ -79,6 +80,7 @@ func New(ctx context.Context) *cobra.Command { ctx = withCommandInUserAgent(ctx, cmd) ctx = withCommandExecIdInUserAgent(ctx) ctx = withUpstreamInUserAgent(ctx) + ctx = testserver.InjectPidToUserAgent(ctx) cmd.SetContext(ctx) return nil } diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 8b8e346e99..d9f3407858 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -9,15 +9,46 @@ import ( "net/http" "net/http/httptest" "net/url" + "os" "reflect" + "regexp" + "strconv" "strings" "sync" + "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go/useragent" "github.com/gorilla/mux" +) - "github.com/databricks/cli/internal/testutil" +const ( + TestPidEnvVar = "DATABRICKS_CLI_TEST_PID" + testPidKey = "test-pid" ) +var testPidRegex = regexp.MustCompile(testPidKey + `/(\d+)`) + +func InjectPidToUserAgent(ctx context.Context) context.Context { + if env.Get(ctx, TestPidEnvVar) != "1" { + return ctx + } + return useragent.InContext(ctx, testPidKey, strconv.Itoa(os.Getpid())) +} + +func ExtractPidFromHeaders(headers http.Header) int { + ua := headers.Get("User-Agent") + matches := testPidRegex.FindStringSubmatch(ua) + if len(matches) < 2 { + return 0 + } + pid, err := strconv.Atoi(matches[1]) + if err != nil { + return 0 + } + return pid +} + type Server struct { *httptest.Server Router *mux.Router @@ -274,6 +305,9 @@ func (s *Server) Handle(method, path string, handler HandlerFunc) { StatusCode: 500, Body: []byte("INJECTED"), } + } else if bytes.Contains(request.Body, []byte("KILL_CALLER")) { + s.handleKillCaller(&request, w) + return } else { respAny := handler(request) if respAny == nil && request.Context.Err() != nil { @@ -322,3 +356,35 @@ func isNil(i any) bool { return false } } + +func (s *Server) handleKillCaller(request *Request, w http.ResponseWriter) { + pid := ExtractPidFromHeaders(request.Headers) + if pid == 0 { + s.t.Errorf("KILL_CALLER requested but test-pid not found in User-Agent") + w.WriteHeader(http.StatusBadRequest) + _, _ = fmt.Fprint(w, "test-pid not found in User-Agent (set DATABRICKS_CLI_TEST_PID=1)") + return + } + + s.t.Logf("KILL_CALLER: killing PID %d", pid) + + process, err := os.FindProcess(pid) + if err != nil { + s.t.Errorf("Failed to find process %d: %s", pid, err) + w.WriteHeader(http.StatusInternalServerError) + _, _ = fmt.Fprintf(w, "failed to find process: %s", err) + return + } + + // Use process.Kill() for cross-platform compatibility. + // On Unix, this sends SIGKILL. On Windows, this calls TerminateProcess. + if err := process.Kill(); err != nil { + s.t.Errorf("Failed to kill process %d: %s", pid, err) + w.WriteHeader(http.StatusInternalServerError) + _, _ = fmt.Fprintf(w, "failed to kill process: %s", err) + return + } + + w.WriteHeader(http.StatusOK) + _, _ = fmt.Fprintf(w, "killed PID %d", pid) +}