add log for cloud fetch speed

Author: shivam2680

internal/config/config.go

@@ -468,6 +468,7 @@ type CloudFetchConfig struct {
 	MaxDownloadThreads int
 	MaxFilesInMemory   int
 	MinTimeToExpiry    time.Duration
+	CloudFetchSpeedThresholdMbps float64 // Minimum download speed in MBps before WARN logging (default: 0.1)
 }
 
 func (cfg CloudFetchConfig) WithDefaults() CloudFetchConfig {
@@ -485,6 +486,10 @@ func (cfg CloudFetchConfig) WithDefaults() CloudFetchConfig {
 		cfg.MinTimeToExpiry = 0 * time.Second
 	}
 
+	if cfg.CloudFetchSpeedThresholdMbps <= 0 {
+		cfg.CloudFetchSpeedThresholdMbps = 0.1
+	}
+
 	return cfg
 }
 
@@ -494,5 +499,6 @@ func (cfg CloudFetchConfig) DeepCopy() CloudFetchConfig {
 		MaxDownloadThreads: cfg.MaxDownloadThreads,
 		MaxFilesInMemory:   cfg.MaxFilesInMemory,
 		MinTimeToExpiry:    cfg.MinTimeToExpiry,
+		CloudFetchSpeedThresholdMbps: cfg.CloudFetchSpeedThresholdMbps,
 	}
 }

internal/rows/arrowbased/batchloader.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"strings"
 	"time"
 
 	"github.com/databricks/databricks-sql-go/internal/config"
@@ -143,6 +144,7 @@ func (bi *cloudBatchIterator) Next() (SparkArrowBatch, error) {
 			link:              link,
 			resultChan:        make(chan cloudFetchDownloadTaskResult),
 			minTimeToExpiry:   bi.cfg.MinTimeToExpiry,
+			speedThresholdMbps: bi.cfg.CloudFetchSpeedThresholdMbps,
 		}
 		task.Run()
 		bi.downloadTasks.Enqueue(task)
@@ -190,6 +192,7 @@ type cloudFetchDownloadTask struct {
 	minTimeToExpiry   time.Duration
 	link              *cli_service.TSparkArrowResultLink
 	resultChan        chan cloudFetchDownloadTaskResult
+	speedThresholdMbps float64
 }
 
 func (cft *cloudFetchDownloadTask) GetResult() (SparkArrowBatch, error) {
@@ -232,7 +235,7 @@ func (cft *cloudFetchDownloadTask) Run() {
 			cft.link.StartRowOffset,
 			cft.link.RowCount,
 		)
-		data, err := fetchBatchBytes(cft.ctx, cft.link, cft.minTimeToExpiry)
+		data, err := fetchBatchBytes(cft.ctx, cft.link, cft.minTimeToExpiry, cft.speedThresholdMbps)
 		if err != nil {
 			cft.resultChan <- cloudFetchDownloadTaskResult{batch: nil, err: err}
 			return
@@ -262,10 +265,30 @@ func (cft *cloudFetchDownloadTask) Run() {
 	}()
 }
 
+// logCloudFetchSpeed calculates and logs download speed metrics
+func logCloudFetchSpeed(fullURL string, contentLength int64, duration time.Duration, speedThresholdMbps float64) {
+	if contentLength > 0 && duration.Seconds() > 0 {
+		// Extract base URL (up to first ?)
+		baseURL := fullURL
+		if idx := strings.Index(baseURL, "?"); idx != -1 {
+			baseURL = baseURL[:idx]
+		}
+
+		speedMbps := float64(contentLength) / (1024 * 1024) / duration.Seconds()
+
+		logger.Info().Msgf("CloudFetch: Result File Download speed from cloud storage %s %.4f Mbps", baseURL, speedMbps)
+
+		if speedMbps < speedThresholdMbps {
+			logger.Warn().Msgf("CloudFetch: Results download is slower than threshold speed of %.4f Mbps", speedThresholdMbps)
+		}
+	}
+}
+
 func fetchBatchBytes(
 	ctx context.Context,
 	link *cli_service.TSparkArrowResultLink,
 	minTimeToExpiry time.Duration,
+	speedThresholdMbps float64,
 ) (io.ReadCloser, error) {
 	if isLinkExpired(link.ExpiryTime, minTimeToExpiry) {
 		return nil, errors.New(dbsqlerr.ErrLinkExpired)
@@ -283,6 +306,7 @@ func fetchBatchBytes(
 		}
 	}
 
+	startTime := time.Now()
 	client := http.DefaultClient
 	res, err := client.Do(req)
 	if err != nil {
@@ -293,6 +317,9 @@ func fetchBatchBytes(
 		return nil, dbsqlerrint.NewDriverError(ctx, msg, err)
 	}
 
+	// Log download speed metrics
+	logCloudFetchSpeed(link.FileLink, res.ContentLength, time.Since(startTime), speedThresholdMbps)
+
 	return res.Body, nil
 }