Swap implementation to use existing transport setting

Author: Multimo

connector.go

@@ -270,13 +270,6 @@ func WithCloudFetch(useCloudFetch bool) ConnOption {
 	}
 }
 
-// WithHTTPClient allows a custom http client to be used for cloud fetch. Default is http.DefaultClient.
-func WithHTTPClient(httpClient *http.Client) ConnOption {
-	return func(c *config.Config) {
-		c.UserConfig.CloudFetchConfig.HTTPClient = httpClient
-	}
-}
-
 // WithMaxDownloadThreads sets up maximum download threads for cloud fetch. Default is 10.
 func WithMaxDownloadThreads(numThreads int) ConnOption {
 	return func(c *config.Config) {

connector_test.go

@@ -246,42 +246,6 @@ func TestNewConnector(t *testing.T) {
 		require.True(t, ok)
 		assert.False(t, coni.cfg.EnableMetricViewMetadata)
 	})
-
-	t.Run("Connector test WithCloudFetchHTTPClient sets custom client", func(t *testing.T) {
-		host := "databricks-host"
-		accessToken := "token"
-		httpPath := "http-path"
-		customClient := &http.Client{Timeout: 5 * time.Second}
-
-		con, err := NewConnector(
-			WithServerHostname(host),
-			WithAccessToken(accessToken),
-			WithHTTPPath(httpPath),
-			WithHTTPClient(customClient),
-		)
-		assert.Nil(t, err)
-
-		coni, ok := con.(*connector)
-		require.True(t, ok)
-		assert.Equal(t, customClient, coni.cfg.UserConfig.CloudFetchConfig.HTTPClient)
-	})
-
-	t.Run("Connector test WithCloudFetchHTTPClient with nil client is accepted", func(t *testing.T) {
-		host := "databricks-host"
-		accessToken := "token"
-		httpPath := "http-path"
-
-		con, err := NewConnector(
-			WithServerHostname(host),
-			WithAccessToken(accessToken),
-			WithHTTPPath(httpPath),
-		)
-		assert.Nil(t, err)
-
-		coni, ok := con.(*connector)
-		require.True(t, ok)
-		assert.Nil(t, coni.cfg.UserConfig.CloudFetchConfig.HTTPClient)
-	})
 }
 
 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)
-	HTTPClient                   *http.Client
+	Transport                    http.RoundTripper
 }
 
 func (cfg CloudFetchConfig) WithDefaults() CloudFetchConfig {

internal/rows/arrowbased/batchloader.go

@@ -40,7 +40,7 @@ func NewCloudIPCStreamIterator(
 		startRowOffset: startRowOffset,
 		pendingLinks:   NewQueue[cli_service.TSparkArrowResultLink](),
 		downloadTasks:  NewQueue[cloudFetchDownloadTask](),
-		httpClient:     cfg.UserConfig.CloudFetchConfig.HTTPClient,
+		transport:      cfg.UserConfig.Transport,
 	}
 
 	for _, link := range files {
@@ -141,7 +141,7 @@ type cloudIPCStreamIterator struct {
 	startRowOffset int64
 	pendingLinks   Queue[cli_service.TSparkArrowResultLink]
 	downloadTasks  Queue[cloudFetchDownloadTask]
-	httpClient     *http.Client
+	transport      http.RoundTripper
 }
 
 var _ IPCStreamIterator = (*cloudIPCStreamIterator)(nil)
@@ -164,7 +164,7 @@ func (bi *cloudIPCStreamIterator) Next() (io.Reader, error) {
 			resultChan:         make(chan cloudFetchDownloadTaskResult),
 			minTimeToExpiry:    bi.cfg.MinTimeToExpiry,
 			speedThresholdMbps: bi.cfg.CloudFetchSpeedThresholdMbps,
-			httpClient:         bi.httpClient,
+			transport:          bi.transport,
 		}
 		task.Run()
 		bi.downloadTasks.Enqueue(task)
@@ -213,7 +213,7 @@ type cloudFetchDownloadTask struct {
 	link               *cli_service.TSparkArrowResultLink
 	resultChan         chan cloudFetchDownloadTaskResult
 	speedThresholdMbps float64
-	httpClient         *http.Client
+	transport          http.RoundTripper
 }
 
 func (cft *cloudFetchDownloadTask) GetResult() (io.Reader, error) {
@@ -256,7 +256,7 @@ func (cft *cloudFetchDownloadTask) Run() {
 			cft.link.StartRowOffset,
 			cft.link.RowCount,
 		)
-		data, err := fetchBatchBytes(cft.ctx, cft.link, cft.minTimeToExpiry, cft.speedThresholdMbps, cft.httpClient)
+		data, err := fetchBatchBytes(cft.ctx, cft.link, cft.minTimeToExpiry, cft.speedThresholdMbps, cft.transport)
 		if err != nil {
 			cft.resultChan <- cloudFetchDownloadTaskResult{data: nil, err: err}
 			return
@@ -304,7 +304,7 @@ func fetchBatchBytes(
 	link *cli_service.TSparkArrowResultLink,
 	minTimeToExpiry time.Duration,
 	speedThresholdMbps float64,
-	httpClient *http.Client,
+	transport http.RoundTripper,
 ) (io.ReadCloser, error) {
 	if isLinkExpired(link.ExpiryTime, minTimeToExpiry) {
 		return nil, errors.New(dbsqlerr.ErrLinkExpired)
@@ -322,8 +322,11 @@ func fetchBatchBytes(
 		}
 	}
 
-	if httpClient == nil {
-		httpClient = http.DefaultClient
+	httpClient := http.DefaultClient
+	if transport != nil {
+		httpClient = &http.Client{
+			Transport: transport,
+		}
 	}
 
 	startTime := time.Now()

internal/rows/arrowbased/batchloader_test.go

@@ -254,8 +254,11 @@ func TestCloudFetchIterator(t *testing.T) {
 		assert.ErrorContains(t, err3, fmt.Sprintf("%s %d", "HTTP error", http.StatusNotFound))
 	})
 
-	t.Run("should use custom HTTPClient when provided", func(t *testing.T) {
-		customClient := &http.Client{Timeout: 5 * time.Second}
+	t.Run("should use custom Transport when provided", func(t *testing.T) {
+		customTransport := &http.Transport{
+			MaxIdleConns:        10,
+			MaxIdleConnsPerHost: 5,
+		}
 		requestCount := 0
 
 		handler = func(w http.ResponseWriter, r *http.Request) {
@@ -281,7 +284,7 @@ func TestCloudFetchIterator(t *testing.T) {
 		cfg := config.WithDefaults()
 		cfg.UseLz4Compression = false
 		cfg.MaxDownloadThreads = 1
-		cfg.UserConfig.CloudFetchConfig.HTTPClient = customClient
+		cfg.UserConfig.Transport = customTransport
 
 		bi, err := NewCloudBatchIterator(
 			context.Background(),
@@ -291,21 +294,21 @@ func TestCloudFetchIterator(t *testing.T) {
 		)
 		assert.Nil(t, err)
 
-		// Verify custom client is passed through the iterator chain
+		// Verify custom transport is passed through the iterator chain
 		wrapper, ok := bi.(*batchIterator)
 		assert.True(t, ok)
 		cbi, ok := wrapper.ipcIterator.(*cloudIPCStreamIterator)
 		assert.True(t, ok)
-		assert.Equal(t, customClient, cbi.httpClient)
+		assert.Equal(t, customTransport, cbi.transport)
 
-		// Fetch should work with custom client
+		// Fetch should work with custom transport
 		sab1, nextErr := bi.Next()
 		assert.Nil(t, nextErr)
 		assert.NotNil(t, sab1)
 		assert.Greater(t, requestCount, 0) // Verify request was made
 	})
 
-	t.Run("should use http.DefaultClient when HTTPClient is nil", func(t *testing.T) {
+	t.Run("should use http.DefaultClient when Transport is nil", func(t *testing.T) {
 		handler = func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 			_, err := w.Write(generateMockArrowBytes(generateArrowRecord()))
@@ -328,7 +331,7 @@ func TestCloudFetchIterator(t *testing.T) {
 		cfg := config.WithDefaults()
 		cfg.UseLz4Compression = false
 		cfg.MaxDownloadThreads = 1
-		// HTTPClient is nil by default
+		// Transport is nil by default
 
 		bi, err := NewCloudBatchIterator(
 			context.Background(),
@@ -338,12 +341,12 @@ func TestCloudFetchIterator(t *testing.T) {
 		)
 		assert.Nil(t, err)
 
-		// Verify nil client is passed through
+		// Verify nil transport is passed through
 		wrapper, ok := bi.(*batchIterator)
 		assert.True(t, ok)
 		cbi, ok := wrapper.ipcIterator.(*cloudIPCStreamIterator)
 		assert.True(t, ok)
-		assert.Nil(t, cbi.httpClient)
+		assert.Nil(t, cbi.transport)
 
 		// Fetch should work (falls back to http.DefaultClient)
 		sab1, nextErr := bi.Next()