add log for cloud fetch speed (#281)

## Description Add logs for cloud fetch download speed relevant ticket: https://databricks.atlassian.net/browse/XTA-11037 ## Testing verified logs are printed when log level is set to INFO

Author: shivam2680

connector_test.go

@@ -43,10 +43,11 @@ func TestNewConnector(t *testing.T) {
 WithSkipTLSHostVerify(),
 )
 expectedCloudFetchConfig := config.CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 15,
-MaxFilesInMemory: 10,
-MinTimeToExpiry: 0 * time.Second,
+UseCloudFetch: true,
+MaxDownloadThreads: 15,
+MaxFilesInMemory: 10,
+MinTimeToExpiry: 0 * time.Second,
+CloudFetchSpeedThresholdMbps: 0.1,
 }
 expectedUserConfig := config.UserConfig{
 Host: host,
@@ -89,10 +90,11 @@ func TestNewConnector(t *testing.T) {
 WithHTTPPath(httpPath),
 )
 expectedCloudFetchConfig := config.CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 10,
-MaxFilesInMemory: 10,
-MinTimeToExpiry: 0 * time.Second,
+UseCloudFetch: true,
+MaxDownloadThreads: 10,
+MaxFilesInMemory: 10,
+MinTimeToExpiry: 0 * time.Second,
+CloudFetchSpeedThresholdMbps: 0.1,
 }
 expectedUserConfig := config.UserConfig{
 Host: host,
@@ -130,10 +132,11 @@ func TestNewConnector(t *testing.T) {
 WithRetries(-1, 0, 0),
 )
 expectedCloudFetchConfig := config.CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 10,
-MaxFilesInMemory: 10,
-MinTimeToExpiry: 0 * time.Second,
+UseCloudFetch: true,
+MaxDownloadThreads: 10,
+MaxFilesInMemory: 10,
+MinTimeToExpiry: 0 * time.Second,
+CloudFetchSpeedThresholdMbps: 0.1,
 }
 expectedUserConfig := config.UserConfig{
 Host: host,

internal/config/config.go

@@ -464,10 +464,11 @@ func (arrowConfig ArrowConfig) DeepCopy() ArrowConfig {
 }
 
 type CloudFetchConfig struct {
-UseCloudFetch bool
-MaxDownloadThreads int
-MaxFilesInMemory int
-MinTimeToExpiry time.Duration
+UseCloudFetch bool
+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,14 +486,19 @@ func (cfg CloudFetchConfig) WithDefaults() CloudFetchConfig {
 cfg.MinTimeToExpiry = 0 * time.Second
 }
 
+if cfg.CloudFetchSpeedThresholdMbps <= 0 {
+cfg.CloudFetchSpeedThresholdMbps = 0.1
+}
+
 return cfg
 }
 
 func (cfg CloudFetchConfig) DeepCopy() CloudFetchConfig {
 return CloudFetchConfig{
-UseCloudFetch: cfg.UseCloudFetch,
-MaxDownloadThreads: cfg.MaxDownloadThreads,
-MaxFilesInMemory: cfg.MaxFilesInMemory,
-MinTimeToExpiry: cfg.MinTimeToExpiry,
+UseCloudFetch: cfg.UseCloudFetch,
+MaxDownloadThreads: cfg.MaxDownloadThreads,
+MaxFilesInMemory: cfg.MaxFilesInMemory,
+MinTimeToExpiry: cfg.MinTimeToExpiry,
+CloudFetchSpeedThresholdMbps: cfg.CloudFetchSpeedThresholdMbps,
 }
 }

internal/config/config_test.go

@@ -247,9 +247,10 @@ func TestParseConfig(t *testing.T) {
 RetryWaitMin: 1 * time.Second,
 RetryWaitMax: 30 * time.Second,
 CloudFetchConfig: CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 10,
-MaxFilesInMemory: 10,
+UseCloudFetch: true,
+MaxDownloadThreads: 10,
+MaxFilesInMemory: 10,
+CloudFetchSpeedThresholdMbps: 0.1,
 },
 },
 wantURL: "https://example.cloud.databricks.com:8000/sql/1.0/endpoints/12346a5b5b0e123b",
@@ -271,9 +272,10 @@ func TestParseConfig(t *testing.T) {
 RetryWaitMin: 1 * time.Second,
 RetryWaitMax: 30 * time.Second,
 CloudFetchConfig: CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 15,
-MaxFilesInMemory: 10,
+UseCloudFetch: true,
+MaxDownloadThreads: 15,
+MaxFilesInMemory: 10,
+CloudFetchSpeedThresholdMbps: 0.1,
 },
 },
 wantURL: "https://example.cloud.databricks.com:8000/sql/1.0/endpoints/12346a5b5b0e123b",
@@ -299,9 +301,10 @@ func TestParseConfig(t *testing.T) {
 RetryWaitMin: 1 * time.Second,
 RetryWaitMax: 30 * time.Second,
 CloudFetchConfig: CloudFetchConfig{
-UseCloudFetch: true,
-MaxDownloadThreads: 15,
-MaxFilesInMemory: 10,
+UseCloudFetch: true,
+MaxDownloadThreads: 15,
+MaxFilesInMemory: 10,
+CloudFetchSpeedThresholdMbps: 0.1,
 },
 },
 wantURL: "https://example.cloud.databricks.com:8000/sql/1.0/endpoints/12346a5b5b0e123a",

internal/rows/arrowbased/batchloader.go

@@ -5,6 +5,7 @@ import (
 "context"
 "fmt"
 "io"
+"strings"
 "time"
 
 "github.com/databricks/databricks-sql-go/internal/config"
@@ -154,12 +155,13 @@ func (bi *cloudIPCStreamIterator) Next() (io.Reader, error) {
 
 cancelCtx, cancelFn := context.WithCancel(bi.ctx)
 task := &cloudFetchDownloadTask{
-ctx: cancelCtx,
-cancel: cancelFn,
-useLz4Compression: bi.cfg.UseLz4Compression,
-link: link,
-resultChan: make(chan cloudFetchDownloadTaskResult),
-minTimeToExpiry: bi.cfg.MinTimeToExpiry,
+ctx: cancelCtx,
+cancel: cancelFn,
+useLz4Compression: bi.cfg.UseLz4Compression,
+link: link,
+resultChan: make(chan cloudFetchDownloadTaskResult),
+minTimeToExpiry: bi.cfg.MinTimeToExpiry,
+speedThresholdMbps: bi.cfg.CloudFetchSpeedThresholdMbps,
 }
 task.Run()
 bi.downloadTasks.Enqueue(task)
@@ -201,12 +203,13 @@ type cloudFetchDownloadTaskResult struct {
 }
 
 type cloudFetchDownloadTask struct {
-ctx context.Context
-cancel context.CancelFunc
-useLz4Compression bool
-minTimeToExpiry time.Duration
-link *cli_service.TSparkArrowResultLink
-resultChan chan cloudFetchDownloadTaskResult
+ctx context.Context
+cancel context.CancelFunc
+useLz4Compression bool
+minTimeToExpiry time.Duration
+link *cli_service.TSparkArrowResultLink
+resultChan chan cloudFetchDownloadTaskResult
+speedThresholdMbps float64
 }
 
 func (cft *cloudFetchDownloadTask) GetResult() (io.Reader, error) {
@@ -249,7 +252,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{data: nil, err: err}
 return
@@ -273,10 +276,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)
@@ -294,6 +317,7 @@ func fetchBatchBytes(
 }
 }
 
+startTime := time.Now()
 client := http.DefaultClient
 res, err := client.Do(req)
 if err != nil {
@@ -304,6 +328,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
 }