Add correctness tests for telemetry gap fixes

Four targeted test suites covering the bug fixes from the previous commit:

aggregator_test.go (new):
- WaitsForInFlightWorkerExports: close() blocks until all HTTP exports complete
- DrainsPendingQueueJobsBeforeCancel: drain step exports pending queue jobs
- InFlightAddBeforeSend: inFlight.Add happens before queue send (no missed jobs)
- SafeToCallMultipleTimes: concurrent close() calls don't deadlock (sync.Once)
- DropWhenQueueFull: drop path calls inFlight.Done() so Wait() is never stuck

integration_test.go:
- OperationLatencyMs_ZeroNotOmitted: 0ms latency serialises as 0 not null (omitempty fix)
- ChunkTotalPresent_DerivedFromChunkCount: chunk_total_present tag propagates to payload

batchloader_test.go:
- OnFileDownloaded callback is called once per file with positive downloadMs
- Nil callback does not panic (non-telemetry paths)

rows_test.go:
- CloseCallback receives correct chunkCount after multi-page iteration
- Nil closeCallback does not panic

Co-authored-by: Isaac

Author: samikshya-db

internal/rows/arrowbased/batchloader_test.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -339,6 +340,115 @@ func TestCloudFetchIterator(t *testing.T) {
 	})
 }
 
+// TestCloudFetchIterator_OnFileDownloaded_CallbackInvokedWithPositiveDuration verifies
+// that the onFileDownloaded telemetry callback is called once per downloaded S3 file with
+// a positive downloadMs value.
+//
+// This covers the CloudFetch timing fix where per-S3-file download durations are measured
+// and reported as initial_chunk_latency_ms / slowest_chunk_latency_ms / sum_chunks_download_time_ms
+// in the telemetry payload.
+func TestCloudFetchIterator_OnFileDownloaded_CallbackInvokedWithPositiveDuration(t *testing.T) {
+	// Serve real arrow bytes so the iterator can parse them successfully.
+	arrowBytes := generateMockArrowBytes(generateArrowRecord())
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Add a tiny sleep so the measured download time is reliably > 0ms.
+		time.Sleep(2 * time.Millisecond)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(arrowBytes)
+	}))
+	defer server.Close()
+
+	startRowOffset := int64(0)
+	links := []*cli_service.TSparkArrowResultLink{
+		{
+			FileLink:       server.URL,
+			ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
+			StartRowOffset: startRowOffset,
+			RowCount:       1,
+		},
+		{
+			FileLink:       server.URL,
+			ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
+			StartRowOffset: startRowOffset + 1,
+			RowCount:       1,
+		},
+	}
+
+	cfg := config.WithDefaults()
+	cfg.UseLz4Compression = false
+	cfg.MaxDownloadThreads = 1
+
+	var callbackMu sync.Mutex
+	var downloadDurations []int64
+
+	onFileDownloaded := func(downloadMs int64) {
+		callbackMu.Lock()
+		downloadDurations = append(downloadDurations, downloadMs)
+		callbackMu.Unlock()
+	}
+
+	bi, err := NewCloudBatchIterator(context.Background(), links, startRowOffset, cfg, onFileDownloaded)
+	assert.Nil(t, err)
+
+	// Consume all batches to trigger the downloads.
+	for bi.HasNext() {
+		_, nextErr := bi.Next()
+		assert.Nil(t, nextErr)
+	}
+
+	callbackMu.Lock()
+	durations := make([]int64, len(downloadDurations))
+	copy(durations, downloadDurations)
+	callbackMu.Unlock()
+
+	// Callback must be invoked once per link.
+	assert.Equal(t, len(links), len(durations),
+		"onFileDownloaded must be called once per downloaded file")
+
+	// Each reported duration must be positive (the server adds a 2ms delay).
+	for i, d := range durations {
+		assert.Greater(t, d, int64(0),
+			"onFileDownloaded must report positive downloadMs for file %d, got %d", i, d)
+	}
+}
+
+// TestCloudFetchIterator_OnFileDownloaded_NilCallbackDoesNotPanic verifies that passing
+// nil for onFileDownloaded (non-telemetry paths) does not cause a panic during iteration.
+func TestCloudFetchIterator_OnFileDownloaded_NilCallbackDoesNotPanic(t *testing.T) {
+	arrowBytes := generateMockArrowBytes(generateArrowRecord())
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(arrowBytes)
+	}))
+	defer server.Close()
+
+	startRowOffset := int64(0)
+	links := []*cli_service.TSparkArrowResultLink{
+		{
+			FileLink:       server.URL,
+			ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
+			StartRowOffset: startRowOffset,
+			RowCount:       1,
+		},
+	}
+
+	cfg := config.WithDefaults()
+	cfg.UseLz4Compression = false
+	cfg.MaxDownloadThreads = 1
+
+	// nil callback — must not panic
+	bi, err := NewCloudBatchIterator(context.Background(), links, startRowOffset, cfg, nil)
+	assert.Nil(t, err)
+
+	assert.NotPanics(t, func() {
+		for bi.HasNext() {
+			_, _ = bi.Next()
+		}
+	}, "nil onFileDownloaded must not cause a panic")
+}
+
 func generateArrowRecord() arrow.Record {
 	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
 

internal/rows/rows_test.go

@@ -1563,3 +1563,121 @@ func TestFetchResultPage_PropagatesGetNextPageError(t *testing.T) {
 
 	assert.ErrorContains(t, actualErr, errorMsg)
 }
+
+// TestRows_CloseCallback_ReceivesChunkCount verifies that when rows.Close() is called,
+// the closeCallback receives the correct chunkCount reflecting the number of result pages
+// that were fetched during iteration.
+//
+// This covers the fix where total_chunks_present in the telemetry payload was always null
+// for paginated CloudFetch queries: the driver now derives it from r.chunkCount and passes
+// it through closeCallback so connection.go can set the "chunk_total_present" tag.
+func TestRows_CloseCallback_ReceivesChunkCount(t *testing.T) {
+	t.Parallel()
+
+	noMoreRows := false
+	moreRows := true
+
+	// Two pages: page 0 (5 rows, has more), page 1 (3 rows, no more).
+	colVals := []*cli_service.TColumn{
+		{BoolVal: &cli_service.TBoolColumn{Values: []bool{true, false, true, false, true}}},
+	}
+	colVals2 := []*cli_service.TColumn{
+		{BoolVal: &cli_service.TBoolColumn{Values: []bool{true, false, true}}},
+	}
+
+	pages := []cli_service.TFetchResultsResp{
+		{
+			Status:      &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS},
+			HasMoreRows: &moreRows,
+			Results:     &cli_service.TRowSet{StartRowOffset: 0, Columns: colVals},
+		},
+		{
+			Status:      &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS},
+			HasMoreRows: &noMoreRows,
+			Results:     &cli_service.TRowSet{StartRowOffset: 5, Columns: colVals2},
+		},
+	}
+
+	pageIndex := -1
+	fetchFn := func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) {
+		pageIndex++
+		p := pages[pageIndex]
+		return &p, nil
+	}
+	metaFn := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (*cli_service.TGetResultSetMetadataResp, error) {
+		return &cli_service.TGetResultSetMetadataResp{
+			Status: &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS},
+			Schema: &cli_service.TTableSchema{
+				Columns: []*cli_service.TColumnDesc{
+					{ColumnName: "flag", Position: 0, TypeDesc: &cli_service.TTypeDesc{
+						Types: []*cli_service.TTypeEntry{{
+							PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE},
+						}},
+					}},
+				},
+			},
+		}, nil
+	}
+
+	testClient := &client.TestClient{
+		FnFetchResults:         fetchFn,
+		FnGetResultSetMetadata: metaFn,
+	}
+
+	var callbackChunkCount int
+	closeCallback := func(latencyMs int64, chunkCount int, err error) {
+		callbackChunkCount = chunkCount
+	}
+
+	ctx := driverctx.NewContextWithConnId(context.Background(), "connId")
+	cfg := config.WithDefaults()
+	cfg.MaxRows = 5 // force paging
+
+	dr, dbErr := NewRows(ctx, nil, testClient, cfg, nil, nil, closeCallback, nil)
+	assert.Nil(t, dbErr)
+
+	// Drain all rows to force two FetchResults calls.
+	dest := make([]driver.Value, 1)
+	for dr.Next(dest) == nil {
+	}
+
+	// Close should invoke the callback with the total chunk count (2 pages fetched).
+	assert.Nil(t, dr.Close())
+
+	// direct results count as chunk 0; two FetchResults calls give chunkCount=2.
+	// (No directResults here so chunkCount starts at 0, then +1 per FetchResults call.)
+	assert.Equal(t, 2, callbackChunkCount,
+		"closeCallback must receive the total number of result pages fetched")
+}
+
+// TestRows_CloseCallback_NilDoesNotPanic verifies that passing nil for closeCallback
+// does not cause a panic when rows.Close() is called.
+func TestRows_CloseCallback_NilDoesNotPanic(t *testing.T) {
+	t.Parallel()
+
+	noMoreRows := false
+	pages := []cli_service.TFetchResultsResp{
+		{
+			Status:      &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS},
+			HasMoreRows: &noMoreRows,
+			Results:     &cli_service.TRowSet{StartRowOffset: 0, Columns: []*cli_service.TColumn{}},
+		},
+	}
+	pageIndex := -1
+	fetchFn := func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) {
+		pageIndex++
+		p := pages[pageIndex]
+		return &p, nil
+	}
+	testClient := &client.TestClient{FnFetchResults: fetchFn}
+
+	ctx := driverctx.NewContextWithConnId(context.Background(), "connId")
+	cfg := config.WithDefaults()
+
+	dr, dbErr := NewRows(ctx, nil, testClient, cfg, nil, nil, nil /* closeCallback=nil */, nil)
+	assert.Nil(t, dbErr)
+
+	assert.NotPanics(t, func() {
+		_ = dr.Close()
+	}, "nil closeCallback must not cause a panic on rows.Close()")
+}