From 02490755510693c99f0560e4466f57498b06a34a Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:27:08 +0000 Subject: [PATCH 1/8] ArrowBatchIterator turns a failed FetchResults into io.EOF (#507) Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> --- .../rows/arrowbased/arrowRecordIterator.go | 40 ++++++++++++++--- .../arrowbased/arrowRecordIterator_test.go | 43 +++++++++++++++++++ .../rows/rowscanner/resultPageIterator.go | 22 +++++++++- 3 files changed, 98 insertions(+), 7 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index 820ba75f..50614345 100644 --- a/internal/rows/arrowbased/arrowRecordIterator.go +++ b/internal/rows/arrowbased/arrowRecordIterator.go @@ -38,6 +38,14 @@ type arrowRecordIterator struct { isFinished bool arrowSchemaBytes []byte arrowSchema *arrow.Schema + // err holds a page-fetch error detected while advancing the result page + // iterator. It must be surfaced through Next() rather than being masked as + // a normal end-of-results (io.EOF). See issue #507. + err error + // errReported is set once the page-fetch error has been surfaced through + // Next() so that checkFinished does not re-capture it from the result page + // iterator (whose Err() keeps returning the stored error). + errReported bool } var _ rows.ArrowBatchIterator = (*arrowRecordIterator)(nil) @@ -49,6 +57,17 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) { return nil, io.EOF } + // If advancing the result page failed, surface that error rather than + // masking a failed page fetch as a normal end-of-results. See #507. The + // error is reported once and then the iterator is finished. + if ri.err != nil { + err := ri.err + ri.err = nil + ri.errReported = true + ri.Close() + return nil, err + } + // make sure we have the current batch err := ri.getCurrentBatch() if err != nil { @@ -66,7 +85,9 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) { // Indicate whether there are any more records available func (ri *arrowRecordIterator) HasNext() bool { ri.checkFinished() - return !ri.isFinished + // A pending page-fetch error still needs to be reported via Next(), so the + // iterator is not considered exhausted until that error is consumed. + return !ri.isFinished || ri.err != nil } // Free any resources associated with this iterator @@ -88,10 +109,19 @@ func (ri *arrowRecordIterator) Close() { } func (ri *arrowRecordIterator) checkFinished() { - finished := ri.isFinished || - ((ri.currentBatch == nil || !ri.currentBatch.HasNext()) && - (ri.batchIterator == nil || !ri.batchIterator.HasNext()) && - (ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext())) + noLocalData := (ri.currentBatch == nil || !ri.currentBatch.HasNext()) && + (ri.batchIterator == nil || !ri.batchIterator.HasNext()) + + noMorePages := ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() + + finished := ri.isFinished || (noLocalData && noMorePages) + + // If the result page iterator reported no more pages because a page fetch + // failed (rather than because the result set was fully consumed), hold on + // to that error so Next() can surface it instead of io.EOF. See #507. + if noLocalData && noMorePages && ri.resultPageIterator != nil && ri.err == nil && !ri.errReported { + ri.err = ri.resultPageIterator.Err() + } if finished { // Reached end of result set so Close diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 13891fc9..a5112d1c 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -193,6 +193,49 @@ func TestArrowRecordIterator(t *testing.T) { }) } +func TestArrowRecordIteratorSwallowsPageError(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + executeStatementResp := cli_service.TExecuteStatementResp{} + loadTestData2(t, "directResultsMultipleFetch/ExecuteStatement.json", &executeStatementResp) + fetchResp1 := cli_service.TFetchResultsResp{} + loadTestData2(t, "directResultsMultipleFetch/FetchResults1.json", &fetchResp1) + + calls := 0 + failing := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + if calls == 1 { + return &fetchResp1, nil + } + return nil, fmt.Errorf("FetchResults: connection reset") + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, 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) + it := NewArrowRecordIterator(ctx, rpi, bi, executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema, cfg) + defer it.Close() + + var lastErr error + rows := int64(0) + for it.HasNext() { + rec, err := it.Next() + if err != nil { + lastErr = err + break + } + rows += rec.NumRows() + rec.Release() + } + t.Logf("fetch calls=%d rows=%d lastErr=%v", calls, rows, lastErr) + assert.ErrorContains(t, lastErr, "connection reset") +} + func TestArrowRecordIteratorSchema(t *testing.T) { // Test with arrowSchemaBytes available t.Run("schema with initial schema bytes", func(t *testing.T) { diff --git a/internal/rows/rowscanner/resultPageIterator.go b/internal/rows/rowscanner/resultPageIterator.go index 22f2af25..82848abc 100644 --- a/internal/rows/rowscanner/resultPageIterator.go +++ b/internal/rows/rowscanner/resultPageIterator.go @@ -23,6 +23,10 @@ func errRowsUnandledFetchDirection(dir string) string { type ResultPageIterator interface { Next() (*cli_service.TFetchResultsResp, error) HasNext() bool + // Err returns the error, if any, that caused HasNext() to report that there + // are no more pages. It returns nil when the iterator finished normally + // (i.e. the result set was fully consumed). + Err() error Close() error Delimiter } @@ -112,12 +116,26 @@ type resultPageIterator struct { var _ ResultPageIterator = (*resultPageIterator)(nil) +// Err returns the error, if any, that caused HasNext() to report that there are +// no more pages. A normal end-of-results (io.EOF) is not considered an error and +// returns nil, so callers can distinguish a failed page fetch from a fully +// consumed result set. +func (rpf *resultPageIterator) Err() error { + if rpf == nil || rpf.err == nil || rpf.err == io.EOF { + return nil + } + return rpf.err +} + // Returns true if there are more pages in the result set. func (rpf *resultPageIterator) HasNext() bool { if rpf.isFinished && rpf.nextResultPage == nil { // There are no more pages to load and there isn't an already fetched - // page waiting to retrieved by Next() - rpf.err = io.EOF + // page waiting to retrieved by Next(). Preserve a previously stored + // fetch error rather than masking it as a normal end-of-results. + if rpf.err == nil { + rpf.err = io.EOF + } return false } From bd22f6d6fcfa33653a56226d6a365ee2279c71a9 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Wed, 21 Oct 2026 20:35:53 +0000 Subject: [PATCH 2/8] ai: apply changes for #509 (1 review thread) Addresses: - #4341629617 at internal/rows/arrowbased/arrowRecordIterator.go:115 Signed-off-by: peco-engineer-bot[bot] --- internal/rows/arrowbased/arrowRecordIterator.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index 50614345..7c11c7a7 100644 --- a/internal/rows/arrowbased/arrowRecordIterator.go +++ b/internal/rows/arrowbased/arrowRecordIterator.go @@ -112,14 +112,18 @@ func (ri *arrowRecordIterator) checkFinished() { noLocalData := (ri.currentBatch == nil || !ri.currentBatch.HasNext()) && (ri.batchIterator == nil || !ri.batchIterator.HasNext()) - noMorePages := ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() + // resultPageIterator.HasNext() issues a real FetchResults round-trip, so it + // is evaluated lazily — only once all local data is drained — to avoid + // eagerly prefetching the next page while the current one is still in use. + noMoreData := noLocalData && + (ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext()) - finished := ri.isFinished || (noLocalData && noMorePages) + finished := ri.isFinished || noMoreData // If the result page iterator reported no more pages because a page fetch // failed (rather than because the result set was fully consumed), hold on // to that error so Next() can surface it instead of io.EOF. See #507. - if noLocalData && noMorePages && ri.resultPageIterator != nil && ri.err == nil && !ri.errReported { + if noMoreData && ri.resultPageIterator != nil && ri.err == nil && !ri.errReported { ri.err = ri.resultPageIterator.Err() } From 89b7320784fa2a610ac96efa97f53e4f4ba55782 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:05:42 +0000 Subject: [PATCH 3/8] ai: apply changes for #509 (2 review threads) Addresses: - #4355373515 at internal/rows/rowscanner/resultPageIterator.go:29 - #4355373521 at internal/rows/arrowbased/arrowRecordIterator_test.go:236 Signed-off-by: peco-engineer-bot[bot] --- .../rows/arrowbased/arrowIPCStreamIterator.go | 43 ++++++++++++- .../arrowbased/arrowRecordIterator_test.go | 64 +++++++++++++++++++ internal/rows/rowscanner/category_test.go | 60 +++++++++++++++++ 3 files changed, 166 insertions(+), 1 deletion(-) diff --git a/internal/rows/arrowbased/arrowIPCStreamIterator.go b/internal/rows/arrowbased/arrowIPCStreamIterator.go index 2aeb0d28..5d136c94 100644 --- a/internal/rows/arrowbased/arrowIPCStreamIterator.go +++ b/internal/rows/arrowbased/arrowIPCStreamIterator.go @@ -30,6 +30,14 @@ type arrowIPCStreamIterator struct { resultPageIterator rowscanner.ResultPageIterator isFinished bool arrowSchemaBytes []byte + // err holds a page-fetch error detected while advancing the result page + // iterator. It must be surfaced through Next() rather than being masked as + // a normal end-of-results (io.EOF). See issue #507. + err error + // errReported is set once the page-fetch error has been surfaced through + // Next() so that HasNext does not re-capture it from the result page + // iterator (whose Err() keeps returning the stored error). + errReported bool } var _ rows.ArrowIPCStreamIterator = (*arrowIPCStreamIterator)(nil) @@ -40,6 +48,17 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) { return nil, io.EOF } + // If advancing the result page failed, surface that error rather than + // masking a failed page fetch as a normal end-of-results. See #507. The + // error is reported once and then the iterator is finished. + if ri.err != nil { + err := ri.err + ri.err = nil + ri.errReported = true + ri.Close() + return nil, err + } + if ri.ipcStreamIterator != nil && ri.ipcStreamIterator.HasNext() { return ri.ipcStreamIterator.Next() } @@ -59,6 +78,12 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) { // HasNext returns true if there are more streams available func (ri *arrowIPCStreamIterator) HasNext() bool { + // A pending page-fetch error still needs to be reported via Next(), so the + // iterator is not considered exhausted until that error is consumed. + if ri.err != nil { + return true + } + if ri.isFinished { return false } @@ -67,7 +92,18 @@ func (ri *arrowIPCStreamIterator) HasNext() bool { return true } - if ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() { + if ri.resultPageIterator == nil { + return false + } + + if !ri.resultPageIterator.HasNext() { + // Distinguish a failed page fetch from a fully consumed result set. + if !ri.errReported { + if err := ri.resultPageIterator.Err(); err != nil { + ri.err = err + return true + } + } return false } @@ -102,6 +138,11 @@ func (ri *arrowIPCStreamIterator) fetchNextData() error { if ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() { ri.isFinished = true + if ri.resultPageIterator != nil { + if err := ri.resultPageIterator.Err(); err != nil { + return err + } + } return io.EOF } diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index a5112d1c..bc9389b3 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -234,6 +234,70 @@ func TestArrowRecordIteratorSwallowsPageError(t *testing.T) { } t.Logf("fetch calls=%d rows=%d lastErr=%v", calls, rows, lastErr) assert.ErrorContains(t, lastErr, "connection reset") + + // Every row from the direct results and the first fetched page must be + // delivered before the error is surfaced. + expectedRows := rowscanner.CountRows(executeStatementResp.DirectResults.ResultSet.Results) + + rowscanner.CountRows(fetchResp1.Results) + assert.Equal(t, expectedRows, rows) + assert.Equal(t, 2, calls) + + // The error is reported exactly once; afterwards the iterator is finished. + assert.False(t, it.HasNext()) + rec, nextErr := it.Next() + assert.Nil(t, rec) + assert.Equal(t, io.EOF, nextErr) + assert.False(t, it.HasNext()) + assert.Equal(t, 2, calls) +} + +func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + executeStatementResp := cli_service.TExecuteStatementResp{} + loadTestData2(t, "directResultsMultipleFetch/ExecuteStatement.json", &executeStatementResp) + fetchResp1 := cli_service.TFetchResultsResp{} + loadTestData2(t, "directResultsMultipleFetch/FetchResults1.json", &fetchResp1) + + calls := 0 + failing := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + if calls == 1 { + return &fetchResp1, nil + } + return nil, fmt.Errorf("FetchResults: connection reset") + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + cfg := *config.WithDefaults() + schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema + ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) + assert.Nil(t, err) + it := NewArrowIPCStreamIterator(ctx, rpi, ipcIt, schemaBytes, cfg) + defer it.Close() + + var lastErr error + streams := 0 + for it.HasNext() { + _, err := it.Next() + if err != nil { + lastErr = err + break + } + streams++ + } + t.Logf("fetch calls=%d streams=%d lastErr=%v", calls, streams, lastErr) + assert.ErrorContains(t, lastErr, "connection reset") + assert.Greater(t, streams, 0) + + // The error is reported once; afterwards the iterator is exhausted. + assert.False(t, it.HasNext()) + _, errAfter := it.Next() + assert.ErrorIs(t, errAfter, io.EOF) } func TestArrowRecordIteratorSchema(t *testing.T) { diff --git a/internal/rows/rowscanner/category_test.go b/internal/rows/rowscanner/category_test.go index d5292da3..ee42f6ff 100644 --- a/internal/rows/rowscanner/category_test.go +++ b/internal/rows/rowscanner/category_test.go @@ -2,6 +2,7 @@ package rowscanner import ( "context" + "io" "testing" "github.com/databricks/databricks-sql-go/internal/cli_service" @@ -35,6 +36,65 @@ func TestResultPageFetchErrorCarriesCategory(t *testing.T) { assert.Equal(t, dbsqlerrint.CategoryResultSet, dbsqlerrint.CategoryFromError(err)) } +// Err() must return nil when the result set is consumed normally (io.EOF is +// not an error). +func TestResultPageIteratorErrNilOnNormalFinish(t *testing.T) { + hasMoreRows := false + calls := 0 + c := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + return &cli_service.TFetchResultsResp{ + HasMoreRows: &hasMoreRows, + Results: &cli_service.TRowSet{ + StartRowOffset: 0, + ArrowBatches: []*cli_service.TSparkArrowBatch{{RowCount: 5}}, + }, + }, nil + }, + } + + rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, false, c, + dbsqllog.WithContext("connId", "correlationId", "")) + + assert.True(t, rpf.HasNext()) + page, err := rpf.Next() + assert.Nil(t, err) + assert.NotNil(t, page) + + assert.False(t, rpf.HasNext()) + assert.Nil(t, rpf.Err()) + _, err = rpf.Next() + assert.Equal(t, io.EOF, err) + assert.Nil(t, rpf.Err()) + assert.Equal(t, 1, calls) +} + +// After a failed page fetch, subsequent HasNext() calls must keep the stored +// error rather than overwriting it with io.EOF. +func TestResultPageIteratorPreservesFetchError(t *testing.T) { + calls := 0 + c := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + return nil, errors.New("connection reset") + }, + } + + rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, false, c, + dbsqllog.WithContext("connId", "correlationId", "")) + + assert.False(t, rpf.HasNext()) + assert.ErrorContains(t, rpf.Err(), "connection reset") + + // A second HasNext() must not mask the error as a normal end-of-results. + assert.False(t, rpf.HasNext()) + assert.ErrorContains(t, rpf.Err(), "connection reset") + _, err := rpf.Next() + assert.ErrorContains(t, err, "connection reset") + assert.Equal(t, 1, calls) +} + // A page fetch that fails when the results context is already aborted must NOT // be tagged result_set_error. func TestResultPageFetchCancelNotTaggedResultSet(t *testing.T) { From 727a602b6699eecaadcd661740bb96e24a105f33 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:13:15 +0000 Subject: [PATCH 4/8] ai: apply changes for #509 (1 review thread) Addresses: - #4355423077 at internal/rows/arrowbased/arrowRecordIterator.go:112 Signed-off-by: peco-engineer-bot[bot] --- .../rows/arrowbased/arrowRecordIterator.go | 7 ++++ .../arrowbased/arrowRecordIterator_test.go | 41 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index 7c11c7a7..f1c1e82d 100644 --- a/internal/rows/arrowbased/arrowRecordIterator.go +++ b/internal/rows/arrowbased/arrowRecordIterator.go @@ -109,6 +109,13 @@ func (ri *arrowRecordIterator) Close() { } func (ri *arrowRecordIterator) checkFinished() { + // Once finished (including after an explicit Close()), never touch the + // result page iterator again: its HasNext() would issue a FetchResults + // round-trip against an operation that may already be closed. + if ri.isFinished { + return + } + noLocalData := (ri.currentBatch == nil || !ri.currentBatch.HasNext()) && (ri.batchIterator == nil || !ri.batchIterator.HasNext()) diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index bc9389b3..793f6838 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -251,6 +251,47 @@ func TestArrowRecordIteratorSwallowsPageError(t *testing.T) { assert.Equal(t, 2, calls) } +func TestArrowRecordIteratorNoFetchAfterEarlyClose(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + executeStatementResp := cli_service.TExecuteStatementResp{} + loadTestData2(t, "directResultsMultipleFetch/ExecuteStatement.json", &executeStatementResp) + + fetchCalls := 0 + closeCalls := 0 + tc := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + fetchCalls++ + return nil, fmt.Errorf("FetchResults: operation closed") + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closeCalls++ + return &cli_service.TCloseOperationResp{}, nil + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger) + cfg := *config.WithDefaults() + schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema + bi, err := NewLocalBatchIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) + assert.Nil(t, err) + it := NewArrowRecordIterator(ctx, rpi, bi, schemaBytes, cfg) + + // Close before any page beyond the direct results has been read. + it.Close() + assert.Equal(t, 1, closeCalls) + + // A closed iterator must not issue FetchResults or report a spurious error. + assert.False(t, it.HasNext()) + rec, nextErr := it.Next() + assert.Nil(t, rec) + assert.Equal(t, io.EOF, nextErr) + assert.False(t, it.HasNext()) + assert.Equal(t, 0, fetchCalls) +} + func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) { logger := dbsqllog.WithContext("connectionId", "correlationId", "") From aa2bd289e2472a189e8e8c778ee72d237086a05d Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:21:11 +0000 Subject: [PATCH 5/8] ai: apply changes for #509 (1 review thread) Addresses: - #4355457211 at internal/rows/arrowbased/arrowRecordIterator_test.go:335 Signed-off-by: peco-engineer-bot[bot] --- .../arrowbased/arrowRecordIterator_test.go | 54 +++++++++++++++++-- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 793f6838..417f032f 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -333,12 +333,60 @@ func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) { } t.Logf("fetch calls=%d streams=%d lastErr=%v", calls, streams, lastErr) assert.ErrorContains(t, lastErr, "connection reset") - assert.Greater(t, streams, 0) - // The error is reported once; afterwards the iterator is exhausted. + // Every stream from the direct results and the first fetched page must be + // delivered before the error is surfaced (one stream per local arrow batch). + expectedStreams := len(executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches) + + len(fetchResp1.Results.ArrowBatches) + assert.Equal(t, expectedStreams, streams) + assert.Equal(t, 2, calls) + + // The error is reported once; afterwards the iterator is exhausted and + // does not fetch again. assert.False(t, it.HasNext()) - _, errAfter := it.Next() + r, errAfter := it.Next() + assert.Nil(t, r) assert.ErrorIs(t, errAfter, io.EOF) + assert.False(t, it.HasNext()) + assert.Equal(t, 2, calls) +} + +func TestArrowIPCStreamIteratorNoFetchAfterEarlyClose(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + executeStatementResp := cli_service.TExecuteStatementResp{} + loadTestData2(t, "directResultsMultipleFetch/ExecuteStatement.json", &executeStatementResp) + + fetchCalls := 0 + tc := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + fetchCalls++ + return nil, fmt.Errorf("FetchResults: operation closed") + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + return &cli_service.TCloseOperationResp{}, nil + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger) + cfg := *config.WithDefaults() + schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema + ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) + assert.Nil(t, err) + it := NewArrowIPCStreamIterator(ctx, rpi, ipcIt, schemaBytes, cfg) + + // Close before any page beyond the direct results has been read. + it.Close() + + // A closed iterator must not issue FetchResults or report a spurious error. + assert.False(t, it.HasNext()) + r, nextErr := it.Next() + assert.Nil(t, r) + assert.Equal(t, io.EOF, nextErr) + assert.False(t, it.HasNext()) + assert.Equal(t, 0, fetchCalls) } func TestArrowRecordIteratorSchema(t *testing.T) { From 207e03a2caabace9d98f4ebc588f879f0286e799 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:29:34 +0000 Subject: [PATCH 6/8] ai: apply changes for #509 (1 review thread) Addresses: - #4355495859 at internal/rows/arrowbased/arrowRecordIterator.go:88 Signed-off-by: peco-engineer-bot[bot] --- .../rows/arrowbased/arrowRecordIterator.go | 24 +++++++++++--- .../arrowbased/arrowRecordIterator_test.go | 31 +++++++++++++++++++ 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index f1c1e82d..b8699b4d 100644 --- a/internal/rows/arrowbased/arrowRecordIterator.go +++ b/internal/rows/arrowbased/arrowRecordIterator.go @@ -60,11 +60,7 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) { // If advancing the result page failed, surface that error rather than // masking a failed page fetch as a normal end-of-results. See #507. The // error is reported once and then the iterator is finished. - if ri.err != nil { - err := ri.err - ri.err = nil - ri.errReported = true - ri.Close() + if err := ri.takePendingErr(); err != nil { return nil, err } @@ -82,6 +78,19 @@ func (ri *arrowRecordIterator) Next() (arrow.Record, error) { return r, err } +// takePendingErr returns the pending page-fetch error, if any, and marks it as +// reported so it is surfaced exactly once; the iterator is then finished. +func (ri *arrowRecordIterator) takePendingErr() error { + if ri.err == nil { + return nil + } + err := ri.err + ri.err = nil + ri.errReported = true + ri.Close() + return err +} + // Indicate whether there are any more records available func (ri *arrowRecordIterator) HasNext() bool { ri.checkFinished() @@ -226,6 +235,11 @@ func (ri *arrowRecordIterator) Schema() (*arrow.Schema, error) { // Try to get schema bytes if not already available if ri.arrowSchemaBytes == nil { if ri.HasNext() { + // A pending page-fetch error is reported here (once) rather than + // via getCurrentBatch, so a later Next() does not repeat it. + if err := ri.takePendingErr(); err != nil { + return nil, err + } if err := ri.getCurrentBatch(); err != nil { return nil, err } diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 417f032f..f4f86ce9 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -251,6 +251,37 @@ func TestArrowRecordIteratorSwallowsPageError(t *testing.T) { assert.Equal(t, 2, calls) } +func TestArrowRecordIteratorSchemaReportsPageErrorOnce(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + calls := 0 + failing := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + return nil, fmt.Errorf("FetchResults: connection reset") + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, 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) + defer it.Close() + + schema, err := it.Schema() + assert.Nil(t, schema) + assert.ErrorContains(t, err, "connection reset") + + // The error was reported by Schema(); Next() must not report it again. + assert.False(t, it.HasNext()) + rec, nextErr := it.Next() + assert.Nil(t, rec) + assert.Equal(t, io.EOF, nextErr) + assert.Equal(t, 1, calls) +} + func TestArrowRecordIteratorNoFetchAfterEarlyClose(t *testing.T) { logger := dbsqllog.WithContext("connectionId", "correlationId", "") From e7f5522c1063bb7572d4c59c6bfc96bd904acd1d Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:38:33 +0000 Subject: [PATCH 7/8] ai: apply changes for #509 (1 review thread) Addresses: - #4355534513 at internal/rows/arrowbased/arrowRecordIterator.go:99 Signed-off-by: peco-engineer-bot[bot] --- .../rows/arrowbased/arrowIPCStreamIterator.go | 6 +- .../rows/arrowbased/arrowRecordIterator.go | 18 ++++- .../arrowbased/arrowRecordIterator_test.go | 68 +++++++++++++++++++ 3 files changed, 88 insertions(+), 4 deletions(-) diff --git a/internal/rows/arrowbased/arrowIPCStreamIterator.go b/internal/rows/arrowbased/arrowIPCStreamIterator.go index 5d136c94..0838c812 100644 --- a/internal/rows/arrowbased/arrowIPCStreamIterator.go +++ b/internal/rows/arrowbased/arrowIPCStreamIterator.go @@ -110,8 +110,12 @@ func (ri *arrowIPCStreamIterator) HasNext() bool { return true } -// Close releases resources +// Close releases resources. It also discards any pending, not-yet-reported +// page-fetch error so a closed iterator reports HasNext()==false and +// Next()==io.EOF. func (ri *arrowIPCStreamIterator) Close() { + ri.err = nil + ri.errReported = true if ri.ipcStreamIterator != nil { ri.ipcStreamIterator.Close() ri.ipcStreamIterator = nil diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index b8699b4d..f081686b 100644 --- a/internal/rows/arrowbased/arrowRecordIterator.go +++ b/internal/rows/arrowbased/arrowRecordIterator.go @@ -99,8 +99,19 @@ func (ri *arrowRecordIterator) HasNext() bool { return !ri.isFinished || ri.err != nil } -// Free any resources associated with this iterator +// Free any resources associated with this iterator. An explicit Close also +// discards any pending, not-yet-reported page-fetch error so a closed iterator +// reports HasNext()==false and Next()==io.EOF. func (ri *arrowRecordIterator) Close() { + ri.err = nil + ri.errReported = true + ri.release() +} + +// release frees the iterator's resources and marks it finished without +// touching any pending page-fetch error, so that error can still be surfaced +// via Next() when the end of the result set is reached internally. +func (ri *arrowRecordIterator) release() { if !ri.isFinished { ri.isFinished = true if ri.currentBatch != nil { @@ -144,8 +155,9 @@ func (ri *arrowRecordIterator) checkFinished() { } if finished { - // Reached end of result set so Close - ri.Close() + // Reached end of result set so release resources, keeping any + // pending page-fetch error for Next() to report. + ri.release() } } diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index f4f86ce9..16a3d8b6 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -323,6 +323,74 @@ func TestArrowRecordIteratorNoFetchAfterEarlyClose(t *testing.T) { assert.Equal(t, 0, fetchCalls) } +func TestArrowRecordIteratorCloseDiscardsPendingPageError(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + calls := 0 + failing := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + return nil, fmt.Errorf("FetchResults: connection reset") + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + return &cli_service.TCloseOperationResp{}, nil + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + cfg := *config.WithDefaults() + it := NewArrowRecordIterator(ctx, rpi, nil, nil, cfg) + + // HasNext captures the page-fetch error as pending. + assert.True(t, it.HasNext()) + assert.Equal(t, 1, calls) + + // An explicit Close discards the pending error. + it.Close() + assert.False(t, it.HasNext()) + rec, nextErr := it.Next() + assert.Nil(t, rec) + assert.Equal(t, io.EOF, nextErr) + assert.False(t, it.HasNext()) + assert.Equal(t, 1, calls) +} + +func TestArrowIPCStreamIteratorCloseDiscardsPendingPageError(t *testing.T) { + logger := dbsqllog.WithContext("connectionId", "correlationId", "") + + calls := 0 + failing := &client.TestClient{ + FnFetchResults: func(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) { + calls++ + return nil, fmt.Errorf("FetchResults: connection reset") + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + return &cli_service.TCloseOperationResp{}, nil + }, + } + + ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") + ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + cfg := *config.WithDefaults() + it := NewArrowIPCStreamIterator(ctx, rpi, nil, nil, cfg) + + // HasNext captures the page-fetch error as pending. + assert.True(t, it.HasNext()) + assert.Equal(t, 1, calls) + + // An explicit Close discards the pending error. + it.Close() + assert.False(t, it.HasNext()) + r, nextErr := it.Next() + assert.Nil(t, r) + assert.Equal(t, io.EOF, nextErr) + assert.False(t, it.HasNext()) + assert.Equal(t, 1, calls) +} + func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) { logger := dbsqllog.WithContext("connectionId", "correlationId", "") From d6a731cadd3f12c53d5557eef420b927446fe1f6 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 02:45:36 +0000 Subject: [PATCH 8/8] ai: apply changes for #509 (1 review thread) Addresses: - #4355577382 at internal/rows/arrowbased/arrowIPCStreamIterator.go:145 Signed-off-by: peco-engineer-bot[bot] --- .../rows/arrowbased/arrowIPCStreamIterator.go | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/internal/rows/arrowbased/arrowIPCStreamIterator.go b/internal/rows/arrowbased/arrowIPCStreamIterator.go index 0838c812..d60144ea 100644 --- a/internal/rows/arrowbased/arrowIPCStreamIterator.go +++ b/internal/rows/arrowbased/arrowIPCStreamIterator.go @@ -51,11 +51,7 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) { // If advancing the result page failed, surface that error rather than // masking a failed page fetch as a normal end-of-results. See #507. The // error is reported once and then the iterator is finished. - if ri.err != nil { - err := ri.err - ri.err = nil - ri.errReported = true - ri.Close() + if err := ri.takePendingErr(); err != nil { return nil, err } @@ -76,6 +72,20 @@ func (ri *arrowIPCStreamIterator) Next() (io.Reader, error) { return nil, io.EOF } +// takePendingErr returns the pending page-fetch error, if any, and marks it as +// reported so it is surfaced exactly once; the iterator is then finished. +// Mirrors arrowRecordIterator.takePendingErr. +func (ri *arrowIPCStreamIterator) takePendingErr() error { + if ri.err == nil { + return nil + } + err := ri.err + ri.err = nil + ri.errReported = true + ri.Close() + return err +} + // HasNext returns true if there are more streams available func (ri *arrowIPCStreamIterator) HasNext() bool { // A pending page-fetch error still needs to be reported via Next(), so the @@ -141,12 +151,10 @@ func (ri *arrowIPCStreamIterator) fetchNextData() error { } if ri.resultPageIterator == nil || !ri.resultPageIterator.HasNext() { + // A page-fetch error (resultPageIterator.Err()) cannot reach here: Next() + // always calls HasNext() first, which captures it into ri.err and + // Next() surfaces it via takePendingErr() before calling fetchNextData. ri.isFinished = true - if ri.resultPageIterator != nil { - if err := ri.resultPageIterator.Err(); err != nil { - return err - } - } return io.EOF }