diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b736e9a..da6c8e8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Release History +## Unreleased + +- Avoid redundant Thrift result fetches when direct results specify `hasMoreRows=false` (databricks/databricks-sql-go#539). + ## v1.19.0 (2026-11-04) - Upgrade the SQL kernel bindings to v1.4.0. - Upgrade go-oidc to v3.21.0 (go-jose v4.1.6) and x/oauth2 to v0.36.0 to address go-jose/v3 security findings (databricks/databricks-sql-go#526). diff --git a/internal/rows/arrowbased/arrowRecordIterator_test.go b/internal/rows/arrowbased/arrowRecordIterator_test.go index 926dcc19..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, + nil, false, simpleClient, logger, ) @@ -140,7 +140,7 @@ func TestArrowRecordIterator(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + nil, false, simpleClient, logger, ) @@ -214,7 +214,7 @@ func TestArrowRecordIteratorSurfacesPageErrorOnce(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger) cfg := *config.WithDefaults() bi, err := NewLocalBatchIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema, &cfg) assert.Nil(t, err) @@ -264,7 +264,7 @@ func TestArrowRecordIteratorSchemaReportsPageErrorOnce(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger) cfg := *config.WithDefaults() // nil batch iterator and nil schema bytes: Schema() must fetch a page. it := NewArrowRecordIterator(ctx, rpi, nil, nil, cfg) @@ -303,7 +303,7 @@ func TestArrowRecordIteratorNoFetchAfterEarlyClose(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, tc, logger) cfg := *config.WithDefaults() schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema bi, err := NewLocalBatchIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) @@ -339,7 +339,7 @@ func TestArrowRecordIteratorCloseDiscardsPendingPageError(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger) cfg := *config.WithDefaults() it := NewArrowRecordIterator(ctx, rpi, nil, nil, cfg) @@ -373,7 +373,7 @@ func TestArrowIPCStreamIteratorCloseDiscardsPendingPageError(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger) cfg := *config.WithDefaults() it := NewArrowIPCStreamIterator(ctx, rpi, nil, nil, cfg) @@ -412,7 +412,7 @@ func TestArrowIPCStreamIteratorSurfacesPageError(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, failing, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, failing, logger) cfg := *config.WithDefaults() schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) @@ -469,7 +469,7 @@ func TestArrowIPCStreamIteratorNoFetchAfterEarlyClose(t *testing.T) { ctx := driverctx.NewContextWithConnId(context.Background(), "connectionId") ctx = driverctx.NewContextWithCorrelationId(ctx, "correlationId") - rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, false, tc, logger) + rpi := rowscanner.NewResultPageIterator(ctx, rowscanner.NewDelimiter(0, 7311), 5000, nil, nil, false, tc, logger) cfg := *config.WithDefaults() schemaBytes := executeStatementResp.DirectResults.ResultSetMetadata.ArrowSchema ipcIt, err := NewLocalIPCStreamIterator(ctx, executeStatementResp.DirectResults.ResultSet.Results.ArrowBatches, 0, schemaBytes, &cfg) @@ -509,7 +509,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + nil, false, simpleClient, logger, ) @@ -563,7 +563,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + nil, false, simpleClient, logger, ) @@ -605,7 +605,7 @@ func TestArrowRecordIteratorSchema(t *testing.T) { rowscanner.NewDelimiter(0, 0), 5000, nil, - false, + nil, false, failingClient, logger, ) diff --git a/internal/rows/arrowbased/arrowRows_test.go b/internal/rows/arrowbased/arrowRows_test.go index 77caff7c..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, + nil, 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..c68d6b1f 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -225,14 +225,18 @@ func NewRows( d = rowscanner.NewDelimiter(0, 0) } - // If the entire query result set fits in direct results the server closes - // the operations. + // CloudFetch may deliver all links without closing the server operation. closedOnServer := directResults != nil && directResults.CloseOperation != nil + var hasMoreRows *bool + if directResults != nil && directResults.ResultSet != nil { + hasMoreRows = directResults.ResultSet.HasMoreRows + } r.ResultPageIterator = rowscanner.NewResultPageIterator( resultsCtx, d, pageSize, opHandle, + hasMoreRows, closedOnServer, client, r.logger(), diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 07d8b8ff..78cba9a1 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -230,6 +230,7 @@ func TestRowsFetchResultPageNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, + nil, false, client, rowSet.logger(), @@ -326,6 +327,7 @@ func TestRowsFetchResultPageWithDirectResults(t *testing.T) { rowscanner.NewDelimiter(rowSet.RowScanner.Start(), rowSet.RowScanner.Count()), 1000, nil, + nil, false, client, rowSet.logger(), @@ -480,6 +482,7 @@ func TestNextNoDirectResults(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, + nil, false, client, rowSet.logger(), @@ -756,7 +759,7 @@ func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { rowSet.resultSetMetadata = metadata } rowSet.ResultPageIterator = rowscanner.NewResultPageIterator( - rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, true, testClient, rowSet.logger()) + rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, nil, true, testClient, rowSet.logger()) checkTypes := func() { t.Helper() @@ -925,6 +928,7 @@ func TestFetchResultsWithRetries(t *testing.T) { rowscanner.NewDelimiter(0, 0), 1000, nil, + nil, false, client, rowSet.logger(), diff --git a/internal/rows/rowscanner/category_test.go b/internal/rows/rowscanner/category_test.go index ee42f6ff..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, c, + rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, nil, false, c, dbsqllog.WithContext("connId", "correlationId", "")) assert.True(t, rpf.HasNext()) @@ -81,7 +81,7 @@ func TestResultPageIteratorPreservesFetchError(t *testing.T) { }, } - rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, false, c, + rpf := NewResultPageIterator(context.Background(), NewDelimiter(0, 0), 100, nil, nil, false, c, dbsqllog.WithContext("connId", "correlationId", "")) assert.False(t, rpf.HasNext()) diff --git a/internal/rows/rowscanner/resultPageIterator.go b/internal/rows/rowscanner/resultPageIterator.go index 82848abc..fed91255 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, + hasMoreRows *bool, closedOnServer bool, client cli_service.TCLIService, logger *dbsqllog.DBSQLLogger, ) ResultPageIterator { - // delimiter and hasMoreRows are used to set up the point in the paginated - // result set that this iterator starts from. return &resultPageIterator{ ctx: ctx, Delimiter: delimiter, - isFinished: closedOnServer, + isFinished: closedOnServer || (hasMoreRows != nil && !*hasMoreRows), maxPageSize: maxPageSize, opHandle: opHandle, closedOnServer: closedOnServer, @@ -130,6 +129,8 @@ func (rpf *resultPageIterator) Err() error { // Returns true if there are more pages in the result set. func (rpf *resultPageIterator) HasNext() bool { if rpf.isFinished && rpf.nextResultPage == nil { + // Preserve best-effort cleanup when direct results skip the final fetch. + rpf.Close() //nolint:errcheck,gosec // G104: match cleanup of the final fetched page // There are no more pages to load and there isn't an already fetched // page waiting to retrieved by Next(). Preserve a previously stored // fetch error rather than masking it as a normal end-of-results. diff --git a/internal/rows/sql_test.go b/internal/rows/sql_test.go new file mode 100644 index 00000000..e11c97ed --- /dev/null +++ b/internal/rows/sql_test.go @@ -0,0 +1,181 @@ +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 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() + 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}, + }}}, + }}}, + } + moreRows, noMoreRows := true, false + fetchErr := errors.New("fetch failed") + for _, tc := range []struct { + name string + hasMoreRows *bool + preclosed bool + empty bool + fetchErr error + earlyClose bool + wantFetches int + wantRows []bool + }{ + {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 + 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) { + closes++ + return nil, closeErr + }, + } + cfg := config.WithDefaults() + cfg.UseLz4Compression = false + direct := &cli_service.TSparkDirectResults{ + ResultSetMetadata: metadata, + ResultSet: &cli_service.TFetchResultsResp{ + HasMoreRows: tc.hasMoreRows, Results: page(0), + }, + } + 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, 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) + 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) + if len(got) <= 2 { + assert.Zero(t, closes, "direct rows must remain readable before cleanup") + } + 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(), tc.fetchErr) + assert.Equal(t, wantCloses, closes) + assert.NoError(t, result.Close()) + } + wantErr := tc.fetchErr + if tc.earlyClose { + wantErr = closeErr + } + assert.ErrorIs(t, result.Err(), wantErr) + assert.Equal(t, tc.wantFetches, fetches) + assert.Equal(t, wantCloses, 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 +}