move the http client instance to the cloudfetch config

Author: Multimo

connector.go

@@ -260,6 +260,12 @@ func WithAuthenticator(authr auth.Authenticator) ConnOption {
 func WithTransport(t http.RoundTripper) ConnOption {
 	return func(c *config.Config) {
 		c.Transport = t
+
+		if c.CloudFetchConfig.HTTPClient == nil {
+			c.CloudFetchConfig.HTTPClient = &http.Client{
+				Transport: t,
+			}
+		}
 	}
 }
 

connector_test.go

@@ -48,6 +48,7 @@ func TestNewConnector(t *testing.T) {
 			MaxFilesInMemory:             10,
 			MinTimeToExpiry:              0 * time.Second,
 			CloudFetchSpeedThresholdMbps: 0.1,
+			HTTPClient:                   &http.Client{Transport: roundTripper},
 		}
 		expectedUserConfig := config.UserConfig{
 			Host:             host,
@@ -246,6 +247,25 @@ func TestNewConnector(t *testing.T) {
 		require.True(t, ok)
 		assert.False(t, coni.cfg.EnableMetricViewMetadata)
 	})
+
+	t.Run("Connector test WithTransport sets HTTPClient in CloudFetchConfig", func(t *testing.T) {
+		host := "databricks-host"
+		accessToken := "token"
+		httpPath := "http-path"
+		customTransport := &http.Transport{MaxIdleConns: 10}
+		con, err := NewConnector(
+			WithServerHostname(host),
+			WithAccessToken(accessToken),
+			WithHTTPPath(httpPath),
+			WithTransport(customTransport),
+		)
+		assert.Nil(t, err)
+
+		coni, ok := con.(*connector)
+		require.True(t, ok)
+		assert.NotNil(t, coni.cfg.CloudFetchConfig.HTTPClient)
+		assert.Equal(t, customTransport, coni.cfg.CloudFetchConfig.HTTPClient.Transport)
+	})
 }
 
 type mockRoundTripper struct{}

internal/config/config.go

@@ -479,7 +479,7 @@ type CloudFetchConfig struct {
 	MaxFilesInMemory             int
 	MinTimeToExpiry              time.Duration
 	CloudFetchSpeedThresholdMbps float64 // Minimum download speed in MBps before WARN logging (default: 0.1)
-	Transport                    http.RoundTripper
+	HTTPClient                   *http.Client
 }
 
 func (cfg CloudFetchConfig) WithDefaults() CloudFetchConfig {

internal/rows/arrowbased/batchloader.go

@@ -34,12 +34,9 @@ func NewCloudIPCStreamIterator(
 	startRowOffset int64,
 	cfg *config.Config,
 ) (IPCStreamIterator, dbsqlerr.DBError) {
-	transport := cfg.UserConfig.Transport
 	httpClient := http.DefaultClient
-	if transport != nil {
-		httpClient = &http.Client{
-			Transport: transport,
-		}
+	if cfg.UserConfig.CloudFetchConfig.HTTPClient != nil {
+		httpClient = cfg.UserConfig.CloudFetchConfig.HTTPClient
 	}
 
 	bi := &cloudIPCStreamIterator{

internal/rows/arrowbased/batchloader_test.go

@@ -254,19 +254,21 @@ func TestCloudFetchIterator(t *testing.T) {
 		assert.ErrorContains(t, err3, fmt.Sprintf("%s %d", "HTTP error", http.StatusNotFound))
 	})
 
-	t.Run("should use custom Transport when provided", func(t *testing.T) {
+	t.Run("should use custom HTTPClient when provided", func(t *testing.T) {
 		handler = func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 			w.Write(generateMockArrowBytes(generateArrowRecord()))
 		}
 
 		startRowOffset := int64(100)
-		customTransport := &http.Transport{MaxIdleConns: 10}
+		customHTTPClient := &http.Client{
+			Transport: &http.Transport{MaxIdleConns: 10},
+		}
 
 		cfg := config.WithDefaults()
 		cfg.UseLz4Compression = false
 		cfg.MaxDownloadThreads = 1
-		cfg.UserConfig.Transport = customTransport
+		cfg.UserConfig.CloudFetchConfig.HTTPClient = customHTTPClient
 
 		bi, err := NewCloudBatchIterator(
 			context.Background(),
@@ -282,16 +284,15 @@ func TestCloudFetchIterator(t *testing.T) {
 		assert.Nil(t, err)
 
 		cbi := bi.(*batchIterator).ipcIterator.(*cloudIPCStreamIterator)
-		assert.NotNil(t, cbi.httpClient)
-		assert.Equal(t, customTransport, cbi.httpClient.Transport)
+		assert.Equal(t, customHTTPClient, cbi.httpClient)
 
 		// Verify fetch works
 		sab, nextErr := bi.Next()
 		assert.Nil(t, nextErr)
 		assert.NotNil(t, sab)
 	})
 
-	t.Run("should fallback to http.DefaultClient when Transport is nil", func(t *testing.T) {
+	t.Run("should fallback to http.DefaultClient when HTTPClient is nil", func(t *testing.T) {
 		handler = func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 			w.Write(generateMockArrowBytes(generateArrowRecord()))
@@ -301,6 +302,8 @@ func TestCloudFetchIterator(t *testing.T) {
 		cfg := config.WithDefaults()
 		cfg.UseLz4Compression = false
 		cfg.MaxDownloadThreads = 1
+		// Explicitly set HTTPClient to nil to verify fallback behavior
+		cfg.UserConfig.CloudFetchConfig.HTTPClient = nil
 
 		bi, err := NewCloudBatchIterator(
 			context.Background(),