diff --git a/internal/rows/arrowbased/arrowIPCStreamIterator.go b/internal/rows/arrowbased/arrowIPCStreamIterator.go index 2aeb0d28..d60144ea 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,13 @@ 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 err := ri.takePendingErr(); err != nil { + return nil, err + } + if ri.ipcStreamIterator != nil && ri.ipcStreamIterator.HasNext() { return ri.ipcStreamIterator.Next() } @@ -57,8 +72,28 @@ 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 + // iterator is not considered exhausted until that error is consumed. + if ri.err != nil { + return true + } + if ri.isFinished { return false } @@ -67,15 +102,30 @@ 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 } 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 @@ -101,6 +151,9 @@ 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 return io.EOF } diff --git a/internal/rows/arrowbased/arrowRecordIterator.go b/internal/rows/arrowbased/arrowRecordIterator.go index 820ba75f..f081686b 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,13 @@ 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 err := ri.takePendingErr(); err != nil { + return nil, err + } + // make sure we have the current batch err := ri.getCurrentBatch() if err != nil { @@ -63,14 +78,40 @@ 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() - 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 +// 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 { @@ -88,14 +129,35 @@ 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())) + // 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()) + + // 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 || 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 noMoreData && ri.resultPageIterator != nil && ri.err == nil && !ri.errReported { + ri.err = ri.resultPageIterator.Err() + } 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() } } @@ -185,6 +247,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 13891fc9..16a3d8b6 100644 --- a/internal/rows/arrowbased/arrowRecordIterator_test.go +++ b/internal/rows/arrowbased/arrowRecordIterator_test.go @@ -193,6 +193,301 @@ 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") + + // 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 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", "") + + 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 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", "") + + 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") + + // 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()) + 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) { // Test with arrowSchemaBytes available t.Run("schema with initial schema bytes", func(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) { 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 }