Fix client-side maxRows enforcement in ResultSet.next()

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

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,14 +285,45 @@ 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.
+   *
+   * <p>{@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. 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",
+          maxRowsLimit,
+          rowsReturned);
+      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);
@@ -312,6 +357,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 +2005,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 {

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,159 @@ 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();
+  }
+
+  @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 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());
+  }
 }