From f5dcb2a338ddd5cbbed39ce202fd15ddfc924b44 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 15 May 2026 15:53:05 +0530 Subject: [PATCH 01/10] Add client-side maxRows truncation in DatabricksResultSet.next() Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 49 ++++++++++- .../api/impl/DatabricksResultSetTest.java | 85 +++++++++++++++++++ 2 files changed, 132 insertions(+), 2 deletions(-) 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 a83956528e..2d16e7474e 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,15 @@ enum ResultSetType { // for the lifetime of a result set. private final TelemetryCollector cachedTelemetryCollector; + // Client-side maxRows enforcement. This provides a single, implementation-agnostic + // truncation point in next(). Per-implementation maxRows enforcement (e.g. in + // StreamingInlineArrowResult, StreamingColumnarResult) is retained as defense-in-depth. + private final long maxRowsLimit; + private long rowsReturned = 0; + // 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 +130,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 +152,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 +200,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 +232,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } @@ -251,6 +264,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } @@ -271,6 +285,7 @@ public DatabricksResultSet( this.updateCount = null; this.parentStatement = null; this.cachedTelemetryCollector = null; + this.maxRowsLimit = 0; this.isClosed = false; this.wasNull = false; } @@ -278,7 +293,21 @@ public DatabricksResultSet( @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. Per-implementation maxRows enforcement is retained as defense-in-depth. + if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) { + LOGGER.debug( + "maxRows limit ({}) reached; returning false from next() after {} rows", + maxRowsLimit, + rowsReturned); + return false; + } boolean hasNext = this.executionResult.next(); + if (hasNext && !countingUpdateRows) { + rowsReturned++; + } if (cachedTelemetryCollector != null) { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); @@ -312,6 +341,17 @@ private static TelemetryCollector resolveTelemetryCollector( return null; } + private static long resolveMaxRowsLimit(IDatabricksStatementInternal parentStatement) { + try { + if (parentStatement != null) { + return parentStatement.getMaxRows(); + } + } catch (Exception e) { + LOGGER.trace("Error resolving maxRows limit: {}", e.getMessage()); + } + return 0; + } + @Override public boolean wasNull() throws SQLException { checkIfClosed(); @@ -1949,8 +1989,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/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java index 15c26c4a17..1318d913e4 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,87 @@ 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.getMaxRows()).thenReturn(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(); + } } From 424da6d35539b312662611b34946e8e9879a53d8 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 15 May 2026 16:16:27 +0530 Subject: [PATCH 02/10] Address round 2 review: add edge case tests and javadoc Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 18 +++++++- .../api/impl/DatabricksResultSetTest.java | 42 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) 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 2d16e7474e..7bcf06441e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -290,13 +290,26 @@ public DatabricksResultSet( 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. Per-implementation maxRows enforcement is retained as defense-in-depth. + // row counting. Individual result-set implementations (e.g. StreamingInlineArrowResult, + // StreamingColumnarResult) also enforce maxRows independently as defense-in-depth, + // so truncation still works even if this central check is somehow bypassed. if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) { LOGGER.debug( "maxRows limit ({}) reached; returning false from next() after {} rows", @@ -305,6 +318,9 @@ public boolean next() throws SQLException { 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++; } 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 1318d913e4..2576918296 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -1492,4 +1492,46 @@ void testNextMaxRowsDoesNotCallExecutionResultAfterLimit() throws Exception { // 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(); + } } From e05f37b9f9d7a4c07f6e1e307df53ca0de4e792e Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 15 May 2026 16:41:58 +0530 Subject: [PATCH 03/10] Add changelog and getUpdateCount bypass test Signed-off-by: Gopal Lal --- NEXT_CHANGELOG.md | 1 + .../api/impl/DatabricksResultSetTest.java | 30 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 5d95473f2e..4e09494ce5 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -23,6 +23,7 @@ - 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/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java index 2576918296..9fb727a837 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -1534,4 +1534,34 @@ void testNextMaxRowsIdempotenceAfterLimit() throws Exception { // executionResult.next() should have been called exactly 2 times (the limit) verify(mockedExecutionResult, times(2)).next(); } + + @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.getMaxRows()).thenReturn(2); + + 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()); + } } From 6be661b2d2005ffe5d8a932f8f31fed7f248abbc Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Mon, 18 May 2026 11:24:09 +0530 Subject: [PATCH 04/10] Address review feedback: telemetry, getLargeMaxRows, isAfterLast, logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (Critical): Record telemetry before returning false on maxRows truncation — dashboards now reflect truncated queries. F2 (Critical): Use getLargeMaxRows() (long) instead of getMaxRows() (int) to prevent silent truncation for values > Integer.MAX_VALUE. Added getLargeMaxRows() to IDatabricksStatementInternal interface. F3 (Critical): isAfterLast() now returns true after client-side maxRows truncation via truncatedByMaxRows flag. F4 (High): Narrowed catch from Exception to SQLException in resolveMaxRowsLimit, log at WARN instead of TRACE. F9 (Medium): Include statementId in truncation debug log. F10 (Medium): Corrected defense-in-depth comment to reflect which implementations actually enforce maxRows. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .gl-implement/verify.md | 35 +++ .../verify/MaxRowsVerificationTest.java | 202 ++++++++++++++++++ .../jdbc/api/impl/DatabricksResultSet.java | 30 ++- .../IDatabricksStatementInternal.java | 2 + .../api/impl/DatabricksResultSetTest.java | 4 +- 5 files changed, 263 insertions(+), 10 deletions(-) create mode 100644 .gl-implement/verify.md create mode 100644 .gl-implement/verify/MaxRowsVerificationTest.java diff --git a/.gl-implement/verify.md b/.gl-implement/verify.md new file mode 100644 index 0000000000..d82c8ffed7 --- /dev/null +++ b/.gl-implement/verify.md @@ -0,0 +1,35 @@ +# MaxRows Verification Plan + +## Objective +Verify that DatabricksResultSet.next() enforces client-side maxRows truncation +per JDBC spec: Statement.setMaxRows() limits the number of rows returned from +any generated ResultSet. Excess rows are silently dropped. + +## Approach +Use the @VisibleForTesting constructor of DatabricksResultSet to create instances +with a mocked IExecutionResult and a mocked IDatabricksStatementInternal that +returns a specific maxRows value via getMaxRows(). + +## Scenarios + +1. **maxRows=5, server returns 100 rows**: next() should return true exactly 5 + times, then return false on the 6th call, even though the mock still has rows. + +2. **maxRows=0 (no limit)**: All rows from the underlying result should be returned. + next() should return true for every row the mock provides. + +3. **maxRows=1**: Only 1 row returned. next() returns true once, then false. + +4. **Empty result set with maxRows set**: If the underlying result has 0 rows, + next() returns false immediately regardless of maxRows setting. + +5. **getUpdateCount bypass**: Internal DML row counting via getUpdateCount() + should NOT be capped by maxRows. This is tested by verifying that the + countingUpdateRows flag bypasses the maxRows check. + +## Mocking Strategy +- Mock IExecutionResult: control next() to return true N times then false +- Mock IDatabricksStatementInternal: return desired maxRows value +- Use real StatementStatus with SUCCEEDED state +- Use real StatementId with a dummy value +- DatabricksResultSetMetaData can be mocked (not exercised by next()) diff --git a/.gl-implement/verify/MaxRowsVerificationTest.java b/.gl-implement/verify/MaxRowsVerificationTest.java new file mode 100644 index 0000000000..2c34bbfee4 --- /dev/null +++ b/.gl-implement/verify/MaxRowsVerificationTest.java @@ -0,0 +1,202 @@ +package com.databricks.jdbc.api.impl; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; +import com.databricks.jdbc.common.StatementType; +import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.model.core.StatementStatus; +import com.databricks.sdk.service.sql.StatementState; +import java.sql.SQLException; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Independent verification test for client-side maxRows truncation in DatabricksResultSet.next(). + * + *

Per JDBC spec for Statement.setMaxRows(): "Sets the limit for the maximum number of rows that + * any ResultSet object generated by this Statement object can contain to the given number. If the + * limit is exceeded, the excess rows are silently dropped." + * + *

This test derives its expectations purely from the spec, not from the implementation. + */ +class MaxRowsVerificationTest { + + /** + * Creates a mock IExecutionResult that returns true from next() for exactly {@code totalRows} + * calls, then returns false. + */ + private IExecutionResult createMockResult(int totalRows) throws SQLException { + IExecutionResult result = mock(IExecutionResult.class); + AtomicInteger counter = new AtomicInteger(0); + when(result.next()) + .thenAnswer( + invocation -> { + return counter.getAndIncrement() < totalRows; + }); + when(result.hasNext()) + .thenAnswer( + invocation -> { + return counter.get() < totalRows; + }); + when(result.getRowCount()).thenReturn((long) totalRows); + when(result.getChunkCount()).thenReturn(1L); + return result; + } + + /** + * Creates a mock parent statement that returns the given maxRows value from getMaxRows(). + */ + private IDatabricksStatementInternal createMockStatement(int maxRows) throws Exception { + IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); + when(stmt.getMaxRows()).thenReturn(maxRows); + return stmt; + } + + /** + * Builds a DatabricksResultSet using the @VisibleForTesting constructor. + */ + private DatabricksResultSet buildResultSet( + IDatabricksStatementInternal parentStatement, + IExecutionResult executionResult) + throws Exception { + StatementStatus status = new StatementStatus().setState(StatementState.SUCCEEDED); + StatementId statementId = new StatementId("test-statement-id"); + DatabricksResultSetMetaData metaData = mock(DatabricksResultSetMetaData.class); + when(metaData.getChunkCount()).thenReturn(1L); + + return new DatabricksResultSet( + status, + statementId, + StatementType.QUERY, + parentStatement, + executionResult, + metaData, + /* complexDatatypeSupport= */ false); + } + + @Test + @DisplayName("maxRows=5, server returns 100 rows -> only 5 returned") + void testMaxRows5ServerReturns100() throws Exception { + IExecutionResult mockResult = createMockResult(100); + IDatabricksStatementInternal mockStmt = createMockStatement(5); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + int count = 0; + while (rs.next()) { + count++; + // Safety: should never exceed maxRows + assertTrue(count <= 5, "Should not return more than maxRows=5 rows"); + } + + assertEquals(5, count, "Expected exactly 5 rows with maxRows=5"); + } + + @Test + @DisplayName("maxRows=0 (no limit) -> all rows returned") + void testMaxRowsZeroNoLimit() throws Exception { + int serverRows = 42; + IExecutionResult mockResult = createMockResult(serverRows); + IDatabricksStatementInternal mockStmt = createMockStatement(0); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + int count = 0; + while (rs.next()) { + count++; + } + + assertEquals(serverRows, count, "maxRows=0 should return all rows from the server"); + } + + @Test + @DisplayName("maxRows=1 -> only 1 row returned") + void testMaxRows1() throws Exception { + IExecutionResult mockResult = createMockResult(50); + IDatabricksStatementInternal mockStmt = createMockStatement(1); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + int count = 0; + while (rs.next()) { + count++; + } + + assertEquals(1, count, "Expected exactly 1 row with maxRows=1"); + } + + @Test + @DisplayName("Empty result set with maxRows set -> returns false immediately") + void testEmptyResultSetWithMaxRows() throws Exception { + IExecutionResult mockResult = createMockResult(0); + IDatabricksStatementInternal mockStmt = createMockStatement(10); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + assertFalse(rs.next(), "Empty result set should return false on first next() call"); + } + + @Test + @DisplayName("maxRows does not affect subsequent next() calls after limit - stays false") + void testMaxRowsStaysFalseAfterLimit() throws Exception { + IExecutionResult mockResult = createMockResult(100); + IDatabricksStatementInternal mockStmt = createMockStatement(3); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + // Consume maxRows rows + for (int i = 0; i < 3; i++) { + assertTrue(rs.next(), "Should return true for row " + (i + 1)); + } + + // All subsequent calls should return false + assertFalse(rs.next(), "4th call should return false (maxRows=3)"); + assertFalse(rs.next(), "5th call should still return false"); + assertFalse(rs.next(), "6th call should still return false"); + } + + @Test + @DisplayName("maxRows larger than actual rows -> all rows returned") + void testMaxRowsLargerThanActualRows() throws Exception { + int serverRows = 3; + IExecutionResult mockResult = createMockResult(serverRows); + IDatabricksStatementInternal mockStmt = createMockStatement(1000); + DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); + + int count = 0; + while (rs.next()) { + count++; + } + + assertEquals( + serverRows, count, "When maxRows > actual rows, all actual rows should be returned"); + } + + @Test + @DisplayName("maxRows with null parent statement -> behaves as no limit (maxRows=0)") + void testNullParentStatementDefaultsToNoLimit() throws Exception { + int serverRows = 25; + IExecutionResult mockResult = createMockResult(serverRows); + + StatementStatus status = new StatementStatus().setState(StatementState.SUCCEEDED); + StatementId statementId = new StatementId("test-null-parent"); + DatabricksResultSetMetaData metaData = mock(DatabricksResultSetMetaData.class); + when(metaData.getChunkCount()).thenReturn(1L); + + DatabricksResultSet rs = + new DatabricksResultSet( + status, + statementId, + StatementType.QUERY, + /* parentStatement= */ null, + mockResult, + metaData, + false); + + int count = 0; + while (rs.next()) { + count++; + } + + assertEquals( + serverRows, count, "Null parent statement should default to no row limit (maxRows=0)"); + } +} 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 7bcf06441e..810b6c941e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -87,6 +87,8 @@ enum ResultSetType { // StreamingInlineArrowResult, StreamingColumnarResult) is retained as defense-in-depth. private final long maxRowsLimit; private long rowsReturned = 0; + private boolean truncatedByMaxRows = + false; // F3: 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; @@ -307,14 +309,21 @@ public boolean next() throws SQLException { // 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. Individual result-set implementations (e.g. StreamingInlineArrowResult, - // StreamingColumnarResult) also enforce maxRows independently as defense-in-depth, - // so truncation still works even if this central check is somehow bypassed. + // row counting. Some streaming implementations (StreamingInlineArrowResult, + // StreamingColumnarResult, LazyThriftResult) also enforce maxRows internally; + // InlineJsonResult and ArrowStreamResult do not — this central check covers them. if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) { LOGGER.debug( - "maxRows limit ({}) reached; returning false from next() after {} rows", + "maxRows limit ({}) reached for statement {}; returning false after {} rows", maxRowsLimit, + statementId, rowsReturned); + truncatedByMaxRows = true; + // F1: 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(); @@ -360,10 +369,13 @@ private static TelemetryCollector resolveTelemetryCollector( private static long resolveMaxRowsLimit(IDatabricksStatementInternal parentStatement) { try { if (parentStatement != null) { - return parentStatement.getMaxRows(); + // F2: Use getLargeMaxRows() to preserve full long precision. + // getMaxRows() returns int and silently truncates values > Integer.MAX_VALUE. + return parentStatement.getLargeMaxRows(); } - } catch (Exception e) { - LOGGER.trace("Error resolving maxRows limit: {}", e.getMessage()); + } catch (SQLException e) { + // F4: Narrow to SQLException (the only checked exception from getLargeMaxRows). + LOGGER.warn("Error resolving maxRows limit: {}", e.getMessage()); } return 0; } @@ -709,7 +721,9 @@ public boolean isBeforeFirst() throws SQLException { @Override public boolean isAfterLast() throws SQLException { checkIfClosed(); - return executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows(); + // F3: account for client-side maxRows truncation + return truncatedByMaxRows + || executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows(); } @Override 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 ccb6f035e6..0f19f0ae78 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(); 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 9fb727a837..2742b6bbbc 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -1415,7 +1415,7 @@ void testTelemetryCollectorCachedOnceNotPerRow() throws SQLException { private DatabricksResultSet getResultSetWithMaxRows(int maxRows, IExecutionResult executionResult) throws Exception { IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); - when(stmt.getMaxRows()).thenReturn(maxRows); + when(stmt.getLargeMaxRows()).thenReturn((long) maxRows); return new DatabricksResultSet( new StatementStatus().setState(StatementState.SUCCEEDED), STATEMENT_ID, @@ -1549,7 +1549,7 @@ void testGetUpdateCountBypassesMaxRows() throws Exception { when(mockMeta.getColumnNameIndex(AFFECTED_ROWS_COUNT)).thenReturn(1); IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); - when(stmt.getMaxRows()).thenReturn(2); + when(stmt.getLargeMaxRows()).thenReturn(2L); DatabricksResultSet resultSet = new DatabricksResultSet( From 3ee53b0b3d9ee0f92adbc5098540a01092de2993 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 19 May 2026 12:33:39 +0530 Subject: [PATCH 05/10] Address review: fix isLast() for maxRows, clarify int/long precision, remove gl-implement artifacts - Fix isLast() to return true when maxRows limit is reached - Update defense-in-depth comment to clarify int vs long precision limitation - Remove .gl-implement/ directory (dead workflow artifacts) Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .gl-implement/verify.md | 35 --- .../verify/MaxRowsVerificationTest.java | 202 ------------------ .../jdbc/api/impl/DatabricksResultSet.java | 14 +- 3 files changed, 11 insertions(+), 240 deletions(-) delete mode 100644 .gl-implement/verify.md delete mode 100644 .gl-implement/verify/MaxRowsVerificationTest.java diff --git a/.gl-implement/verify.md b/.gl-implement/verify.md deleted file mode 100644 index d82c8ffed7..0000000000 --- a/.gl-implement/verify.md +++ /dev/null @@ -1,35 +0,0 @@ -# MaxRows Verification Plan - -## Objective -Verify that DatabricksResultSet.next() enforces client-side maxRows truncation -per JDBC spec: Statement.setMaxRows() limits the number of rows returned from -any generated ResultSet. Excess rows are silently dropped. - -## Approach -Use the @VisibleForTesting constructor of DatabricksResultSet to create instances -with a mocked IExecutionResult and a mocked IDatabricksStatementInternal that -returns a specific maxRows value via getMaxRows(). - -## Scenarios - -1. **maxRows=5, server returns 100 rows**: next() should return true exactly 5 - times, then return false on the 6th call, even though the mock still has rows. - -2. **maxRows=0 (no limit)**: All rows from the underlying result should be returned. - next() should return true for every row the mock provides. - -3. **maxRows=1**: Only 1 row returned. next() returns true once, then false. - -4. **Empty result set with maxRows set**: If the underlying result has 0 rows, - next() returns false immediately regardless of maxRows setting. - -5. **getUpdateCount bypass**: Internal DML row counting via getUpdateCount() - should NOT be capped by maxRows. This is tested by verifying that the - countingUpdateRows flag bypasses the maxRows check. - -## Mocking Strategy -- Mock IExecutionResult: control next() to return true N times then false -- Mock IDatabricksStatementInternal: return desired maxRows value -- Use real StatementStatus with SUCCEEDED state -- Use real StatementId with a dummy value -- DatabricksResultSetMetaData can be mocked (not exercised by next()) diff --git a/.gl-implement/verify/MaxRowsVerificationTest.java b/.gl-implement/verify/MaxRowsVerificationTest.java deleted file mode 100644 index 2c34bbfee4..0000000000 --- a/.gl-implement/verify/MaxRowsVerificationTest.java +++ /dev/null @@ -1,202 +0,0 @@ -package com.databricks.jdbc.api.impl; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; -import com.databricks.jdbc.common.StatementType; -import com.databricks.jdbc.dbclient.impl.common.StatementId; -import com.databricks.jdbc.model.core.StatementStatus; -import com.databricks.sdk.service.sql.StatementState; -import java.sql.SQLException; -import java.util.concurrent.atomic.AtomicInteger; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Independent verification test for client-side maxRows truncation in DatabricksResultSet.next(). - * - *

Per JDBC spec for Statement.setMaxRows(): "Sets the limit for the maximum number of rows that - * any ResultSet object generated by this Statement object can contain to the given number. If the - * limit is exceeded, the excess rows are silently dropped." - * - *

This test derives its expectations purely from the spec, not from the implementation. - */ -class MaxRowsVerificationTest { - - /** - * Creates a mock IExecutionResult that returns true from next() for exactly {@code totalRows} - * calls, then returns false. - */ - private IExecutionResult createMockResult(int totalRows) throws SQLException { - IExecutionResult result = mock(IExecutionResult.class); - AtomicInteger counter = new AtomicInteger(0); - when(result.next()) - .thenAnswer( - invocation -> { - return counter.getAndIncrement() < totalRows; - }); - when(result.hasNext()) - .thenAnswer( - invocation -> { - return counter.get() < totalRows; - }); - when(result.getRowCount()).thenReturn((long) totalRows); - when(result.getChunkCount()).thenReturn(1L); - return result; - } - - /** - * Creates a mock parent statement that returns the given maxRows value from getMaxRows(). - */ - private IDatabricksStatementInternal createMockStatement(int maxRows) throws Exception { - IDatabricksStatementInternal stmt = mock(IDatabricksStatementInternal.class); - when(stmt.getMaxRows()).thenReturn(maxRows); - return stmt; - } - - /** - * Builds a DatabricksResultSet using the @VisibleForTesting constructor. - */ - private DatabricksResultSet buildResultSet( - IDatabricksStatementInternal parentStatement, - IExecutionResult executionResult) - throws Exception { - StatementStatus status = new StatementStatus().setState(StatementState.SUCCEEDED); - StatementId statementId = new StatementId("test-statement-id"); - DatabricksResultSetMetaData metaData = mock(DatabricksResultSetMetaData.class); - when(metaData.getChunkCount()).thenReturn(1L); - - return new DatabricksResultSet( - status, - statementId, - StatementType.QUERY, - parentStatement, - executionResult, - metaData, - /* complexDatatypeSupport= */ false); - } - - @Test - @DisplayName("maxRows=5, server returns 100 rows -> only 5 returned") - void testMaxRows5ServerReturns100() throws Exception { - IExecutionResult mockResult = createMockResult(100); - IDatabricksStatementInternal mockStmt = createMockStatement(5); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - int count = 0; - while (rs.next()) { - count++; - // Safety: should never exceed maxRows - assertTrue(count <= 5, "Should not return more than maxRows=5 rows"); - } - - assertEquals(5, count, "Expected exactly 5 rows with maxRows=5"); - } - - @Test - @DisplayName("maxRows=0 (no limit) -> all rows returned") - void testMaxRowsZeroNoLimit() throws Exception { - int serverRows = 42; - IExecutionResult mockResult = createMockResult(serverRows); - IDatabricksStatementInternal mockStmt = createMockStatement(0); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - int count = 0; - while (rs.next()) { - count++; - } - - assertEquals(serverRows, count, "maxRows=0 should return all rows from the server"); - } - - @Test - @DisplayName("maxRows=1 -> only 1 row returned") - void testMaxRows1() throws Exception { - IExecutionResult mockResult = createMockResult(50); - IDatabricksStatementInternal mockStmt = createMockStatement(1); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - int count = 0; - while (rs.next()) { - count++; - } - - assertEquals(1, count, "Expected exactly 1 row with maxRows=1"); - } - - @Test - @DisplayName("Empty result set with maxRows set -> returns false immediately") - void testEmptyResultSetWithMaxRows() throws Exception { - IExecutionResult mockResult = createMockResult(0); - IDatabricksStatementInternal mockStmt = createMockStatement(10); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - assertFalse(rs.next(), "Empty result set should return false on first next() call"); - } - - @Test - @DisplayName("maxRows does not affect subsequent next() calls after limit - stays false") - void testMaxRowsStaysFalseAfterLimit() throws Exception { - IExecutionResult mockResult = createMockResult(100); - IDatabricksStatementInternal mockStmt = createMockStatement(3); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - // Consume maxRows rows - for (int i = 0; i < 3; i++) { - assertTrue(rs.next(), "Should return true for row " + (i + 1)); - } - - // All subsequent calls should return false - assertFalse(rs.next(), "4th call should return false (maxRows=3)"); - assertFalse(rs.next(), "5th call should still return false"); - assertFalse(rs.next(), "6th call should still return false"); - } - - @Test - @DisplayName("maxRows larger than actual rows -> all rows returned") - void testMaxRowsLargerThanActualRows() throws Exception { - int serverRows = 3; - IExecutionResult mockResult = createMockResult(serverRows); - IDatabricksStatementInternal mockStmt = createMockStatement(1000); - DatabricksResultSet rs = buildResultSet(mockStmt, mockResult); - - int count = 0; - while (rs.next()) { - count++; - } - - assertEquals( - serverRows, count, "When maxRows > actual rows, all actual rows should be returned"); - } - - @Test - @DisplayName("maxRows with null parent statement -> behaves as no limit (maxRows=0)") - void testNullParentStatementDefaultsToNoLimit() throws Exception { - int serverRows = 25; - IExecutionResult mockResult = createMockResult(serverRows); - - StatementStatus status = new StatementStatus().setState(StatementState.SUCCEEDED); - StatementId statementId = new StatementId("test-null-parent"); - DatabricksResultSetMetaData metaData = mock(DatabricksResultSetMetaData.class); - when(metaData.getChunkCount()).thenReturn(1L); - - DatabricksResultSet rs = - new DatabricksResultSet( - status, - statementId, - StatementType.QUERY, - /* parentStatement= */ null, - mockResult, - metaData, - false); - - int count = 0; - while (rs.next()) { - count++; - } - - assertEquals( - serverRows, count, "Null parent statement should default to no row limit (maxRows=0)"); - } -} 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 810b6c941e..0421dd701c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -82,9 +82,12 @@ enum ResultSetType { // for the lifetime of a result set. private final TelemetryCollector cachedTelemetryCollector; - // Client-side maxRows enforcement. This provides a single, implementation-agnostic - // truncation point in next(). Per-implementation maxRows enforcement (e.g. in - // StreamingInlineArrowResult, StreamingColumnarResult) is retained as defense-in-depth. + // Client-side maxRows enforcement. This central check in next() is the primary + // enforcement point using the full long precision from getLargeMaxRows(). Some + // streaming implementations (StreamingInlineArrowResult, StreamingColumnarResult, + // LazyThriftResult) also enforce maxRows internally using int precision — those + // per-impl checks are secondary and only correct for values within int range. + // For values > Integer.MAX_VALUE, only this central check is reliable. private final long maxRowsLimit; private long rowsReturned = 0; private boolean truncatedByMaxRows = @@ -752,6 +755,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 From b07ae45372989fcb1ea73e6782cfe4b2402f7b69 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 19 May 2026 14:09:43 +0530 Subject: [PATCH 06/10] Fix telemetry re-firing on repeated next() after maxRows and strip review tags - Guard truncation telemetry with if (!truncatedByMaxRows) so it fires once - Remove stale // F1: / F2: / F3: / F4: review-iteration prefixes from comments Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) 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 0421dd701c..2fa1ca87ee 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -90,8 +90,7 @@ enum ResultSetType { // For values > Integer.MAX_VALUE, only this central check is reliable. private final long maxRowsLimit; private long rowsReturned = 0; - private boolean truncatedByMaxRows = - false; // F3: tracks client-side truncation for cursor methods + 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; @@ -321,11 +320,13 @@ public boolean next() throws SQLException { maxRowsLimit, statementId, rowsReturned); - truncatedByMaxRows = true; - // F1: record telemetry for truncated queries so dashboards reflect the truncation - if (cachedTelemetryCollector != null) { - cachedTelemetryCollector.recordResultSetIteration( - statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), false); + 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; } @@ -372,12 +373,12 @@ private static TelemetryCollector resolveTelemetryCollector( private static long resolveMaxRowsLimit(IDatabricksStatementInternal parentStatement) { try { if (parentStatement != null) { - // F2: Use getLargeMaxRows() to preserve full long precision. + // Use getLargeMaxRows() to preserve full long precision. // getMaxRows() returns int and silently truncates values > Integer.MAX_VALUE. return parentStatement.getLargeMaxRows(); } } catch (SQLException e) { - // F4: Narrow to SQLException (the only checked exception from getLargeMaxRows). + // Narrow to SQLException (the only checked exception from getLargeMaxRows). LOGGER.warn("Error resolving maxRows limit: {}", e.getMessage()); } return 0; @@ -724,7 +725,7 @@ public boolean isBeforeFirst() throws SQLException { @Override public boolean isAfterLast() throws SQLException { checkIfClosed(); - // F3: account for client-side maxRows truncation + // Account for client-side maxRows truncation return truncatedByMaxRows || executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows(); } From 60e4d07e9672ffe4a2cbdc4d6ecfa488934fdd11 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 19 May 2026 14:30:03 +0530 Subject: [PATCH 07/10] Fix getRow() after maxRows truncation, add cursor method tests, complete impl list in comments - getRow() now returns 0 when cursor is past last row due to maxRows truncation - Add test verifying isLast(), isAfterLast(), getRow() after truncation - Add LazyThriftInlineArrowResult to per-impl enforcement list in comments Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 12 ++++++++--- .../api/impl/DatabricksResultSetTest.java | 21 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) 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 2fa1ca87ee..d6205eda73 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -85,7 +85,8 @@ enum ResultSetType { // Client-side maxRows enforcement. This central check in next() is the primary // enforcement point using the full long precision from getLargeMaxRows(). Some // streaming implementations (StreamingInlineArrowResult, StreamingColumnarResult, - // LazyThriftResult) also enforce maxRows internally using int precision — those + // LazyThriftResult, LazyThriftInlineArrowResult) also enforce maxRows internally using int + // precision — those // per-impl checks are secondary and only correct for values within int range. // For values > Integer.MAX_VALUE, only this central check is reliable. private final long maxRowsLimit; @@ -312,8 +313,9 @@ public boolean next() throws SQLException { // implementation when the limit has been reached. This is skipped during // getUpdateCount() internal iteration (countingUpdateRows) to avoid breaking DML // row counting. Some streaming implementations (StreamingInlineArrowResult, - // StreamingColumnarResult, LazyThriftResult) also enforce maxRows internally; - // InlineJsonResult and ArrowStreamResult do not — this central check covers them. + // StreamingColumnarResult, LazyThriftResult, LazyThriftInlineArrowResult) also + // enforce maxRows internally; InlineJsonResult and ArrowStreamResult do not — + // this central check covers all paths. if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) { LOGGER.debug( "maxRows limit ({}) reached for statement {}; returning false after {} rows", @@ -801,6 +803,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; } 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 2742b6bbbc..9ecf0747f6 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java @@ -1535,6 +1535,27 @@ void testNextMaxRowsIdempotenceAfterLimit() throws Exception { 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. From 2976c11882163a85b0a74843729f2fbe98e98252 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 15:20:22 +0530 Subject: [PATCH 08/10] =?UTF-8?q?Remove=20per-impl=20int=20maxRows=20enfor?= =?UTF-8?q?cement=20=E2=80=94=20central=20long=20check=20is=20sole=20enfor?= =?UTF-8?q?cement=20point?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-impl maxRows checks in StreamingInlineArrowResult, StreamingColumnarResult, LazyThriftResult, and LazyThriftInlineArrowResult used int precision from getMaxRows(), which silently truncates values > Integer.MAX_VALUE via int cast. For setLargeMaxRows(5_000_000_000L), the int cast becomes 705032704 and per-impl truncates at ~705M rows before the central long check can see the 5B value. Remove per-impl enforcement entirely — the central check in DatabricksResultSet.next() using getLargeMaxRows() (long) is now the single enforcement point for all 6 IExecutionResult implementations. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 15 ++++--------- .../jdbc/api/impl/LazyThriftResult.java | 16 -------------- .../arrow/LazyThriftInlineArrowResult.java | 16 -------------- .../arrow/StreamingInlineArrowResult.java | 20 +----------------- .../impl/thrift/StreamingColumnarResult.java | 21 +------------------ 5 files changed, 6 insertions(+), 82 deletions(-) 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 5acd5b23f8..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,13 +82,9 @@ enum ResultSetType { // for the lifetime of a result set. private final TelemetryCollector cachedTelemetryCollector; - // Client-side maxRows enforcement. This central check in next() is the primary - // enforcement point using the full long precision from getLargeMaxRows(). Some - // streaming implementations (StreamingInlineArrowResult, StreamingColumnarResult, - // LazyThriftResult, LazyThriftInlineArrowResult) also enforce maxRows internally using int - // precision — those - // per-impl checks are secondary and only correct for values within int range. - // For values > Integer.MAX_VALUE, only this central check is reliable. + // 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 @@ -312,10 +308,7 @@ public boolean next() throws SQLException { // 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. Some streaming implementations (StreamingInlineArrowResult, - // StreamingColumnarResult, LazyThriftResult, LazyThriftInlineArrowResult) also - // enforce maxRows internally; InlineJsonResult and ArrowStreamResult do not — - // this central check covers all paths. + // 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", 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(); From c88d16db1e0cbbe3b9da446d0c6ce47e6c1d774a Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 16:22:26 +0530 Subject: [PATCH 09/10] Remove duplicate getLargeMaxRows() from merge with main The method was added by this PR at line 18 and also came in via the merge with main (from PR #1444) at line 52. Remove the duplicate. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/internal/IDatabricksStatementInternal.java | 2 -- 1 file changed, 2 deletions(-) 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 059d02cdd7..62b893c6b9 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java @@ -48,6 +48,4 @@ default void markDirectResultsReceived() { default void closeServerOperation() { // no-op by default } - - long getLargeMaxRows() throws DatabricksSQLException; } From e23ef7ef09152de094c0106d831e60181f93e83b Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 16:41:25 +0530 Subject: [PATCH 10/10] Remove per-impl maxRows tests and stale getMaxRows() stubs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-impl maxRows enforcement was removed in the previous commit — the central check in DatabricksResultSet.next() is the single enforcement point. Remove the now-stale getMaxRows() stubs from all 4 streaming result test files, and delete the per-impl maxRows test methods that tested the removed enforcement (testMaxRowsLimit, testMaxRowsLimitEnforced, testMaxRowsLimitAcrossBatches). Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/LazyThriftResultTest.java | 65 ---------------- .../LazyThriftInlineArrowResultTest.java | 35 --------- .../arrow/StreamingInlineArrowResultTest.java | 66 ----------------- .../thrift/StreamingColumnarResultTest.java | 74 ------------------- 4 files changed, 240 deletions(-) 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 =