fix: skip redundant fetches for complete direct results

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+## Unreleased
+
+- Avoid redundant Thrift result fetches when direct results specify `hasMoreRows=false` (databricks/databricks-sql-go#539).
+
 ## v1.19.0 (2026-11-04)
 - Upgrade the SQL kernel bindings to v1.4.0.
 - Upgrade go-oidc to v3.21.0 (go-jose v4.1.6) and x/oauth2 to v0.36.0 to address go-jose/v3 security findings (databricks/databricks-sql-go#526).

internal/rows/arrowbased/arrowRecordIterator_test.go

@@ -42,7 +42,7 @@ func TestArrowRecordIterator(t *testing.T) {
 			rowscanner.NewDelimiter(0, 7311),
 			5000,
 			nil,
-			false,
+			nil, false,
 			simpleClient,
 			logger,
 		)
@@ -140,7 +140,7 @@ func TestArrowRecordIterator(t *testing.T) {
 			rowscanner.NewDelimiter(0, 0),
 			5000,
 			nil,
-			false,
+			nil, false,
 			simpleClient,
 			logger,
 		)
@@ -214,7 +214,7 @@ func TestArrowRecordIteratorSurfacesPageErrorOnce(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger)
 	cfg := *config.WithDefaults()
 	bi, err := NewLocalBatchIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema, &cfg)
 	assert.Nil(t, err)
@@ -264,7 +264,7 @@ func TestArrowRecordIteratorSchemaReportsPageErrorOnce(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger)
 	cfg := *config.WithDefaults()
 	// nil batch iterator and nil schema bytes: Schema() must fetch a page.
 	it := NewArrowRecordIterator(ctx, rpi, nil, nil, cfg)
@@ -303,7 +303,7 @@ func TestArrowRecordIteratorNoFetchAfterEarlyClose(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, tc, logger)
 	cfg := *config.WithDefaults()
 	schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema
 	bi, err := NewLocalBatchIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg)
@@ -339,7 +339,7 @@ func TestArrowRecordIteratorCloseDiscardsPendingPageError(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger)
 	cfg := *config.WithDefaults()
 	it := NewArrowRecordIterator(ctx, rpi, nil, nil, cfg)
 
@@ -373,7 +373,7 @@ func TestArrowIPCStreamIteratorCloseDiscardsPendingPageError(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger)
 	cfg := *config.WithDefaults()
 	it := NewArrowIPCStreamIterator(ctx, rpi, nil, nil, cfg)
 
@@ -412,7 +412,7 @@ func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger)
 	cfg := *config.WithDefaults()
 	schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema
 	ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg)
@@ -469,7 +469,7 @@ func TestArrowIPCStreamIteratorNoFetchAfterEarlyClose(t *testing.T) {
 
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId")
-	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger)
+	rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, tc, logger)
 	cfg := *config.WithDefaults()
 	schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema
 	ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg)
@@ -509,7 +509,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) {
 			rowscanner.NewDelimiter(0, 0),
 			5000,
 			nil,
-			false,
+			nil, false,
 			simpleClient,
 			logger,
 		)
@@ -563,7 +563,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) {
 			rowscanner.NewDelimiter(0, 0),
 			5000,
 			nil,
-			false,
+			nil, false,
 			simpleClient,
 			logger,
 		)
@@ -605,7 +605,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) {
 			rowscanner.NewDelimiter(0, 0),
 			5000,
 			nil,
-			false,
+			nil, false,
 			failingClient,
 			logger,
 		)

internal/rows/arrowbased/arrowRows_test.go

@@ -1545,7 +1545,7 @@ func TestArrowRowScanner(t *testing.T) {
 			rowscanner.NewDelimiter(0, 7311),
 			5000,
 			nil,
-			false,
+			nil, false,
 			client,
 			logger)
 

internal/rows/arrowbased/batchloader_test.go

@@ -1004,7 +1004,11 @@ func TestCloudFetchIterator_CloseReleasesInFlightDownloads(t *testing.T) {
 
 	// Kick off the first batch download. The iterator schedules
 	// MaxDownloadThreads concurrent fetches behind the scenes.
-	go func() { _, _ = bi.Next() }()
+	nextDone := make(chan struct{})
+	go func() {
+		defer close(nextDone)
+		_, _ = bi.Next()
+	}()
 
 	// Wait for all MaxDownloadThreads goroutines to be blocked inside the
 	// server handler (they've issued the GET and are waiting for the body).
@@ -1018,9 +1022,12 @@ func TestCloudFetchIterator_CloseReleasesInFlightDownloads(t *testing.T) {
 	// will be queued, blocked on the send.
 	close(release)
 
-	// Give the goroutines time to finish their HTTP work and reach the
-	// channel send.
-	time.Sleep(200 * time.Millisecond)
+	// Serialize iterator access while the other downloads remain in flight.
+	select {
+	case <-nextDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("foreground Next did not finish")
+	}
 
 	// Close the iterator without consuming the remaining batches.
 	bi.Close()
@@ -1114,7 +1121,11 @@ func TestCloudFetchIterator_CloseReleasesAfterRetry(t *testing.T) {
 
 	// One foreground reader. It drains exactly one batch and then exits;
 	// the remaining MaxDownloadThreads-1 in-flight tasks have nobody reading.
-	go func() { _, _ = bi.Next() }()
+	nextDone := make(chan struct{})
+	go func() {
+		defer close(nextDone)
+		_, _ = bi.Next()
+	}()
 
 	// Wait for every concurrent task to have completed its 503 retry, slept
 	// the backoff, and reached its second-attempt handler (where it's
@@ -1131,8 +1142,12 @@ func TestCloudFetchIterator_CloseReleasesAfterRetry(t *testing.T) {
 	// goroutines park on resultChan (via sendResult).
 	close(release)
 
-	// Give the goroutines time to reach the channel send.
-	time.Sleep(200 * time.Millisecond)
+	// Serialize iterator access while the other downloads remain in flight.
+	select {
+	case <-nextDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("foreground Next did not finish")
+	}
 
 	// Close without draining. cft.ctx is cancelled for every queued task,
 	// so each parked sendResult must unblock via its ctx.Done arm.

internal/rows/rows.go

@@ -225,14 +225,18 @@ func NewRows(
 		d = rowscanner.NewDelimiter(0, 0)
 	}
 
-	// If the entire query result set fits in direct results the server closes
-	// the operations.
+	// CloudFetch may deliver all links without closing the server operation.
 	closedOnServer := directResults != nil && directResults.CloseOperation != nil
+	var hasMoreRows *bool
+	if directResults != nil && directResults.ResultSet != nil {
+		hasMoreRows = directResults.ResultSet.HasMoreRows
+	}
 	r.ResultPageIterator = rowscanner.NewResultPageIterator(
 		resultsCtx,
 		d,
 		pageSize,
 		opHandle,
+		hasMoreRows,
 		closedOnServer,
 		client,
 		r.logger(),

internal/rows/rows_test.go

@@ -230,6 +230,7 @@ func TestRowsFetchResultPageNoDirectResults(t *testing.T) {
 		rowscanner.NewDelimiter(0, 0),
 		1000,
 		nil,
+		nil,
 		false,
 		client,
 		rowSet.logger(),
@@ -326,6 +327,7 @@ func TestRowsFetchResultPageWithDirectResults(t *testing.T) {
 		rowscanner.NewDelimiter(rowSet.RowScanner.Start(), rowSet.RowScanner.Count()),
 		1000,
 		nil,
+		nil,
 		false,
 		client,
 		rowSet.logger(),
@@ -480,6 +482,7 @@ func TestNextNoDirectResults(t *testing.T) {
 		rowscanner.NewDelimiter(0, 0),
 		1000,
 		nil,
+		nil,
 		false,
 		client,
 		rowSet.logger(),
@@ -756,7 +759,7 @@ func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) {
 				rowSet.resultSetMetadata = metadata
 			}
 			rowSet.ResultPageIterator = rowscanner.NewResultPageIterator(
-				rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, true, testClient, rowSet.logger())
+				rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, nil, true, testClient, rowSet.logger())
 
 			checkTypes := func() {
 				t.Helper()
@@ -925,6 +928,7 @@ func TestFetchResultsWithRetries(t *testing.T) {
 		rowscanner.NewDelimiter(0, 0),
 		1000,
 		nil,
+		nil,
 		false,
 		client,
 		rowSet.logger(),

internal/rows/rowscanner/category_test.go

@@ -54,7 +54,7 @@ func TestResultPageIteratorErrNilOnNormalFinish(t *testing.T) {
 		},
 	}
 
-	rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, false, c,
+	rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, nil, false, c,
 		dbsqllog.WithContext("connId", "correlationId", ""))
 
 	assert.True(t, rpf.HasNext())
@@ -81,7 +81,7 @@ func TestResultPageIteratorPreservesFetchError(t *testing.T) {
 		},
 	}
 
-	rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, false, c,
+	rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, nil, false, c,
 		dbsqllog.WithContext("connId", "correlationId", ""))
 
 	assert.False(t, rpf.HasNext())

internal/rows/rowscanner/resultPageIterator.go

@@ -53,17 +53,16 @@ func NewResultPageIterator(
 	delimiter Delimiter,
 	maxPageSize int64,
 	opHandle *cli_service.TOperationHandle,
+	hasMoreRows *bool,
 	closedOnServer bool,
 	client cli_service.TCLIService,
 	logger *dbsqllog.DBSQLLogger,
 ) ResultPageIterator {
 
-	// delimiter and hasMoreRows are used to set up the point in the paginated
-	// result set that this iterator starts from.
 	return &resultPageIterator{
 		ctx:            ctx,
 		Delimiter:      delimiter,
-		isFinished:     closedOnServer,
+		isFinished:     closedOnServer || (hasMoreRows != nil && !*hasMoreRows),
 		maxPageSize:    maxPageSize,
 		opHandle:       opHandle,
 		closedOnServer: closedOnServer,
@@ -130,6 +129,8 @@ func (rpf *resultPageIterator) Err() error {
 // Returns true if there are more pages in the result set.
 func (rpf *resultPageIterator) HasNext() bool {
 	if rpf.isFinished && rpf.nextResultPage == nil {
+		// Preserve best-effort cleanup when direct results skip the final fetch.
+		rpf.Close() //nolint:errcheck,gosec // G104: match cleanup of the final fetched page
 		// There are no more pages to load and there isn't an already fetched
 		// page waiting to retrieved by Next(). Preserve a previously stored
 		// fetch error rather than masking it as a normal end-of-results.