From 788aca5bb7ee104a32a54bd83014b13117930a93 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 28 Oct 2026 01:13:59 +0000 Subject: [PATCH 1/2] air: send submit's config validation over the CLI's existing connection The submit-path config:validate request used its own clone of http.DefaultTransport, so every `air run` opened a second TCP and TLS connection to the workspace. Send it through an SDK client instead, so it reuses the connection the CLI's other calls already opened. The request is still built the same way and gets one attempt. The SDK reports a failed response-body read without the HTTP status, so a caller-error response whose body stalls now counts as validation being unavailable and the submission continues; the server validates the run again when it starts. Co-authored-by: Isaac --- cmd/air/validateconfig.go | 18 ++++++++---------- cmd/air/validateconfig_test.go | 32 +++++++++++++++++++++++++++----- 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/cmd/air/validateconfig.go b/cmd/air/validateconfig.go index 62dad63f8ab..9bb5aa1a2cc 100644 --- a/cmd/air/validateconfig.go +++ b/cmd/air/validateconfig.go @@ -2,7 +2,6 @@ package aircmd import ( "context" - "crypto/tls" "encoding/json" "errors" "fmt" @@ -147,15 +146,14 @@ func validateConfigOnce(ctx context.Context, w *databricks.WorkspaceClient, cfg req.Header.Set("User-Agent", useragent.FromContext(req.Context())) traceparent.AddTraceparent(req) - transport := clientCfg.Transport - if transport == nil { - defaultTransport := http.DefaultTransport.(*http.Transport).Clone() - if clientCfg.InsecureSkipVerify { - defaultTransport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} - } - transport = defaultTransport - } - resp, err := (&http.Client{Transport: transport}).Do(req) + // Send through an SDK client so the request reuses the connection the CLI's + // other calls already opened, instead of a new TCP and TLS handshake. It makes + // one attempt: the SDK retries some connection errors whatever ErrorRetriable + // says, so a minimal retry window stops it after the first. + clientCfg.ErrorRetriable = func(context.Context, error) bool { return false } + clientCfg.TransientErrors = nil + clientCfg.RetryTimeout = time.Nanosecond + resp, err := (&http.Client{Transport: httpclient.NewApiClient(clientCfg)}).Do(req) if err != nil { return fmt.Errorf("failed to validate config: %w", err) } diff --git a/cmd/air/validateconfig_test.go b/cmd/air/validateconfig_test.go index f0a0f3e38c8..e20f83285e4 100644 --- a/cmd/air/validateconfig_test.go +++ b/cmd/air/validateconfig_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "io" + "net" "net/http" "net/http/httptest" "net/url" @@ -19,6 +20,7 @@ import ( sdkconfig "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/credentials" sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" @@ -185,7 +187,9 @@ func TestPreflightValidationFailsOpenOnInvalidSuccessResponse(t *testing.T) { } } -func TestPreflightValidationBodyReadFailureOnCallerErrorBlocks(t *testing.T) { +// The SDK reports a failed body read without the response status, so even a +// caller-error response counts as validation being unavailable. +func TestPreflightValidationBodyReadFailureOnCallerErrorFailsOpen(t *testing.T) { w := validationTestWorkspaceClientWithTransport(t, func(req *http.Request) (*http.Response, error) { body := io.NopCloser(strings.NewReader(`{}`)) status := http.StatusOK @@ -204,10 +208,28 @@ func TestPreflightValidationBodyReadFailureOnCallerErrorBlocks(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") - require.Error(t, err) - apiErr, ok := errors.AsType[*apierr.APIError](err) - require.True(t, ok) - assert.Equal(t, http.StatusUnauthorized, apiErr.StatusCode) + require.NoError(t, err) +} + +func TestPreflightValidationReusesSDKConnection(t *testing.T) { + var connections atomic.Int32 + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{}`)) + })) + srv.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + connections.Add(1) + } + } + srv.Start() + t.Cleanup(srv.Close) + w := validationTestWorkspaceClient(t, srv.URL) + + _, err := w.CurrentUser.Me(t.Context(), iam.MeRequest{}) + require.NoError(t, err) + err = preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + require.NoError(t, err) + assert.Equal(t, int32(1), connections.Load()) } func TestPreflightValidationOAuthRefreshUsesDeadline(t *testing.T) { From 0b62bdf8172f9c0e5ef84951f3260966fc313daf Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 28 Oct 2026 01:14:45 +0000 Subject: [PATCH 2/2] Add changelog entry Co-authored-by: Isaac --- .nextchanges/air/7210.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 .nextchanges/air/7210.md diff --git a/.nextchanges/air/7210.md b/.nextchanges/air/7210.md new file mode 100644 index 00000000000..c71d364f35a --- /dev/null +++ b/.nextchanges/air/7210.md @@ -0,0 +1 @@ +* `air run` reuses its existing workspace connection for config validation instead of opening a second one, which saves a TCP and TLS handshake per submission. ([#7210](https://github.com/databricks/cli/pull/7210))