From a41aabda990b7d3a7b624890b5519baea3ea3dd1 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 07:15:59 +0000 Subject: [PATCH 1/7] fix: honor CloudFetch end-of-results before closing operations Signed-off-by: Vu Anh Phung --- .../arrowbased/arrowRecordIterator_test.go | 24 +- internal/rows/arrowbased/arrowRows_test.go | 2 +- internal/rows/arrowbased/batchloader_test.go | 29 ++- internal/rows/rows.go | 14 +- internal/rows/rows_test.go | 216 +++++++++++++++++- internal/rows/rowscanner/category_test.go | 4 +- .../rows/rowscanner/resultPageIterator.go | 9 +- 7 files changed, 262 insertions(+), 36 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 926dcc19..385c26af 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -42,7 +42,7 @@ func TestArrowRecordIterator(t *testing.T) { rowscanner.NewDelimiter(0, 7311), 5000, nil, - false, + false, false, simpleClient, logger, ) @@ -140,7 +140,7 @@ func TestArrowRecordIterator(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + false, 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, false, 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, false, 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, false, 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, false, 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, false, 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, false, 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, false, 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, + false, false, simpleClient, logger, ) @@ -563,7 +563,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + false, false, simpleClient, logger, ) @@ -605,7 +605,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + false, false, failingClient, logger, ) diff --git a/internal/rows/arrowbased/arrowRows_test.go b/internal/rows/arrowbased/arrowRows_test.go index 77caff7c..8903d782 100644 --- a/internal/rows/arrowbased/arrowRows_test.go +++ b/internal/rows/arrowbased/arrowRows_test.go @@ -1545,7 +1545,7 @@ func TestArrowRowScanner(t *testing.T) { rowscanner.NewDelimiter(0, 7311), 5000, nil, - false, + false, false, client, logger) diff --git a/internal/rows/arrowbased/batchloader_test.go b/internal/rows/arrowbased/batchloader_test.go index 3e3e4831..77ba1629 100644 --- a/internal/rows/arrowbased/batchloader_test.go +++ b/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. diff --git a/internal/rows/rows.go b/internal/rows/rows.go index f7a7a0f4..82e7c8b4 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -76,6 +76,8 @@ type rows struct { // resultsCancel aborts in-flight result RPCs/downloads when Close() is // called, so the detached ctx never leaves an operation uncancellable. resultsCancel context.CancelFunc + closeOnce sync.Once + closeErr error // Telemetry tracking // telemetryUpdate is called after each chunk is fetched with: @@ -225,14 +227,16 @@ 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 + isFinished := closedOnServer || (directResults != nil && directResults.ResultSet != nil && + directResults.ResultSet.HasMoreRows != nil && !*directResults.ResultSet.HasMoreRows) r.ResultPageIterator = rowscanner.NewResultPageIterator( resultsCtx, d, pageSize, opHandle, + isFinished, closedOnServer, client, r.logger(), @@ -271,7 +275,13 @@ func (r *rows) Close() error { if r == nil { return nil } + r.closeOnce.Do(func() { + r.closeErr = r.close() + }) + return r.closeErr +} +func (r *rows) close() error { // Release the detached results context after the close RPC runs, aborting // any in-flight FetchResults/CloudFetch downloads still referencing it. if r.resultsCancel != nil { diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 07d8b8ff..15b55821 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -9,13 +9,17 @@ import ( "fmt" "io" "math" + "net/http" + "net/http/httptest" "os" "reflect" "testing" "time" "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/ipc" + "github.com/apache/arrow/go/v12/arrow/memory" "github.com/databricks/databricks-sql-go/driverctx" "github.com/databricks/databricks-sql-go/internal/cli_service" "github.com/databricks/databricks-sql-go/internal/client" @@ -230,7 +234,7 @@ func TestRowsFetchResultPageNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, + false, false, client, rowSet.logger(), ) @@ -326,7 +330,7 @@ func TestRowsFetchResultPageWithDirectResults(t *testing.T) { rowscanner.NewDelimiter(rowSet.RowScanner.Start(), rowSet.RowScanner.Count()), 1000, nil, - false, + false, false, client, rowSet.logger(), ) @@ -480,7 +484,7 @@ func TestNextNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, + false, false, client, rowSet.logger(), ) @@ -756,7 +760,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, true, true, testClient, rowSet.logger()) checkTypes := func() { t.Helper() @@ -902,6 +906,147 @@ func TestRowsCloseOptimization(t *testing.T) { assert.Equal(t, 0, closeCount) } +func TestRowsDirectResultsLifecycle(t *testing.T) { + moreRows, noMoreRows := true, false + for _, tc := range []struct { + name string + cloudFetch bool + hasMoreRows *bool + closedOnServer bool + empty bool + wantFetches int + }{ + {name: "cloudfetch complete", cloudFetch: true, hasMoreRows: &noMoreRows}, + {name: "cloudfetch paginated", cloudFetch: true, hasMoreRows: &moreRows, wantFetches: 1}, + {name: "cloudfetch unknown", cloudFetch: true, wantFetches: 1}, + {name: "cloudfetch empty", cloudFetch: true, hasMoreRows: &noMoreRows, empty: true}, + {name: "inline closed", closedOnServer: true}, + {name: "inline complete", hasMoreRows: &noMoreRows}, + {name: "inline unknown", wantFetches: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + cfg := config.WithDefaults() + cfg.UseLz4Compression = false + metadata := &cli_service.TGetResultSetMetadataResp{ + Schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ + ColumnName: "flag", + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE}, + }}}, + }}}, + } + var fileURL string + if tc.cloudFetch { + schema := arrow.NewSchema([]arrow.Field{{Name: "flag", Type: arrow.FixedWidthTypes.Boolean}}, nil) + builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer builder.Release() + builder.Field(0).(*array.BooleanBuilder).AppendValues([]bool{true, false}, nil) + record := builder.NewRecord() + defer record.Release() + var data bytes.Buffer + writer := ipc.NewWriter(&data, ipc.WithSchema(schema)) + require.NoError(t, writer.Write(record)) + require.NoError(t, writer.Close()) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write(data.Bytes()) + })) + defer server.Close() + fileURL = server.URL + } + page := func(start int64, empty bool) *cli_service.TRowSet { + rs := &cli_service.TRowSet{StartRowOffset: start} + if tc.cloudFetch { + rs.ResultLinks = []*cli_service.TSparkArrowResultLink{} + if !empty { + rs.ResultLinks = append(rs.ResultLinks, &cli_service.TSparkArrowResultLink{ + FileLink: fileURL, StartRowOffset: start, RowCount: 2, + ExpiryTime: time.Now().Add(time.Hour).Unix(), + }) + } + } else { + rs.Columns = []*cli_service.TColumn{{BoolVal: &cli_service.TBoolColumn{Values: []bool{true, false}}}} + } + return rs + } + opHandle := &cli_service.TOperationHandle{OperationId: &cli_service.THandleIdentifier{GUID: make([]byte, 16)}} + fetches, closes, callbacks, downloads := 0, 0, 0, 0 + testClient := &client.TestClient{ + FnFetchResults: func(_ context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + fetches++ + require.LessOrEqual(t, fetches, tc.wantFetches) + assert.Same(t, opHandle, req.OperationHandle) + return &cli_service.TFetchResultsResp{HasMoreRows: &noMoreRows, Results: page(2, false)}, nil + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closes++ + assert.NoError(t, ctx.Err()) + assert.Same(t, opHandle, req.OperationHandle) + return &cli_service.TCloseOperationResp{}, nil + }, + } + direct := &cli_service.TSparkDirectResults{ + ResultSetMetadata: metadata, + ResultSet: &cli_service.TFetchResultsResp{ + HasMoreRows: tc.hasMoreRows, Results: page(0, tc.empty), + }, + } + if tc.closedOnServer { + direct.CloseOperation = &cli_service.TCloseOperationResp{} + } + var chunkCounts []int + dr, err := NewRows(context.Background(), opHandle, testClient, cfg, direct, &TelemetryCallbacks{ + OnChunkFetched: func(count int, _ int64, index int, latency int64, total int32) { + chunkCounts = append(chunkCounts, count) + assert.Equal(t, count-1, index) + if tc.cloudFetch { + assert.Zero(t, latency) + if !tc.empty { + assert.EqualValues(t, 1, total) + } + } + }, + OnCloudFetchFile: func(int64) { downloads++ }, + OnClose: func(latency int64, count int, iterErr, closeErr error) { + callbacks++ + assert.GreaterOrEqual(t, latency, int64(0)) + assert.Equal(t, 1+tc.wantFetches, count) + assert.NoError(t, iterErr) + assert.NoError(t, closeErr) + }, + }) + require.NoError(t, err) + defer dr.Close() //nolint:errcheck + dest := make([]driver.Value, 1) + wantRows := 2 * (1 + tc.wantFetches) + if tc.empty { + wantRows = 0 + } + for i := 0; i < wantRows; i++ { + require.NoError(t, dr.Next(dest)) + assert.Equal(t, i%2 == 0, dest[0]) + assert.Zero(t, closes, "iteration must not close the operation") + } + require.ErrorIs(t, dr.Next(dest), io.EOF) + require.ErrorIs(t, dr.Next(dest), io.EOF) + assert.Zero(t, closes) + assert.Zero(t, callbacks) + assert.Equal(t, tc.wantFetches, fetches) + assert.Len(t, chunkCounts, 1+tc.wantFetches) + if tc.cloudFetch { + assert.Equal(t, wantRows/2, downloads) + } + require.NoError(t, dr.Close()) + require.NoError(t, dr.Close()) + wantCloses := 1 + if tc.closedOnServer { + wantCloses = 0 + } + assert.Equal(t, wantCloses, closes) + assert.Equal(t, 1, callbacks) + }) + } +} + type fetch struct { direction cli_service.TFetchOrientation resultStartRec int @@ -925,7 +1070,7 @@ func TestFetchResultsWithRetries(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, + false, false, client, rowSet.logger(), ) @@ -1823,6 +1968,8 @@ func TestNewRows_DetachesResultRPCContextFromQueryContextCancellation(t *testing assert.Equal(t, true, dest[0]) assert.True(t, fetchCalled, "FetchResults should use the detached result context") assert.True(t, metaCalled, "GetResultSetMetadata should use the detached result context") + assert.False(t, closeCalled, "iteration must not close the operation") + assert.NoError(t, dr.Close()) assert.True(t, closeCalled, "CloseOperation should use the detached result context") } @@ -1942,6 +2089,9 @@ func TestRows_CloseCallback_ReceivesChunkCount(t *testing.T) { testClient := &client.TestClient{ FnFetchResults: fetchFn, FnGetResultSetMetadata: metaFn, + FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + return &cli_service.TCloseOperationResp{}, nil + }, } var callbackChunkCount int @@ -1970,6 +2120,62 @@ func TestRows_CloseCallback_ReceivesChunkCount(t *testing.T) { "closeCallback must receive the total number of result pages fetched") } +func TestRowsCloseReportsErrors(t *testing.T) { + moreRows := true + for _, fetchErr := range []error{nil, errors.New("fetch failed")} { + for _, closeErr := range []error{nil, errors.New("close failed")} { + t.Run(fmt.Sprintf("fetch=%v/close=%v", fetchErr, closeErr), func(t *testing.T) { + closes, callbacks := 0, 0 + testClient := &client.TestClient{ + FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + return nil, fetchErr + }, + FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closes++ + return nil, closeErr + }, + } + direct := &cli_service.TSparkDirectResults{ + ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ + Schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ + ColumnName: "flag", + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE}, + }}}, + }}}, + }, + ResultSet: &cli_service.TFetchResultsResp{ + HasMoreRows: &moreRows, + Results: &cli_service.TRowSet{Columns: []*cli_service.TColumn{ + {BoolVal: &cli_service.TBoolColumn{Values: []bool{true}}}, + }}, + }, + } + dr, err := NewRows(context.Background(), nil, testClient, nil, direct, &TelemetryCallbacks{ + OnClose: func(_ int64, count int, iterErr, rpcErr error) { + callbacks++ + assert.Equal(t, 1, closes, "telemetry must run after the close RPC") + assert.Equal(t, 1, count) + assert.ErrorIs(t, iterErr, fetchErr) + assert.ErrorIs(t, rpcErr, closeErr) + }, + }) + require.NoError(t, err) + if fetchErr != nil { + dest := make([]driver.Value, 1) + require.NoError(t, dr.Next(dest)) + require.ErrorIs(t, dr.Next(dest), fetchErr) + } + assert.Zero(t, closes) + assert.ErrorIs(t, dr.Close(), closeErr) + assert.ErrorIs(t, dr.Close(), closeErr) + assert.Equal(t, 1, closes) + assert.Equal(t, 1, callbacks) + }) + } + } +} + // 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) { diff --git a/internal/rows/rowscanner/category_test.go b/internal/rows/rowscanner/category_test.go index ee42f6ff..b1305612 100644 --- a/internal/rows/rowscanner/category_test.go +++ b/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, false, 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, false, false, c, dbsqllog.WithContext("connId", "correlationId", "")) assert.False(t, rpf.HasNext()) diff --git a/internal/rows/rowscanner/resultPageIterator.go b/internal/rows/rowscanner/resultPageIterator.go index 82848abc..deeba305 100644 --- a/internal/rows/rowscanner/resultPageIterator.go +++ b/internal/rows/rowscanner/resultPageIterator.go @@ -53,17 +53,16 @@ func NewResultPageIterator( delimiter Delimiter, maxPageSize int64, opHandle *cli_service.TOperationHandle, + isFinished 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: isFinished || closedOnServer, maxPageSize: maxPageSize, opHandle: opHandle, closedOnServer: closedOnServer, @@ -143,7 +142,6 @@ func (rpf *resultPageIterator) HasNext() bool { if rpf.nextResultPage == nil { nrp, err := rpf.getNextPage() if err != nil { - rpf.Close() //nolint:errcheck,gosec // G104: close in error path rpf.isFinished = true rpf.err = err return false @@ -151,9 +149,6 @@ func (rpf *resultPageIterator) HasNext() bool { rpf.err = nil rpf.nextResultPage = nrp - if !nrp.GetHasMoreRows() { - rpf.Close() //nolint:errcheck,gosec // G104: close in error path - } } return rpf.nextResultPage != nil From 58d2937a5b2bf48a161b63fa5a6ca2626f5de19e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 04:39:09 +0000 Subject: [PATCH 2/7] fix: distinguish direct result state in iterator API Signed-off-by: Vu Anh Phung --- .../arrowbased/arrowRecordIterator_test.go | 24 +++++++++---------- internal/rows/arrowbased/arrowRows_test.go | 2 +- internal/rows/rows.go | 8 ++++--- internal/rows/rows_test.go | 10 ++++---- internal/rows/rowscanner/category_test.go | 4 ++-- .../rows/rowscanner/resultPageIterator.go | 4 ++-- 6 files changed, 27 insertions(+), 25 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 385c26af..3db6b526 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -42,7 +42,7 @@ func TestArrowRecordIterator(t *testing.T) { rowscanner.NewDelimiter(0, 7311), 5000, nil, - false, false, + nil, false, simpleClient, logger, ) @@ -140,7 +140,7 @@ func TestArrowRecordIterator(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, 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, 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, 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, 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, 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, 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, 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, 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, false, + nil, false, simpleClient, logger, ) @@ -563,7 +563,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, false, + nil, false, simpleClient, logger, ) @@ -605,7 +605,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, false, + nil, false, failingClient, logger, ) diff --git a/internal/rows/arrowbased/arrowRows_test.go b/internal/rows/arrowbased/arrowRows_test.go index 8903d782..27d28d61 100644 --- a/internal/rows/arrowbased/arrowRows_test.go +++ b/internal/rows/arrowbased/arrowRows_test.go @@ -1545,7 +1545,7 @@ func TestArrowRowScanner(t *testing.T) { rowscanner.NewDelimiter(0, 7311), 5000, nil, - false, false, + nil, false, client, logger) diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 82e7c8b4..1cd5b2fd 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -229,14 +229,16 @@ func NewRows( // CloudFetch may deliver all links without closing the server operation. closedOnServer := directResults != nil && directResults.CloseOperation != nil - isFinished := closedOnServer || (directResults != nil && directResults.ResultSet != nil && - directResults.ResultSet.HasMoreRows != nil && !*directResults.ResultSet.HasMoreRows) + var hasMoreRows *bool + if directResults != nil && directResults.ResultSet != nil { + hasMoreRows = directResults.ResultSet.HasMoreRows + } r.ResultPageIterator = rowscanner.NewResultPageIterator( resultsCtx, d, pageSize, opHandle, - isFinished, + hasMoreRows, closedOnServer, client, r.logger(), diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 15b55821..37f341c6 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -234,7 +234,7 @@ func TestRowsFetchResultPageNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, false, + nil, false, client, rowSet.logger(), ) @@ -330,7 +330,7 @@ func TestRowsFetchResultPageWithDirectResults(t *testing.T) { rowscanner.NewDelimiter(rowSet.RowScanner.Start(), rowSet.RowScanner.Count()), 1000, nil, - false, false, + nil, false, client, rowSet.logger(), ) @@ -484,7 +484,7 @@ func TestNextNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, false, + nil, false, client, rowSet.logger(), ) @@ -760,7 +760,7 @@ func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { rowSet.resultSetMetadata = metadata } rowSet.ResultPageIterator = rowscanner.NewResultPageIterator( - rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, true, true, testClient, rowSet.logger()) + rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, nil, true, testClient, rowSet.logger()) checkTypes := func() { t.Helper() @@ -1070,7 +1070,7 @@ func TestFetchResultsWithRetries(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - false, false, + nil, false, client, rowSet.logger(), ) diff --git a/internal/rows/rowscanner/category_test.go b/internal/rows/rowscanner/category_test.go index b1305612..1ed909a1 100644 --- a/internal/rows/rowscanner/category_test.go +++ b/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, 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, false, c, + rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, nil, false, c, dbsqllog.WithContext("connId", "correlationId", "")) assert.False(t, rpf.HasNext()) diff --git a/internal/rows/rowscanner/resultPageIterator.go b/internal/rows/rowscanner/resultPageIterator.go index deeba305..0b39ac64 100644 --- a/internal/rows/rowscanner/resultPageIterator.go +++ b/internal/rows/rowscanner/resultPageIterator.go @@ -53,7 +53,7 @@ func NewResultPageIterator( delimiter Delimiter, maxPageSize int64, opHandle *cli_service.TOperationHandle, - isFinished bool, + hasMoreRows *bool, closedOnServer bool, client cli_service.TCLIService, logger *dbsqllog.DBSQLLogger, @@ -62,7 +62,7 @@ func NewResultPageIterator( return &resultPageIterator{ ctx: ctx, Delimiter: delimiter, - isFinished: isFinished || closedOnServer, + isFinished: closedOnServer || (hasMoreRows != nil && !*hasMoreRows), maxPageSize: maxPageSize, opHandle: opHandle, closedOnServer: closedOnServer, From df3af976c20c98cefc16b56ecc58c596699c9496 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 04:45:24 +0000 Subject: [PATCH 3/7] docs: note CloudFetch lifecycle fix in changelog Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b736e9a..6102fbf7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Release History +## Unreleased + +- Avoid redundant CloudFetch result fetches and defer operation cleanup to `Rows.Close()` for `database/sql` iteration on the Thrift backend (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). From 426df48f791117a0f53a20fa760a391cd5fc350c Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 04:51:02 +0000 Subject: [PATCH 4/7] test: document and verify IPC stream close ownership Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- internal/rows/rows_test.go | 82 ++++++++++++++++++++++++++++++++++++++ rows/rows.go | 3 ++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6102fbf7..2a6867aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- Avoid redundant CloudFetch result fetches and defer operation cleanup to `Rows.Close()` for `database/sql` iteration on the Thrift backend (databricks/databricks-sql-go#539). +- Avoid redundant CloudFetch result fetches on the Thrift backend. Row and Arrow IPC stream iteration defer operation cleanup to the parent `Rows.Close()`, which callers must invoke even after draining or closing an IPC iterator (databricks/databricks-sql-go#539). ## v1.19.0 (2026-11-04) - Upgrade the SQL kernel bindings to v1.4.0. diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 37f341c6..7192ad0f 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -1289,6 +1289,88 @@ func TestGetArrowBatches(t *testing.T) { }) } +func TestRowsArrowIPCStreamsLifecycle(t *testing.T) { + moreRows, noMoreRows := true, false + for _, tc := range []struct { + name string + hasMoreRows *bool + wantFetches int + }{ + {name: "complete", hasMoreRows: &noMoreRows}, + {name: "paginated", hasMoreRows: &moreRows, wantFetches: 1}, + {name: "unknown", wantFetches: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + var execute cli_service.TExecuteStatementResp + loadTestData(t, "directResultsMultipleFetch/ExecuteStatement.json", &execute) + execute.DirectResults.CloseOperation = nil + execute.DirectResults.ResultSet.HasMoreRows = tc.hasMoreRows + var fetched cli_service.TFetchResultsResp + loadTestData(t, "directResultsMultipleFetch/FetchResults1.json", &fetched) + fetched.HasMoreRows = &noMoreRows + + fetches, closes, callbacks := 0, 0, 0 + closeErr := errors.New("close failed") + testClient := &client.TestClient{ + FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + fetches++ + require.LessOrEqual(t, fetches, tc.wantFetches) + return &fetched, nil + }, + FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closes++ + return nil, closeErr + }, + } + dr, err := NewRows(context.Background(), nil, testClient, config.WithDefaults(), execute.DirectResults, &TelemetryCallbacks{ + OnClose: func(_ int64, _ int, iterErr, rpcErr error) { + callbacks++ + assert.Equal(t, 1, closes) + assert.NoError(t, iterErr) + assert.ErrorIs(t, rpcErr, closeErr) + }, + }) + require.NoError(t, err) + defer dr.Close() //nolint:errcheck + streams, streamErr := dr.(dbsqlrows.Rows).GetArrowIPCStreams(context.Background()) + require.NoError(t, streamErr) + defer streams.Close() + + var gotRows, wantRows int64 + for _, batch := range execute.DirectResults.ResultSet.Results.ArrowBatches { + wantRows += batch.RowCount + } + if tc.wantFetches > 0 { + for _, batch := range fetched.Results.ArrowBatches { + wantRows += batch.RowCount + } + } + for streams.HasNext() { + stream, nextErr := streams.Next() + require.NoError(t, nextErr) + reader, readErr := ipc.NewReader(stream) + require.NoError(t, readErr) + for reader.Next() { + gotRows += reader.Record().NumRows() + } + readErr = reader.Err() + reader.Release() + require.NoError(t, readErr) + assert.Zero(t, closes, "iteration must not close the operation") + } + assert.Equal(t, wantRows, gotRows) + assert.Equal(t, tc.wantFetches, fetches) + streams.Close() + assert.Zero(t, closes, "closing the IPC iterator must leave operation cleanup to Rows.Close") + assert.Zero(t, callbacks) + assert.ErrorIs(t, dr.Close(), closeErr) + assert.ErrorIs(t, dr.Close(), closeErr) + assert.Equal(t, 1, closes) + assert.Equal(t, 1, callbacks) + }) + } +} + type rowTestPagingResult struct { getMetadataCount int fetchResultsCount int diff --git a/rows/rows.go b/rows/rows.go index 98ba71da..7e8dfa1e 100644 --- a/rows/rows.go +++ b/rows/rows.go @@ -9,6 +9,9 @@ import ( type Rows interface { GetArrowBatches(context.Context) (ArrowBatchIterator, error) + // GetArrowIPCStreams returns an iterator over Arrow IPC streams. + // Close the originating driver.Rows to release the server operation, + // even after the iterator is exhausted or closed. GetArrowIPCStreams(context.Context) (ArrowIPCStreamIterator, error) } From 4a9556f38901619df417b34ff7af6e660d75bee3 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 05:41:33 +0000 Subject: [PATCH 5/7] test: cover database/sql close errors after EOF Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + internal/rows/sql_test.go | 146 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 internal/rows/sql_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a6867aa..49cbc29e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## Unreleased - Avoid redundant CloudFetch result fetches on the Thrift backend. Row and Arrow IPC stream iteration defer operation cleanup to the parent `Rows.Close()`, which callers must invoke even after draining or closing an IPC iterator (databricks/databricks-sql-go#539). +- **Behavior change:** On the Thrift backend, `Rows.Err()` now reports a failed `CloseOperation` during automatic cleanup after all rows have been read. Check `Rows.Err()` even after consuming the full result (databricks/databricks-sql-go#539). ## v1.19.0 (2026-11-04) - Upgrade the SQL kernel bindings to v1.4.0. diff --git a/internal/rows/sql_test.go b/internal/rows/sql_test.go new file mode 100644 index 00000000..f5aed59d --- /dev/null +++ b/internal/rows/sql_test.go @@ -0,0 +1,146 @@ +package rows + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/ipc" + "github.com/apache/arrow/go/v12/arrow/memory" + "github.com/databricks/databricks-sql-go/internal/cli_service" + "github.com/databricks/databricks-sql-go/internal/client" + "github.com/databricks/databricks-sql-go/internal/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSQLRowsReportsCloseOperationError(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{{Name: "flag", Type: arrow.FixedWidthTypes.Boolean}}, nil) + builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer builder.Release() + builder.Field(0).(*array.BooleanBuilder).AppendValues([]bool{true, false}, nil) + record := builder.NewRecord() + defer record.Release() + var data bytes.Buffer + writer := ipc.NewWriter(&data, ipc.WithSchema(schema)) + require.NoError(t, writer.Write(record)) + require.NoError(t, writer.Close()) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write(data.Bytes()) + })) + defer server.Close() + page := func(offset int64) *cli_service.TRowSet { + return &cli_service.TRowSet{StartRowOffset: offset, ResultLinks: []*cli_service.TSparkArrowResultLink{{ + FileLink: server.URL, StartRowOffset: offset, RowCount: 2, + ExpiryTime: time.Now().Add(time.Hour).Unix(), + }}} + } + metadata := &cli_service.TGetResultSetMetadataResp{ + Schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ + ColumnName: "flag", + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE}, + }}}, + }}}, + } + for _, tc := range []struct { + name string + hasMoreRows bool + earlyClose bool + wantFetches int + wantRows []bool + }{ + {name: "complete direct results", wantRows: []bool{true, false}}, + {name: "paginated results", hasMoreRows: true, wantFetches: 1, wantRows: []bool{true, false, true, false}}, + {name: "explicit early close", hasMoreRows: true, earlyClose: true, wantRows: []bool{true}}, + } { + for _, closeErr := range []error{nil, errors.New("close failed")} { + t.Run(fmt.Sprintf("%s/close=%v", tc.name, closeErr), func(t *testing.T) { + fetches, closes, callbacks := 0, 0, 0 + noMoreRows := false + testClient := &client.TestClient{ + FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + fetches++ + require.LessOrEqual(t, fetches, tc.wantFetches) + return &cli_service.TFetchResultsResp{HasMoreRows: &noMoreRows, Results: page(2)}, nil + }, + FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closes++ + return nil, closeErr + }, + } + cfg := config.WithDefaults() + cfg.UseLz4Compression = false + dr, err := NewRows(context.Background(), nil, testClient, cfg, &cli_service.TSparkDirectResults{ + ResultSetMetadata: metadata, + ResultSet: &cli_service.TFetchResultsResp{ + HasMoreRows: &tc.hasMoreRows, Results: page(0), + }, + }, &TelemetryCallbacks{OnClose: func(_ int64, count int, iterErr, rpcErr error) { + callbacks++ + assert.Equal(t, 1, closes) + assert.Equal(t, 1+tc.wantFetches, count) + assert.NoError(t, iterErr) + assert.ErrorIs(t, rpcErr, closeErr) + }}) + require.NoError(t, err) + defer dr.Close() //nolint:errcheck + db := sql.OpenDB(&sqlRowsTestDriver{result: dr}) + t.Cleanup(func() { assert.NoError(t, db.Close()) }) + result, queryErr := db.QueryContext(context.Background(), "select flag") + require.NoError(t, queryErr) + defer result.Close() //nolint:errcheck + var got []bool + for result.Next() { + var value bool + require.NoError(t, result.Scan(&value)) + got = append(got, value) + assert.Zero(t, closes) + if tc.earlyClose { + break + } + } + assert.Equal(t, tc.wantRows, got) + if tc.earlyClose { + assert.ErrorIs(t, result.Close(), closeErr) + } else { + assert.ErrorIs(t, result.Err(), closeErr) + assert.Equal(t, 1, closes) + assert.NoError(t, result.Close()) + } + assert.ErrorIs(t, result.Err(), closeErr) + assert.ErrorIs(t, dr.Close(), closeErr) + assert.Equal(t, tc.wantFetches, fetches) + assert.Equal(t, 1, closes) + assert.Equal(t, 1, callbacks) + }) + } + } +} + +type sqlRowsTestDriver struct { + result driver.Rows +} + +func (d *sqlRowsTestDriver) Connect(context.Context) (driver.Conn, error) { return d, nil } +func (d *sqlRowsTestDriver) Driver() driver.Driver { return d } +func (d *sqlRowsTestDriver) Open(string) (driver.Conn, error) { return d, nil } +func (d *sqlRowsTestDriver) Close() error { return nil } +func (d *sqlRowsTestDriver) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare is unsupported") +} +func (d *sqlRowsTestDriver) Begin() (driver.Tx, error) { + return nil, errors.New("transactions are unsupported") +} +func (d *sqlRowsTestDriver) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) { + return d.result, nil +} From e7992df6527f42fdb2c0c2a240eeba3d0f95e418 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 05:46:06 +0000 Subject: [PATCH 6/7] docs: clarify cleanup errors and direct driver ownership Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 49cbc29e..fed25cf5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,8 @@ ## Unreleased -- Avoid redundant CloudFetch result fetches on the Thrift backend. Row and Arrow IPC stream iteration defer operation cleanup to the parent `Rows.Close()`, which callers must invoke even after draining or closing an IPC iterator (databricks/databricks-sql-go#539). -- **Behavior change:** On the Thrift backend, `Rows.Err()` now reports a failed `CloseOperation` during automatic cleanup after all rows have been read. Check `Rows.Err()` even after consuming the full result (databricks/databricks-sql-go#539). +- Avoid redundant CloudFetch result fetches on the Thrift backend. Row and Arrow IPC stream iteration defer operation cleanup to the parent `Rows.Close()`. Callers using driver rows or IPC streams directly must close the parent rows after use, including after iteration errors (databricks/databricks-sql-go#539). +- **Behavior change:** For Thrift queries through `database/sql`, `sql.Rows.Err()` now reports a failed `CloseOperation` during automatic cleanup after all rows have been read (databricks/databricks-sql-go#539). ## v1.19.0 (2026-11-04) - Upgrade the SQL kernel bindings to v1.4.0. From 9101d3bde31e93bed16b2a45c020b1832eb1fb3b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 06:09:47 +0000 Subject: [PATCH 7/7] fix: preserve close semantics when skipping redundant result fetches Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 3 +- internal/rows/rows.go | 8 - internal/rows/rows_test.go | 300 +----------------- .../rows/rowscanner/resultPageIterator.go | 6 + internal/rows/sql_test.go | 75 +++-- rows/rows.go | 3 - 6 files changed, 70 insertions(+), 325 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fed25cf5..da6c8e8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,7 @@ ## Unreleased -- Avoid redundant CloudFetch result fetches on the Thrift backend. Row and Arrow IPC stream iteration defer operation cleanup to the parent `Rows.Close()`. Callers using driver rows or IPC streams directly must close the parent rows after use, including after iteration errors (databricks/databricks-sql-go#539). -- **Behavior change:** For Thrift queries through `database/sql`, `sql.Rows.Err()` now reports a failed `CloseOperation` during automatic cleanup after all rows have been read (databricks/databricks-sql-go#539). +- 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. diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 1cd5b2fd..c68d6b1f 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -76,8 +76,6 @@ type rows struct { // resultsCancel aborts in-flight result RPCs/downloads when Close() is // called, so the detached ctx never leaves an operation uncancellable. resultsCancel context.CancelFunc - closeOnce sync.Once - closeErr error // Telemetry tracking // telemetryUpdate is called after each chunk is fetched with: @@ -277,13 +275,7 @@ func (r *rows) Close() error { if r == nil { return nil } - r.closeOnce.Do(func() { - r.closeErr = r.close() - }) - return r.closeErr -} -func (r *rows) close() error { // Release the detached results context after the close RPC runs, aborting // any in-flight FetchResults/CloudFetch downloads still referencing it. if r.resultsCancel != nil { diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 7192ad0f..78cba9a1 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -9,17 +9,13 @@ import ( "fmt" "io" "math" - "net/http" - "net/http/httptest" "os" "reflect" "testing" "time" "github.com/apache/arrow/go/v12/arrow" - "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/ipc" - "github.com/apache/arrow/go/v12/arrow/memory" "github.com/databricks/databricks-sql-go/driverctx" "github.com/databricks/databricks-sql-go/internal/cli_service" "github.com/databricks/databricks-sql-go/internal/client" @@ -234,7 +230,8 @@ func TestRowsFetchResultPageNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - nil, false, + nil, + false, client, rowSet.logger(), ) @@ -330,7 +327,8 @@ func TestRowsFetchResultPageWithDirectResults(t *testing.T) { rowscanner.NewDelimiter(rowSet.RowScanner.Start(), rowSet.RowScanner.Count()), 1000, nil, - nil, false, + nil, + false, client, rowSet.logger(), ) @@ -484,7 +482,8 @@ func TestNextNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - nil, false, + nil, + false, client, rowSet.logger(), ) @@ -906,147 +905,6 @@ func TestRowsCloseOptimization(t *testing.T) { assert.Equal(t, 0, closeCount) } -func TestRowsDirectResultsLifecycle(t *testing.T) { - moreRows, noMoreRows := true, false - for _, tc := range []struct { - name string - cloudFetch bool - hasMoreRows *bool - closedOnServer bool - empty bool - wantFetches int - }{ - {name: "cloudfetch complete", cloudFetch: true, hasMoreRows: &noMoreRows}, - {name: "cloudfetch paginated", cloudFetch: true, hasMoreRows: &moreRows, wantFetches: 1}, - {name: "cloudfetch unknown", cloudFetch: true, wantFetches: 1}, - {name: "cloudfetch empty", cloudFetch: true, hasMoreRows: &noMoreRows, empty: true}, - {name: "inline closed", closedOnServer: true}, - {name: "inline complete", hasMoreRows: &noMoreRows}, - {name: "inline unknown", wantFetches: 1}, - } { - t.Run(tc.name, func(t *testing.T) { - cfg := config.WithDefaults() - cfg.UseLz4Compression = false - metadata := &cli_service.TGetResultSetMetadataResp{ - Schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ - ColumnName: "flag", - TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ - PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE}, - }}}, - }}}, - } - var fileURL string - if tc.cloudFetch { - schema := arrow.NewSchema([]arrow.Field{{Name: "flag", Type: arrow.FixedWidthTypes.Boolean}}, nil) - builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) - defer builder.Release() - builder.Field(0).(*array.BooleanBuilder).AppendValues([]bool{true, false}, nil) - record := builder.NewRecord() - defer record.Release() - var data bytes.Buffer - writer := ipc.NewWriter(&data, ipc.WithSchema(schema)) - require.NoError(t, writer.Write(record)) - require.NoError(t, writer.Close()) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - _, _ = w.Write(data.Bytes()) - })) - defer server.Close() - fileURL = server.URL - } - page := func(start int64, empty bool) *cli_service.TRowSet { - rs := &cli_service.TRowSet{StartRowOffset: start} - if tc.cloudFetch { - rs.ResultLinks = []*cli_service.TSparkArrowResultLink{} - if !empty { - rs.ResultLinks = append(rs.ResultLinks, &cli_service.TSparkArrowResultLink{ - FileLink: fileURL, StartRowOffset: start, RowCount: 2, - ExpiryTime: time.Now().Add(time.Hour).Unix(), - }) - } - } else { - rs.Columns = []*cli_service.TColumn{{BoolVal: &cli_service.TBoolColumn{Values: []bool{true, false}}}} - } - return rs - } - opHandle := &cli_service.TOperationHandle{OperationId: &cli_service.THandleIdentifier{GUID: make([]byte, 16)}} - fetches, closes, callbacks, downloads := 0, 0, 0, 0 - testClient := &client.TestClient{ - FnFetchResults: func(_ context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { - fetches++ - require.LessOrEqual(t, fetches, tc.wantFetches) - assert.Same(t, opHandle, req.OperationHandle) - return &cli_service.TFetchResultsResp{HasMoreRows: &noMoreRows, Results: page(2, false)}, nil - }, - FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { - closes++ - assert.NoError(t, ctx.Err()) - assert.Same(t, opHandle, req.OperationHandle) - return &cli_service.TCloseOperationResp{}, nil - }, - } - direct := &cli_service.TSparkDirectResults{ - ResultSetMetadata: metadata, - ResultSet: &cli_service.TFetchResultsResp{ - HasMoreRows: tc.hasMoreRows, Results: page(0, tc.empty), - }, - } - if tc.closedOnServer { - direct.CloseOperation = &cli_service.TCloseOperationResp{} - } - var chunkCounts []int - dr, err := NewRows(context.Background(), opHandle, testClient, cfg, direct, &TelemetryCallbacks{ - OnChunkFetched: func(count int, _ int64, index int, latency int64, total int32) { - chunkCounts = append(chunkCounts, count) - assert.Equal(t, count-1, index) - if tc.cloudFetch { - assert.Zero(t, latency) - if !tc.empty { - assert.EqualValues(t, 1, total) - } - } - }, - OnCloudFetchFile: func(int64) { downloads++ }, - OnClose: func(latency int64, count int, iterErr, closeErr error) { - callbacks++ - assert.GreaterOrEqual(t, latency, int64(0)) - assert.Equal(t, 1+tc.wantFetches, count) - assert.NoError(t, iterErr) - assert.NoError(t, closeErr) - }, - }) - require.NoError(t, err) - defer dr.Close() //nolint:errcheck - dest := make([]driver.Value, 1) - wantRows := 2 * (1 + tc.wantFetches) - if tc.empty { - wantRows = 0 - } - for i := 0; i < wantRows; i++ { - require.NoError(t, dr.Next(dest)) - assert.Equal(t, i%2 == 0, dest[0]) - assert.Zero(t, closes, "iteration must not close the operation") - } - require.ErrorIs(t, dr.Next(dest), io.EOF) - require.ErrorIs(t, dr.Next(dest), io.EOF) - assert.Zero(t, closes) - assert.Zero(t, callbacks) - assert.Equal(t, tc.wantFetches, fetches) - assert.Len(t, chunkCounts, 1+tc.wantFetches) - if tc.cloudFetch { - assert.Equal(t, wantRows/2, downloads) - } - require.NoError(t, dr.Close()) - require.NoError(t, dr.Close()) - wantCloses := 1 - if tc.closedOnServer { - wantCloses = 0 - } - assert.Equal(t, wantCloses, closes) - assert.Equal(t, 1, callbacks) - }) - } -} - type fetch struct { direction cli_service.TFetchOrientation resultStartRec int @@ -1070,7 +928,8 @@ func TestFetchResultsWithRetries(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, - nil, false, + nil, + false, client, rowSet.logger(), ) @@ -1289,88 +1148,6 @@ func TestGetArrowBatches(t *testing.T) { }) } -func TestRowsArrowIPCStreamsLifecycle(t *testing.T) { - moreRows, noMoreRows := true, false - for _, tc := range []struct { - name string - hasMoreRows *bool - wantFetches int - }{ - {name: "complete", hasMoreRows: &noMoreRows}, - {name: "paginated", hasMoreRows: &moreRows, wantFetches: 1}, - {name: "unknown", wantFetches: 1}, - } { - t.Run(tc.name, func(t *testing.T) { - var execute cli_service.TExecuteStatementResp - loadTestData(t, "directResultsMultipleFetch/ExecuteStatement.json", &execute) - execute.DirectResults.CloseOperation = nil - execute.DirectResults.ResultSet.HasMoreRows = tc.hasMoreRows - var fetched cli_service.TFetchResultsResp - loadTestData(t, "directResultsMultipleFetch/FetchResults1.json", &fetched) - fetched.HasMoreRows = &noMoreRows - - fetches, closes, callbacks := 0, 0, 0 - closeErr := errors.New("close failed") - testClient := &client.TestClient{ - FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { - fetches++ - require.LessOrEqual(t, fetches, tc.wantFetches) - return &fetched, nil - }, - FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { - closes++ - return nil, closeErr - }, - } - dr, err := NewRows(context.Background(), nil, testClient, config.WithDefaults(), execute.DirectResults, &TelemetryCallbacks{ - OnClose: func(_ int64, _ int, iterErr, rpcErr error) { - callbacks++ - assert.Equal(t, 1, closes) - assert.NoError(t, iterErr) - assert.ErrorIs(t, rpcErr, closeErr) - }, - }) - require.NoError(t, err) - defer dr.Close() //nolint:errcheck - streams, streamErr := dr.(dbsqlrows.Rows).GetArrowIPCStreams(context.Background()) - require.NoError(t, streamErr) - defer streams.Close() - - var gotRows, wantRows int64 - for _, batch := range execute.DirectResults.ResultSet.Results.ArrowBatches { - wantRows += batch.RowCount - } - if tc.wantFetches > 0 { - for _, batch := range fetched.Results.ArrowBatches { - wantRows += batch.RowCount - } - } - for streams.HasNext() { - stream, nextErr := streams.Next() - require.NoError(t, nextErr) - reader, readErr := ipc.NewReader(stream) - require.NoError(t, readErr) - for reader.Next() { - gotRows += reader.Record().NumRows() - } - readErr = reader.Err() - reader.Release() - require.NoError(t, readErr) - assert.Zero(t, closes, "iteration must not close the operation") - } - assert.Equal(t, wantRows, gotRows) - assert.Equal(t, tc.wantFetches, fetches) - streams.Close() - assert.Zero(t, closes, "closing the IPC iterator must leave operation cleanup to Rows.Close") - assert.Zero(t, callbacks) - assert.ErrorIs(t, dr.Close(), closeErr) - assert.ErrorIs(t, dr.Close(), closeErr) - assert.Equal(t, 1, closes) - assert.Equal(t, 1, callbacks) - }) - } -} - type rowTestPagingResult struct { getMetadataCount int fetchResultsCount int @@ -2050,8 +1827,6 @@ func TestNewRows_DetachesResultRPCContextFromQueryContextCancellation(t *testing assert.Equal(t, true, dest[0]) assert.True(t, fetchCalled, "FetchResults should use the detached result context") assert.True(t, metaCalled, "GetResultSetMetadata should use the detached result context") - assert.False(t, closeCalled, "iteration must not close the operation") - assert.NoError(t, dr.Close()) assert.True(t, closeCalled, "CloseOperation should use the detached result context") } @@ -2171,9 +1946,6 @@ func TestRows_CloseCallback_ReceivesChunkCount(t *testing.T) { testClient := &client.TestClient{ FnFetchResults: fetchFn, FnGetResultSetMetadata: metaFn, - FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { - return &cli_service.TCloseOperationResp{}, nil - }, } var callbackChunkCount int @@ -2202,62 +1974,6 @@ func TestRows_CloseCallback_ReceivesChunkCount(t *testing.T) { "closeCallback must receive the total number of result pages fetched") } -func TestRowsCloseReportsErrors(t *testing.T) { - moreRows := true - for _, fetchErr := range []error{nil, errors.New("fetch failed")} { - for _, closeErr := range []error{nil, errors.New("close failed")} { - t.Run(fmt.Sprintf("fetch=%v/close=%v", fetchErr, closeErr), func(t *testing.T) { - closes, callbacks := 0, 0 - testClient := &client.TestClient{ - FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { - return nil, fetchErr - }, - FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { - closes++ - return nil, closeErr - }, - } - direct := &cli_service.TSparkDirectResults{ - ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ - ColumnName: "flag", - TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ - PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_BOOLEAN_TYPE}, - }}}, - }}}, - }, - ResultSet: &cli_service.TFetchResultsResp{ - HasMoreRows: &moreRows, - Results: &cli_service.TRowSet{Columns: []*cli_service.TColumn{ - {BoolVal: &cli_service.TBoolColumn{Values: []bool{true}}}, - }}, - }, - } - dr, err := NewRows(context.Background(), nil, testClient, nil, direct, &TelemetryCallbacks{ - OnClose: func(_ int64, count int, iterErr, rpcErr error) { - callbacks++ - assert.Equal(t, 1, closes, "telemetry must run after the close RPC") - assert.Equal(t, 1, count) - assert.ErrorIs(t, iterErr, fetchErr) - assert.ErrorIs(t, rpcErr, closeErr) - }, - }) - require.NoError(t, err) - if fetchErr != nil { - dest := make([]driver.Value, 1) - require.NoError(t, dr.Next(dest)) - require.ErrorIs(t, dr.Next(dest), fetchErr) - } - assert.Zero(t, closes) - assert.ErrorIs(t, dr.Close(), closeErr) - assert.ErrorIs(t, dr.Close(), closeErr) - assert.Equal(t, 1, closes) - assert.Equal(t, 1, callbacks) - }) - } - } -} - // 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) { diff --git a/internal/rows/rowscanner/resultPageIterator.go b/internal/rows/rowscanner/resultPageIterator.go index 0b39ac64..fed91255 100644 --- a/internal/rows/rowscanner/resultPageIterator.go +++ b/internal/rows/rowscanner/resultPageIterator.go @@ -129,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. @@ -142,6 +144,7 @@ func (rpf *resultPageIterator) HasNext() bool { if rpf.nextResultPage == nil { nrp, err := rpf.getNextPage() if err != nil { + rpf.Close() //nolint:errcheck,gosec // G104: close in error path rpf.isFinished = true rpf.err = err return false @@ -149,6 +152,9 @@ func (rpf *resultPageIterator) HasNext() bool { rpf.err = nil rpf.nextResultPage = nrp + if !nrp.GetHasMoreRows() { + rpf.Close() //nolint:errcheck,gosec // G104: close in error path + } } return rpf.nextResultPage != nil diff --git a/internal/rows/sql_test.go b/internal/rows/sql_test.go index f5aed59d..e11c97ed 100644 --- a/internal/rows/sql_test.go +++ b/internal/rows/sql_test.go @@ -23,7 +23,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestSQLRowsReportsCloseOperationError(t *testing.T) { +func TestSQLRowsDirectResultsPreserveCloseSemantics(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{{Name: "flag", Type: arrow.FixedWidthTypes.Boolean}}, nil) builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) defer builder.Release() @@ -52,25 +52,45 @@ func TestSQLRowsReportsCloseOperationError(t *testing.T) { }}}, }}}, } + moreRows, noMoreRows := true, false + fetchErr := errors.New("fetch failed") for _, tc := range []struct { name string - hasMoreRows bool + hasMoreRows *bool + preclosed bool + empty bool + fetchErr error earlyClose bool wantFetches int wantRows []bool }{ - {name: "complete direct results", wantRows: []bool{true, false}}, - {name: "paginated results", hasMoreRows: true, wantFetches: 1, wantRows: []bool{true, false, true, false}}, - {name: "explicit early close", hasMoreRows: true, earlyClose: true, wantRows: []bool{true}}, + {name: "complete direct results", hasMoreRows: &noMoreRows, wantRows: []bool{true, false}}, + {name: "empty direct results", hasMoreRows: &noMoreRows, empty: true}, + {name: "paginated results", hasMoreRows: &moreRows, wantFetches: 1, wantRows: []bool{true, false, true, false}}, + {name: "unknown completion", wantFetches: 1, wantRows: []bool{true, false, true, false}}, + {name: "already closed", hasMoreRows: &noMoreRows, preclosed: true, wantRows: []bool{true, false}}, + {name: "early close complete direct results", hasMoreRows: &noMoreRows, earlyClose: true, wantRows: []bool{true}}, + {name: "early close paginated results", hasMoreRows: &moreRows, earlyClose: true, wantRows: []bool{true}}, + {name: "fetch error", hasMoreRows: &moreRows, fetchErr: fetchErr, wantFetches: 1, wantRows: []bool{true, false}}, } { for _, closeErr := range []error{nil, errors.New("close failed")} { t.Run(fmt.Sprintf("%s/close=%v", tc.name, closeErr), func(t *testing.T) { fetches, closes, callbacks := 0, 0, 0 - noMoreRows := false + wantCloses := 1 + if tc.preclosed { + wantCloses = 0 + } + var wantCloseErr error + if tc.earlyClose { + wantCloseErr = closeErr + } testClient := &client.TestClient{ FnFetchResults: func(context.Context, *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { fetches++ require.LessOrEqual(t, fetches, tc.wantFetches) + if tc.fetchErr != nil { + return nil, tc.fetchErr + } return &cli_service.TFetchResultsResp{HasMoreRows: &noMoreRows, Results: page(2)}, nil }, FnCloseOperation: func(context.Context, *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { @@ -80,20 +100,30 @@ func TestSQLRowsReportsCloseOperationError(t *testing.T) { } cfg := config.WithDefaults() cfg.UseLz4Compression = false - dr, err := NewRows(context.Background(), nil, testClient, cfg, &cli_service.TSparkDirectResults{ + direct := &cli_service.TSparkDirectResults{ ResultSetMetadata: metadata, ResultSet: &cli_service.TFetchResultsResp{ - HasMoreRows: &tc.hasMoreRows, Results: page(0), + HasMoreRows: tc.hasMoreRows, Results: page(0), }, - }, &TelemetryCallbacks{OnClose: func(_ int64, count int, iterErr, rpcErr error) { + } + if tc.empty { + direct.ResultSet.Results.ResultLinks = []*cli_service.TSparkArrowResultLink{} + } + if tc.preclosed { + direct.CloseOperation = &cli_service.TCloseOperationResp{} + } + dr, err := NewRows(context.Background(), nil, testClient, cfg, direct, &TelemetryCallbacks{OnClose: func(_ int64, count int, iterErr, rpcErr error) { callbacks++ - assert.Equal(t, 1, closes) - assert.Equal(t, 1+tc.wantFetches, count) - assert.NoError(t, iterErr) - assert.ErrorIs(t, rpcErr, closeErr) + assert.Equal(t, wantCloses, closes) + wantChunks := 1 + tc.wantFetches + if tc.fetchErr != nil { + wantChunks = 1 + } + assert.Equal(t, wantChunks, count) + assert.ErrorIs(t, iterErr, tc.fetchErr) + assert.ErrorIs(t, rpcErr, wantCloseErr) }}) require.NoError(t, err) - defer dr.Close() //nolint:errcheck db := sql.OpenDB(&sqlRowsTestDriver{result: dr}) t.Cleanup(func() { assert.NoError(t, db.Close()) }) result, queryErr := db.QueryContext(context.Background(), "select flag") @@ -104,7 +134,9 @@ func TestSQLRowsReportsCloseOperationError(t *testing.T) { var value bool require.NoError(t, result.Scan(&value)) got = append(got, value) - assert.Zero(t, closes) + if len(got) <= 2 { + assert.Zero(t, closes, "direct rows must remain readable before cleanup") + } if tc.earlyClose { break } @@ -113,14 +145,17 @@ func TestSQLRowsReportsCloseOperationError(t *testing.T) { if tc.earlyClose { assert.ErrorIs(t, result.Close(), closeErr) } else { - assert.ErrorIs(t, result.Err(), closeErr) - assert.Equal(t, 1, closes) + assert.ErrorIs(t, result.Err(), tc.fetchErr) + assert.Equal(t, wantCloses, closes) assert.NoError(t, result.Close()) } - assert.ErrorIs(t, result.Err(), closeErr) - assert.ErrorIs(t, dr.Close(), closeErr) + wantErr := tc.fetchErr + if tc.earlyClose { + wantErr = closeErr + } + assert.ErrorIs(t, result.Err(), wantErr) assert.Equal(t, tc.wantFetches, fetches) - assert.Equal(t, 1, closes) + assert.Equal(t, wantCloses, closes) assert.Equal(t, 1, callbacks) }) } diff --git a/rows/rows.go b/rows/rows.go index 7e8dfa1e..98ba71da 100644 --- a/rows/rows.go +++ b/rows/rows.go @@ -9,9 +9,6 @@ import ( type Rows interface { GetArrowBatches(context.Context) (ArrowBatchIterator, error) - // GetArrowIPCStreams returns an iterator over Arrow IPC streams. - // Close the originating driver.Rows to release the server operation, - // even after the iterator is exhausted or closed. GetArrowIPCStreams(context.Context) (ArrowIPCStreamIterator, error) }