air: send submit's config validation over the CLI's existing connection

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

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

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