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)) 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) {