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

## Summary
Add client-side `maxRows` enforcement in `DatabricksResultSet.next()`.
When `statement.setMaxRows()` is set, `next()` returns `false` once the
row limit is reached, even if the server returns more rows. Applies
uniformly to all result types (Thrift, SEA, inline, CloudFetch).

### Problem
The server-side `row_limit`/`resultRowLimit` is not always honored — the
server may return more rows than requested. Some `IExecutionResult`
implementations (`LazyThriftResult`, `StreamingColumnarResult`,
`StreamingInlineArrowResult`) enforce maxRows internally, but
`InlineJsonResult` and `ArrowStreamResult` do not, leaving a gap.

### Fix
Single enforcement point at `DatabricksResultSet.next()`:
- `maxRowsLimit` field cached from `parentStatement.getMaxRows()` at
construction
- `rowsReturned` counter incremented on each successful `next()`
- Check `maxRowsLimit > 0 && rowsReturned >= maxRowsLimit` before
delegating to `executionResult.next()`
- `maxRows == 0` means no limit (JDBC convention)
- `getUpdateCount()` internal iteration bypasses the limit via
`countingUpdateRows` flag

Existing per-implementation maxRows enforcement retained as
defense-in-depth.


## Test plan
- [x] `testNextRespectsMaxRows` — limit=3, 4th next() returns false
- [x] `testNextMaxRowsZeroNoLimit` — no truncation with limit=0
- [x] `testNextMaxRowsNullParentNoLimit` — null parent = no limit
- [x] `testNextMaxRowsOneEdge` — limit=1 edge case
- [x] `testNextMaxRowsDoesNotCallExecutionResultAfterLimit` — verify no
unnecessary delegation
- [x] `testNextMaxRowsWithEmptyResultSet` — empty result + maxRows
- [x] `testNextMaxRowsGreaterThanActualRows` — natural exhaustion before
limit
- [x] `testNextMaxRowsIdempotenceAfterLimit` — repeated calls after
limit
- [x] `testGetUpdateCountBypassesMaxRows` — DML counting unaffected
- [x] Independent verification: 7/7 PASS (Agent C, black-box from spec)
- [x] Full suite: 62 tests pass

This pull request was AI-assisted by Isaac.

---------

Signed-off-by: Gopal Lal <gopal.lal@databricks.com>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -55,6 +55,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters.
 - Fixed `?` characters inside SQL comments, string literals, and quoted identifiers being incorrectly counted as parameter placeholders when `supportManyParameters=1`. `SQLInterpolator` now uses `SqlCommentParser` to locate only real placeholders. Fixes #1331.
 - Fixed `MetadataOperationTimeout` not being applied when metadata operations use SHOW commands. Operations like `getTables`, `getSchemas`, and `getColumns` now respect the `MetadataOperationTimeout` connection property instead of hanging indefinitely with no timeout.
 - Reclassify transient server errors to standard SQL states (08S01, 40001) across all Thrift error sites. This ensures UC unavailability and concurrent modification errors surface consistently for better retry handling. Note: Dashboards and branching logic keyed on legacy XXUCC or 42000 must be updated.
+- Fixed client-side enforcement of `maxRows` limit. When `statement.setMaxRows()` is set, `ResultSet.next()` now returns false once the row limit is reached, even if the server returns more rows. Applies to all result types (Thrift, SEA, inline, CloudFetch).
 
 ---
 *Note: When making changes, please add your change under the appropriate section

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java

@@ -82,6 +82,16 @@ enum ResultSetType {
   // for the lifetime of a result set.
   private final TelemetryCollector cachedTelemetryCollector;
 
+  // Client-side maxRows enforcement. This central check in next() is the single
+  // enforcement point using the full long precision from getLargeMaxRows().
+  // No per-impl maxRows enforcement exists — all implementations delegate to this check.
+  private final long maxRowsLimit;
+  private long rowsReturned = 0;
+  private boolean truncatedByMaxRows = false; // tracks client-side truncation for cursor methods
+  // Flag to bypass maxRows check during getUpdateCount() internal iteration,
+  // which calls next() to sum affected-row counts for DML statements.
+  private boolean countingUpdateRows = false;
+
   // Constructor for SEA result set
   public DatabricksResultSet(
       StatementStatus statementStatus,
@@ -121,6 +131,7 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = parentStatement;
     this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement);
+    this.maxRowsLimit = resolveMaxRowsLimit(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
   }
@@ -142,6 +153,7 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = parentStatement;
     this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement);
+    this.maxRowsLimit = resolveMaxRowsLimit(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
     this.complexDatatypeSupport = complexDatatypeSupport;
@@ -189,6 +201,7 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = parentStatement;
     this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement);
+    this.maxRowsLimit = resolveMaxRowsLimit(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
   }
@@ -220,6 +233,7 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = null;
     this.cachedTelemetryCollector = null;
+    this.maxRowsLimit = 0;
     this.isClosed = false;
     this.wasNull = false;
   }
@@ -251,6 +265,7 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = null;
     this.cachedTelemetryCollector = null;
+    this.maxRowsLimit = 0;
     this.isClosed = false;
     this.wasNull = false;
   }
@@ -271,14 +286,52 @@ public DatabricksResultSet(
     this.updateCount = null;
     this.parentStatement = null;
     this.cachedTelemetryCollector = null;
+    this.maxRowsLimit = 0;
     this.isClosed = false;
     this.wasNull = false;
   }
 
+  /**
+   * Advances the cursor to the next row of the result set. Returns {@code false} when no more rows
+   * are available or when the client-side {@code maxRows} limit has been reached.
+   *
+   * <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. This is the single maxRows enforcement point for all implementations.
+    if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit && !countingUpdateRows) {
+      LOGGER.debug(
+          "maxRows limit ({}) reached for statement {}; returning false after {} rows",
+          maxRowsLimit,
+          statementId,
+          rowsReturned);
+      if (!truncatedByMaxRows) {
+        truncatedByMaxRows = true;
+        // Record telemetry for truncated queries so dashboards reflect the truncation
+        if (cachedTelemetryCollector != null) {
+          cachedTelemetryCollector.recordResultSetIteration(
+              statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), false);
+        }
+      }
+      return false;
+    }
     boolean hasNext = this.executionResult.next();
+    // Only count rows for customer iteration, not internal DML counting
+    // (getUpdateCount() sets countingUpdateRows=true to iterate over affected-row counts
+    // without inflating the user-visible row counter).
+    if (hasNext && !countingUpdateRows) {
+      rowsReturned++;
+    }
     if (cachedTelemetryCollector != null) {
       cachedTelemetryCollector.recordResultSetIteration(
           statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext);
@@ -321,6 +374,20 @@ private static TelemetryCollector resolveTelemetryCollector(
     return null;
   }
 
+  private static long resolveMaxRowsLimit(IDatabricksStatementInternal parentStatement) {
+    try {
+      if (parentStatement != null) {
+        // Use getLargeMaxRows() to preserve full long precision.
+        // getMaxRows() returns int and silently truncates values > Integer.MAX_VALUE.
+        return parentStatement.getLargeMaxRows();
+      }
+    } catch (SQLException e) {
+      // Narrow to SQLException (the only checked exception from getLargeMaxRows).
+      LOGGER.warn("Error resolving maxRows limit: {}", e.getMessage());
+    }
+    return 0;
+  }
+
   @Override
   public boolean wasNull() throws SQLException {
     checkIfClosed();
@@ -662,7 +729,9 @@ public boolean isBeforeFirst() throws SQLException {
   @Override
   public boolean isAfterLast() throws SQLException {
     checkIfClosed();
-    return executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows();
+    // Account for client-side maxRows truncation
+    return truncatedByMaxRows
+        || executionResult.getCurrentRow() >= resultSetMetaData.getTotalRows();
   }
 
   @Override
@@ -691,6 +760,11 @@ public boolean isFirst() throws SQLException {
   @Override
   public boolean isLast() throws SQLException {
     checkIfClosed();
+    // Account for client-side maxRows truncation: if the next next() call would
+    // hit the limit, this is the last row the caller will see.
+    if (maxRowsLimit > 0 && rowsReturned >= maxRowsLimit) {
+      return true;
+    }
     if (executionResult instanceof LazyThriftResult
         || executionResult instanceof StreamingColumnarResult
         || executionResult instanceof LazyThriftInlineArrowResult
@@ -731,6 +805,10 @@ public boolean last() throws SQLException {
   @Override
   public int getRow() throws SQLException {
     checkIfClosed();
+    // JDBC spec: getRow() returns 0 when cursor is not on a valid row (after last)
+    if (truncatedByMaxRows) {
+      return 0;
+    }
     return (int) executionResult.getCurrentRow() + 1;
   }
 
@@ -1958,8 +2036,13 @@ public long getUpdateCount() throws SQLException {
       updateCount = 0L;
     } else if (hasUpdateCount()) {
       long rowsUpdated = 0;
-      while (next()) {
-        rowsUpdated += this.getLong(AFFECTED_ROWS_COUNT);
+      countingUpdateRows = true;
+      try {
+        while (next()) {
+          rowsUpdated += this.getLong(AFFECTED_ROWS_COUNT);
+        }
+      } finally {
+        countingUpdateRows = false;
       }
       updateCount = rowsUpdated;
     } else {

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;

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;

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<ColumnInfo> 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;