diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 2c3e267d56..4f02fac720 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -55,6 +55,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters. - Fixed `?` characters inside SQL comments, string literals, and quoted identifiers being incorrectly counted as parameter placeholders when `supportManyParameters=1`. `SQLInterpolator` now uses `SqlCommentParser` to locate only real placeholders. Fixes #1331. - Fixed `MetadataOperationTimeout` not being applied when metadata operations use SHOW commands. Operations like `getTables`, `getSchemas`, and `getColumns` now respect the `MetadataOperationTimeout` connection property instead of hanging indefinitely with no timeout. - Reclassify transient server errors to standard SQL states (08S01, 40001) across all Thrift error sites. This ensures UC unavailability and concurrent modification errors surface consistently for better retry handling. Note: Dashboards and branching logic keyed on legacy XXUCC or 42000 must be updated. +- Fixed client-side enforcement of `maxRows` limit. When `statement.setMaxRows()` is set, `ResultSet.next()` now returns false once the row limit is reached, even if the server returns more rows. Applies to all result types (Thrift, SEA, inline, CloudFetch). --- *Note: When making changes, please add your change under the appropriate section diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java index db436247c8..04e2cc97f8 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -82,6 +82,16 @@ enum ResultSetType { // for the lifetime of a result set. private final TelemetryCollector cachedTelemetryCollector; + // Client-side maxRows enforcement. This central check in next() is the single + // enforcement point using the full long precision from getLargeMaxRows(). + // No per-impl maxRows enforcement exists — all implementations delegate to this check. + private final long maxRowsLimit; + private long rowsReturned = 0; + private boolean truncatedByMaxRows = false; // tracks client-side truncation for cursor methods + // Flag to bypass maxRows check during getUpdateCount() internal iteration, + // which calls next() to sum affected-row counts for DML statements. + private boolean countingUpdateRows = false; + // Constructor for SEA result set public DatabricksResultSet( StatementStatus statementStatus, @@ -121,6 +131,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = parentStatement; this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement); + this.maxRowsLimit = resolveMaxRowsLimit(parentStatement); this.isClosed = false; this.wasNull = false; } @@ -142,6 +153,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = parentStatement; this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement); + this.maxRowsLimit = resolveMaxRowsLimit(parentStatement); this.isClosed = false; this.wasNull = false; this.complexDatatypeSupport = complexDatatypeSupport; @@ -189,6 +201,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = parentStatement; this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement); + this.maxRowsLimit = resolveMaxRowsLimit(parentStatement); this.isClosed = false; this.wasNull = false; } @@ -220,6 +233,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } @@ -251,6 +265,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } @@ -271,14 +286,52 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } + /** + * Advances the cursor to the next row of the result set. Returns {@code false} when no more rows + * are available or when the client-side {@code maxRows} limit has been reached. + * + *

{@code rowsReturned} tracks how many rows have been delivered to the caller through this + * ResultSet instance. It is only incremented during normal (non-DML-counting) iteration and + * assumes single-threaded access, which is the standard JDBC contract. + * + * @return {@code true} if the new current row is valid; {@code false} if there are no more rows + * @throws SQLException if the result set is closed or an error occurs + */ @Override public boolean next() throws SQLException { checkIfClosed(); + // Client-side maxRows truncation: stop before delegating to the underlying result + // implementation when the limit has been reached. This is skipped during + // getUpdateCount() internal iteration (countingUpdateRows) to avoid breaking DML + // row counting. This is the single maxRows enforcement point for all implementations. + if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) { + LOGGER.debug( + "maxRows limit ({}) reached for statement {}; returning false after {} rows", + maxRowsLimit, + statementId, + rowsReturned); + if (!truncatedByMaxRows) { + truncatedByMaxRows = true; + // Record telemetry for truncated queries so dashboards reflect the truncation + if (cachedTelemetryCollector != null) { + cachedTelemetryCollector.recordResultSetIteration( + statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), false); + } + } + return false; + } boolean hasNext = this.executionResult.next(); + // Only count rows for customer iteration, not internal DML counting + // (getUpdateCount() sets countingUpdateRows=true to iterate over affected-row counts + // without inflating the user-visible row counter). + if (hasNext && !countingUpdateRows) { + rowsReturned++; + } if (cachedTelemetryCollector != null) { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); @@ -321,6 +374,20 @@ private static TelemetryCollector resolveTelemetryCollector( return null; } + private static long resolveMaxRowsLimit(IDatabricksStatementInternal parentStatement) { + try { + if (parentStatement != null) { + // Use getLargeMaxRows() to preserve full long precision. + // getMaxRows() returns int and silently truncates values > Integer.MAX_VALUE. + return parentStatement.getLargeMaxRows(); + } + } catch (SQLException e) { + // Narrow to SQLException (the only checked exception from getLargeMaxRows). + LOGGER.warn("Error resolving maxRows limit: {}", e.getMessage()); + } + return 0; + } + @Override public boolean wasNull() throws SQLException { checkIfClosed(); @@ -662,7 +729,9 @@ public boolean isBeforeFirst() throws SQLException { @Override public boolean isAfterLast() throws SQLException { checkIfClosed(); - return executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows(); + // Account for client-side maxRows truncation + return truncatedByMaxRows + || executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows(); } @Override @@ -691,6 +760,11 @@ public boolean isFirst() throws SQLException { @Override public boolean isLast() throws SQLException { checkIfClosed(); + // Account for client-side maxRows truncation: if the next next() call would + // hit the limit, this is the last row the caller will see. + if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit) { + return true; + } if (executionResult instanceof LazyThriftResult || executionResult instanceof StreamingColumnarResult || executionResult instanceof LazyThriftInlineArrowResult @@ -731,6 +805,10 @@ public boolean last() throws SQLException { @Override public int getRow() throws SQLException { checkIfClosed(); + // JDBC spec: getRow() returns 0 when cursor is not on a valid row (after last) + if (truncatedByMaxRows) { + return 0; + } return (int) executionResult.getCurrentRow() + 1; } @@ -1958,8 +2036,13 @@ public long getUpdateCount() throws SQLException { updateCount = 0L; } else if (hasUpdateCount()) { long rowsUpdated = 0; - while (next()) { - rowsUpdated += this.getLong(AFFECTED_ROWS_COUNT); + countingUpdateRows = true; + try { + while (next()) { + rowsUpdated += this.getLong(AFFECTED_ROWS_COUNT); + } + } finally { + countingUpdateRows = false; } updateCount = rowsUpdated; } else { diff --git a/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java b/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java index 854da86566..6f6faf4102 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/LazyThriftResult.java @@ -1,6 +1,5 @@ package com.databricks.jdbc.api.impl; -import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createColumnarView; import com.databricks.jdbc.api.internal.IDatabricksSession; @@ -21,7 +20,6 @@ public class LazyThriftResult implements IExecutionResult { private long globalRowIndex; private final IDatabricksSession session; private final IDatabricksStatementInternal statement; - private final int maxRows; private boolean hasReachedEnd; private boolean isClosed; private long totalRowsFetched; @@ -42,7 +40,6 @@ public LazyThriftResult( this.currentResponse = initialResponse; this.statement = statement; this.session = session; - this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; this.globalRowIndex = -1; this.currentBatchIndex = -1; this.hasReachedEnd = false; @@ -115,13 +112,6 @@ public boolean next() throws SQLException { return false; } - // Check if we've reached the maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - hasReachedEnd = true; - return false; - } - // Move to next row in current batch currentBatchIndex++; globalRowIndex++; @@ -163,12 +153,6 @@ public boolean hasNext() { return false; } - // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - return false; - } - // Check if there are more rows in current batch if (currentBatchIndex + 1 < currentBatch.getRowCount()) { return true; diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java index 98d97a8992..2e63b32892 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java @@ -1,6 +1,5 @@ package com.databricks.jdbc.api.impl.arrow; -import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.util.ArrowUtil.createArrowByteStream; import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList; import static com.databricks.jdbc.common.util.ArrowUtil.getSerializedSchema; @@ -35,7 +34,6 @@ public class LazyThriftInlineArrowResult implements IExecutionResult { private long globalRowIndex; private final IDatabricksSession session; private final IDatabricksStatementInternal statement; - private final int maxRows; private boolean hasReachedEnd; private boolean isClosed; private long totalRowsFetched; @@ -58,7 +56,6 @@ public LazyThriftInlineArrowResult( this.currentResponse = initialResponse; this.statement = statement; this.session = session; - this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; this.globalRowIndex = -1; this.hasReachedEnd = false; this.isClosed = false; @@ -163,13 +160,6 @@ public boolean next() throws SQLException { return false; } - // Check if we've reached the maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - hasReachedEnd = true; - return false; - } - // Try to advance in current chunk if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { boolean advanced = currentChunkIterator.nextRow(); @@ -209,12 +199,6 @@ public boolean hasNext() { return false; } - // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - return false; - } - // Check if there are more rows in current chunk if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { return true; diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java index 7c9025c3f0..5f2812284d 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java @@ -1,6 +1,5 @@ package com.databricks.jdbc.api.impl.arrow; -import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS; import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList; @@ -54,7 +53,6 @@ public class StreamingInlineArrowResult implements IExecutionResult { // Metadata private List columnInfos; - private final int maxRows; // State private boolean hasReachedEnd; @@ -78,8 +76,6 @@ public StreamingInlineArrowResult( throws DatabricksSQLException { this.session = session; - this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; - this.globalRowIndex = -1; this.hasReachedEnd = false; this.isClosed = false; @@ -105,9 +101,8 @@ public StreamingInlineArrowResult( } LOGGER.debug( - "StreamingInlineArrowResult initialized - firstBatchRows={}, maxRows={}, maxBatchesInMemory={}", + "StreamingInlineArrowResult initialized - firstBatchRows={}, maxBatchesInMemory={}", currentBatch != null ? currentBatch.getRowCount() : 0, - maxRows, session.getConnectionContext().getThriftMaxBatchesInMemory()); } @@ -184,13 +179,6 @@ public boolean next() throws DatabricksSQLException { return false; } - // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - hasReachedEnd = true; - return false; - } - globalRowIndex++; // Try to move to next row in current chunk @@ -243,12 +231,6 @@ public boolean hasNext() { return false; } - // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - return false; - } - // Check current chunk if (currentChunkIterator != null && currentChunkIterator.hasNextRow()) { return true; diff --git a/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java b/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java index bf5a95d346..1a69de952e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResult.java @@ -1,6 +1,5 @@ package com.databricks.jdbc.api.impl.thrift; -import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS; import com.databricks.jdbc.api.impl.ColumnarRowView; @@ -47,9 +46,6 @@ public class StreamingColumnarResult implements IExecutionResult { private int currentBatchRowIndex; private long globalRowIndex; - // Limits - private final int maxRows; - // State private boolean hasReachedEnd; private volatile boolean isClosed; @@ -71,8 +67,6 @@ public StreamingColumnarResult( IDatabricksSession session) throws DatabricksSQLException { - this.maxRows = statement != null ? statement.getMaxRows() : DEFAULT_RESULT_ROW_LIMIT; - this.globalRowIndex = -1; this.currentBatchRowIndex = -1; this.hasReachedEnd = false; @@ -93,9 +87,8 @@ public StreamingColumnarResult( } LOGGER.debug( - "StreamingColumnarResult initialized - firstBatchRows={}, maxRows={}, maxBatchesInMemory={}", + "StreamingColumnarResult initialized - firstBatchRows={}, maxBatchesInMemory={}", currentBatch != null ? currentBatch.getRowCount() : 0, - maxRows, session.getConnectionContext().getThriftMaxBatchesInMemory()); } @@ -168,12 +161,6 @@ public boolean next() throws DatabricksSQLException { } // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - hasReachedEnd = true; - return false; - } - // Move to next row currentBatchRowIndex++; globalRowIndex++; @@ -226,12 +213,6 @@ public boolean hasNext() { return false; } - // Check maxRows limit - boolean hasRowLimit = maxRows > 0; - if (hasRowLimit && globalRowIndex + 1 >= maxRows) { - return false; - } - // Check current batch - type-safe getData() returns ColumnarRowView if (currentBatch != null) { ColumnarRowView view = currentBatch.getData(); diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java index 965a971e77..62b893c6b9 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java @@ -15,6 +15,8 @@ public interface IDatabricksStatementInternal { int getMaxRows() throws DatabricksSQLException; + long getLargeMaxRows() throws DatabricksSQLException; + void setStatementId(StatementId statementId); StatementId getStatementId(); @@ -46,6 +48,4 @@ default void markDirectResultsReceived() { default void closeServerOperation() { // no-op by default } - - long getLargeMaxRows() throws DatabricksSQLException; } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java index 15c26c4a17..9ecf0747f6 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -6,6 +6,8 @@ import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.databricks.jdbc.api.ExecutionState; @@ -1407,4 +1409,180 @@ void testTelemetryCollectorCachedOnceNotPerRow() throws SQLException { TelemetryCollectorManager.getInstance().getOrCreateCollector(verifyContext); assertNotNull(collector); } + + // --- maxRows truncation tests --- + + private DatabricksResultSet getResultSetWithMaxRows(int maxRows, IExecutionResult executionResult) + throws Exception { + IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); + when(stmt.getLargeMaxRows()).thenReturn((long) maxRows); + return new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.QUERY, + stmt, + executionResult, + mockedResultSetMetadata, + false); + } + + @Test + void testNextRespectsMaxRows() throws Exception { + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(3, mockedExecutionResult); + + assertTrue(resultSet.next()); // row 1 + assertTrue(resultSet.next()); // row 2 + assertTrue(resultSet.next()); // row 3 + assertFalse(resultSet.next()); // limit reached + assertFalse(resultSet.next()); // still false + } + + @Test + void testNextMaxRowsZeroNoLimit() throws Exception { + // maxRows=0 means no limit; all 100 next() calls should succeed + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(0, mockedExecutionResult); + + for (int i = 0; i < 100; i++) { + assertTrue(resultSet.next(), "next() should return true at iteration " + i); + } + } + + @Test + void testNextMaxRowsNullParentNoLimit() throws Exception { + // parentStatement=null means maxRowsLimit=0 (no limit) + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.QUERY, + null, // null parentStatement + mockedExecutionResult, + mockedResultSetMetadata, + false); + + for (int i = 0; i < 100; i++) { + assertTrue(resultSet.next(), "next() should return true at iteration " + i); + } + } + + @Test + void testNextMaxRowsOneEdge() throws Exception { + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(1, mockedExecutionResult); + + assertTrue(resultSet.next()); // row 1 + assertFalse(resultSet.next()); // limit reached + } + + @Test + void testNextMaxRowsDoesNotCallExecutionResultAfterLimit() throws Exception { + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(3, mockedExecutionResult); + + resultSet.next(); // row 1 + resultSet.next(); // row 2 + resultSet.next(); // row 3 + // These should NOT delegate to executionResult + resultSet.next(); + resultSet.next(); + + // executionResult.next() should have been called exactly 3 times + verify(mockedExecutionResult, times(3)).next(); + } + + @Test + void testNextMaxRowsWithEmptyResultSet() throws Exception { + // maxRows > 0 but the underlying result set is empty; next() should return false immediately + when(mockedExecutionResult.next()).thenReturn(false); + DatabricksResultSet resultSet = getResultSetWithMaxRows(5, mockedExecutionResult); + + assertFalse(resultSet.next(), "next() should return false on an empty result set"); + assertFalse( + resultSet.next(), + "next() should still return false on subsequent calls to an empty result set"); + // executionResult.next() should have been called because rowsReturned (0) < maxRows (5) + verify(mockedExecutionResult, times(2)).next(); + } + + @Test + void testNextMaxRowsGreaterThanActualRows() throws Exception { + // maxRows=10 but only 3 rows exist; result set should be naturally exhausted before the limit + when(mockedExecutionResult.next()).thenReturn(true, true, true, false); + DatabricksResultSet resultSet = getResultSetWithMaxRows(10, mockedExecutionResult); + + assertTrue(resultSet.next()); // row 1 + assertTrue(resultSet.next()); // row 2 + assertTrue(resultSet.next()); // row 3 + assertFalse(resultSet.next(), "next() should return false when data is exhausted before limit"); + } + + @Test + void testNextMaxRowsIdempotenceAfterLimit() throws Exception { + // Calling next() many times after the limit is reached should always return false + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(2, mockedExecutionResult); + + assertTrue(resultSet.next()); // row 1 + assertTrue(resultSet.next()); // row 2 (limit reached) + // All subsequent calls must consistently return false (idempotent behavior) + for (int i = 0; i < 10; i++) { + assertFalse(resultSet.next(), "next() must return false on repeated call #" + (i + 1)); + } + // executionResult.next() should have been called exactly 2 times (the limit) + verify(mockedExecutionResult, times(2)).next(); + } + + @Test + void testCursorMethodsAfterMaxRowsTruncation() throws Exception { + when(mockedExecutionResult.next()).thenReturn(true); + DatabricksResultSet resultSet = getResultSetWithMaxRows(3, mockedExecutionResult); + + // Consume all 3 allowed rows + assertTrue(resultSet.next()); + assertTrue(resultSet.next()); + assertTrue(resultSet.next()); + + // On the last allowed row, isLast() should be true + assertTrue(resultSet.isLast(), "isLast() should be true on the last allowed row"); + + // next() returns false — truncated + assertFalse(resultSet.next()); + + // After truncation: cursor is logically after last row + assertTrue(resultSet.isAfterLast(), "isAfterLast() should be true after truncation"); + assertEquals(0, resultSet.getRow(), "getRow() should return 0 when cursor is after last row"); + } + + @Test + void testGetUpdateCountBypassesMaxRows() throws Exception { + // Mock an UPDATE result set with maxRows=2 but 5 affected rows across 5 result rows. + // getUpdateCount() should sum all 5 rows (returning 5), proving the + // countingUpdateRows flag bypasses the maxRows limit during internal iteration. + InlineJsonResult mockExec = mock(InlineJsonResult.class); + when(mockExec.next()).thenReturn(true, true, true, true, true, false); + when(mockExec.getObject(0)).thenReturn(1L, 1L, 1L, 1L, 1L); + + DatabricksResultSetMetaData mockMeta = mock(DatabricksResultSetMetaData.class); + when(mockMeta.getColumnType(1)).thenReturn(Types.BIGINT); + when(mockMeta.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); + + IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); + when(stmt.getLargeMaxRows()).thenReturn(2L); + + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.UPDATE, + stmt, + mockExec, + mockMeta, + false); + + // getUpdateCount() must iterate all 5 rows despite maxRows=2 + assertEquals(5L, resultSet.getUpdateCount()); + } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java index 523a0fae6b..7d4f913516 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/LazyThriftResultTest.java @@ -29,7 +29,6 @@ public class LazyThriftResultTest { void setUp() throws SQLException { // Lenient stubbing to avoid unnecessary strictness lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); - lenient().when(statement.getMaxRows()).thenReturn(0); // No limit by default } @Test @@ -174,70 +173,6 @@ void testMultipleEmptyBatches() throws SQLException { assertEquals("final_col1", result.getObject(0)); } - @Test - void testMaxRowsLimit() throws SQLException { - when(statement.getMaxRows()).thenReturn(2); // Limit to 2 rows - - TFetchResultsResp response = - createResponseWithStringData( - Arrays.asList("row1_col1", "row1_col2"), - Arrays.asList("row2_col1", "row2_col2"), - Arrays.asList("row3_col1", "row3_col2"), - true); // hasMoreRows = true - - LazyThriftResult result = new LazyThriftResult(response, statement, session); - - // Should only get first 2 rows due to maxRows limit - assertTrue(result.next()); - assertEquals("row1_col1", result.getObject(0)); - - assertTrue(result.next()); - assertEquals("row2_col1", result.getObject(0)); - - // Should not proceed to next row due to limit - assertFalse(result.hasNext()); - assertFalse(result.next()); - - // Should not attempt to fetch more results from server - verify(databricksClient, never()).getMoreResults(any()); - } - - @Test - void testMaxRowsLimitAcrossBatches() throws SQLException { - when(statement.getMaxRows()).thenReturn(3); // Limit to 3 rows - - // First batch has 2 rows - TFetchResultsResp firstBatch = - createResponseWithStringData( - Arrays.asList("row1_col1", "row1_col2"), - Arrays.asList("row2_col1", "row2_col2"), - true); // hasMoreRows = true - - // Second batch has 2 rows - TFetchResultsResp secondBatch = - createResponseWithStringData( - Arrays.asList("row3_col1", "row3_col2"), - Arrays.asList("row4_col1", "row4_col2"), - false); // hasMoreRows = false - - when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); - - LazyThriftResult result = new LazyThriftResult(firstBatch, statement, session); - - // Consume first batch - assertTrue(result.next()); // row 1 - assertTrue(result.next()); // row 2 - - // Should get one more row from second batch then stop - assertTrue(result.next()); // row 3 - assertEquals("row3_col1", result.getObject(0)); - - // Should stop due to maxRows limit - assertFalse(result.hasNext()); - assertFalse(result.next()); - assertEquals(4, result.getTotalRowsFetched()); // 2 from first batch + 2 from second batch - } - @Test void testErrorHandlingDuringFetch() throws SQLException { TFetchResultsResp firstBatch = diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java index 3cac2c52c0..eb10b7cd1a 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java @@ -145,7 +145,6 @@ void testEmptyResultSet() throws SQLException { byte[] arrowData = createValidArrowData(1, 0); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 0, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -166,7 +165,6 @@ void testGetObjectThrowsWhenClosed() throws SQLException { byte[] arrowData = createValidArrowData(1, 1); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 1, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -189,7 +187,6 @@ void testGetObjectThrowsWhenBeforeFirstRow() throws SQLException { byte[] arrowData = createValidArrowData(1, 1); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 1, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -206,7 +203,6 @@ void testIsCompletelyFetchedWithMoreRows() throws SQLException { byte[] arrowData = createValidArrowData(1, 0); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, 0, true); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -222,7 +218,6 @@ void testIterateThroughRowsWithValidArrowData() throws SQLException { byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -251,7 +246,6 @@ void testGetObjectReturnsCorrectIntegerValue() throws SQLException { byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); when(session.getConnectionContext()).thenReturn(connectionContext); @@ -282,7 +276,6 @@ void testGetObjectWithTwoColumns() throws SQLException { byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); when(session.getConnectionContext()).thenReturn(connectionContext); @@ -309,7 +302,6 @@ void testGetObjectThrowsForColumnIndexOutOfBounds() throws SQLException { byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); // Note: session.getConnectionContext() is not stubbed here because the column index // validation happens before the connection context is accessed @@ -332,37 +324,12 @@ void testGetObjectThrowsForColumnIndexOutOfBounds() throws SQLException { assertEquals(DatabricksDriverErrorCode.INVALID_STATE.name(), beyondException.getSQLState()); } - @Test - void testMaxRowsLimitEnforced() throws SQLException { - int totalRows = 10; - int maxRows = 3; - byte[] arrowData = createValidArrowData(1, totalRows); - TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, totalRows, false); - - when(statement.getMaxRows()).thenReturn(maxRows); - when(statement.getStatementId()).thenReturn(STATEMENT_ID); - - LazyThriftInlineArrowResult result = - new LazyThriftInlineArrowResult(initialResponse, statement, session); - - // Should only be able to iterate up to maxRows - int rowsRetrieved = 0; - while (result.next()) { - rowsRetrieved++; - } - - assertEquals(maxRows, rowsRetrieved); - assertFalse(result.hasNext()); - assertEquals(maxRows - 1, result.getCurrentRow()); // 0-based index - } - @Test void testGetArrowMetadataReturnsMetadata() throws SQLException { int rowCount = 1; byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = @@ -386,7 +353,6 @@ void testFetchNextChunkFromServer() throws SQLException { // Second chunk with hasMoreRows = false TFetchResultsResp secondResponse = createFetchResultsResp(arrowData2, rowsPerChunk, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); when(session.getDatabricksClient()).thenReturn(databricksClient); when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse); @@ -418,7 +384,6 @@ void testGetRowCountReturnsCurrentChunkRowCount() throws SQLException { byte[] arrowData = createValidArrowData(1, rowCount); TFetchResultsResp initialResponse = createFetchResultsResp(arrowData, rowCount, false); - when(statement.getMaxRows()).thenReturn(0); when(statement.getStatementId()).thenReturn(STATEMENT_ID); LazyThriftInlineArrowResult result = diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java index 0280faa923..448a02de60 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java @@ -56,7 +56,6 @@ void setUp() throws Exception { connectionContext = DatabricksConnectionContextFactory.create(JDBC_URL, new Properties()); lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); lenient().when(session.getConnectionContext()).thenReturn(connectionContext); - lenient().when(statement.getMaxRows()).thenReturn(0); lenient().when(statement.getStatementId()).thenReturn(STATEMENT_ID); } @@ -147,32 +146,6 @@ void testMultiColumnAccess() throws SQLException { } } - @Test - void testMaxRowsLimit() throws SQLException { - int totalRows = 10; - int maxRows = 3; - when(statement.getMaxRows()).thenReturn(maxRows); - - byte[] arrowData = createValidArrowData(1, totalRows); - TFetchResultsResp response = createFetchResultsResp(arrowData, totalRows, false); - - StreamingInlineArrowResult result = - new StreamingInlineArrowResult(response, statement, session); - - try { - int rowsRetrieved = 0; - while (result.next()) { - rowsRetrieved++; - } - - assertEquals(maxRows, rowsRetrieved); - assertFalse(result.hasNext()); - assertEquals(maxRows - 1, result.getCurrentRow()); - } finally { - result.close(); - } - } - @Test void testColumnIndexBounds() throws SQLException { byte[] arrowData = createValidArrowData(1, 1); @@ -324,45 +297,6 @@ void testErrorDuringFetch() throws Exception { "Exception message should contain error details: " + caughtException.getMessage()); } - @Test - void testMaxRowsLimitAcrossBatches() throws SQLException, InterruptedException { - // MaxRows limit of 3, spanning across 2 batches (2 rows each) - int rowsPerChunk = 2; - when(statement.getMaxRows()).thenReturn(3); - - byte[] arrowData1 = createValidArrowData(1, rowsPerChunk); - byte[] arrowData2 = createValidArrowData(1, rowsPerChunk); - - TFetchResultsResp firstResponse = createFetchResultsResp(arrowData1, rowsPerChunk, true); - TFetchResultsResp secondResponse = createFetchResultsResp(arrowData2, rowsPerChunk, false); - - when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse); - - StreamingInlineArrowResult result = - new StreamingInlineArrowResult(firstResponse, statement, session); - - try { - // Give prefetch thread time to start - Thread.sleep(100); - - // Consume first batch (rows 0 and 1) - assertTrue(result.next()); - assertEquals(0, result.getCurrentRow()); - assertTrue(result.next()); - assertEquals(1, result.getCurrentRow()); - - // Get one row from second batch (row 2) - assertTrue(result.next()); - assertEquals(2, result.getCurrentRow()); - - // Should stop at maxRows=3 - assertFalse(result.hasNext()); - assertFalse(result.next()); - } finally { - result.close(); - } - } - @Test void testGetArrowMetadata() throws SQLException { byte[] arrowData = createValidArrowData(1, 2); diff --git a/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java index f044ec8546..8edcf22464 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/thrift/StreamingColumnarResultTest.java @@ -35,7 +35,6 @@ void setUp() throws SQLException { lenient().when(session.getDatabricksClient()).thenReturn(databricksClient); lenient().when(session.getConnectionContext()).thenReturn(connectionContext); lenient().when(connectionContext.getThriftMaxBatchesInMemory()).thenReturn(3); - lenient().when(statement.getMaxRows()).thenReturn(0); // No limit by default } @Test @@ -117,36 +116,6 @@ void testMultiBatchFetching() throws SQLException, InterruptedException { } } - @Test - void testMaxRowsLimit() throws SQLException { - when(statement.getMaxRows()).thenReturn(2); - - // Use hasMoreRows=false so the prefetch thread doesn't try to fetch - TFetchResultsResp response = - createResponseWithStringData( - Arrays.asList("row1_col1", "row1_col2"), - Arrays.asList("row2_col1", "row2_col2"), - Arrays.asList("row3_col1", "row3_col2"), - false); - - StreamingColumnarResult result = new StreamingColumnarResult(response, statement, session); - - try { - // Should only get 2 rows due to maxRows limit - assertTrue(result.next()); - assertEquals("row1_col1", result.getObject(0)); - - assertTrue(result.next()); - assertEquals("row2_col1", result.getObject(0)); - - // Should stop due to maxRows limit even though there's more data - assertFalse(result.hasNext()); - assertFalse(result.next()); - } finally { - result.close(); - } - } - @Test void testAccessAfterClose() throws SQLException { TFetchResultsResp response = @@ -286,49 +255,6 @@ void testErrorDuringFetch() throws SQLException, InterruptedException { "Exception should contain error details: " + caughtException.getMessage()); } - @Test - void testMaxRowsLimitAcrossBatches() throws SQLException, InterruptedException { - // MaxRows limit of 3, spanning across 2 batches (2 rows each) - when(statement.getMaxRows()).thenReturn(3); - - TFetchResultsResp firstBatch = - createResponseWithStringData( - Arrays.asList("row1_col1", "row1_col2"), - Arrays.asList("row2_col1", "row2_col2"), - true); // hasMoreRows = true - - TFetchResultsResp secondBatch = - createResponseWithStringData( - Arrays.asList("row3_col1", "row3_col2"), - Arrays.asList("row4_col1", "row4_col2"), - false); - - when(databricksClient.getMoreResults(statement)).thenReturn(secondBatch); - - StreamingColumnarResult result = new StreamingColumnarResult(firstBatch, statement, session); - - try { - // Give prefetch thread time to start - Thread.sleep(100); - - // Consume first batch (rows 0 and 1) - assertTrue(result.next()); - assertEquals("row1_col1", result.getObject(0)); - assertTrue(result.next()); - assertEquals("row2_col1", result.getObject(0)); - - // Get one row from second batch (row 2) - assertTrue(result.next()); - assertEquals("row3_col1", result.getObject(0)); - - // Should stop at maxRows=3 - assertFalse(result.hasNext()); - assertFalse(result.next()); - } finally { - result.close(); - } - } - @Test void testBatchesInMemoryTracking() throws SQLException { TFetchResultsResp response =