Proactive server-side statement close when results consumed (#1444)

## Summary
Proactively closes the server-side operation when all results have been
consumed or the ResultSet is explicitly closed, while keeping the
client-side Statement open for reuse. This reduces server resource usage
and warehouse cost for slow consumers who don't explicitly close
statements.

### When server operation is closed:
1. `next()` returns false — all rows consumed
2. `ResultSet.close()` called explicitly

### What stays open:
- Client-side Statement — fully reusable for re-execution
- `Statement.close()` becomes a no-op for the server RPC (operation
already closed), just cleans up client state

### New flag: `serverOperationClosed`
- Set after proactive `closeStatement` RPC
- Prevents duplicate RPCs when `Statement.close()` is called later
- Reset on re-execution (`resetForNewExecution()`)
- Skipped when `directResultsReceived` is already true (server already
closed)
- Best-effort: server errors during proactive close are swallowed
(logged at DEBUG)

## Test plan
- [x] `testCloseServerOperation_closesServerAndSkipsRpcOnStatementClose`
— proactive close sends RPC, Statement.close() skips it
- [x] `testCloseServerOperation_idempotent` — double call sends only one
RPC
- [x] `testCloseServerOperation_skippedForDirectResults` — no-op when
server already closed
- [x] `testCloseServerOperation_skippedWhenNoStatementId` — no-op when
no statement ID
- [x] `testCloseServerOperation_resetsAfterFlagClear` — flag prevents
all subsequent RPCs
- [x] `testCloseServerOperation_errorSwallowed` — server errors don't
propagate
- [x] Full suite: 193 tests pass (DatabricksStatementTest +
DatabricksResultSetTest + DatabricksConnectionTest)

This pull request was AI-assisted by Isaac.

---------

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

Author: gopalldb

NEXT_CHANGELOG.md

@@ -43,6 +43,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters.
 - `EnableGeoSpatialSupport` no longer requires `EnableComplexDatatypeSupport=1`. Geospatial types (GEOMETRY, GEOGRAPHY) can now be enabled independently of complex type support (ARRAY, MAP, STRUCT).
 - Arrow schema deserialization failures (Thrift metadata path) now surface a dedicated driver error code `ARROW_SCHEMA_PARSING_ERROR` (vendor code `22000`) and a proper SQLSTATE `22000` (Data Exception) on the thrown `SQLException`, instead of the generic `RESULT_SET_ERROR` (1004) and the enum name as SQLSTATE. The exception message is unchanged.
 - When a Statement is re-executed, the previous server-side operation is now explicitly closed before starting the new execution, preventing orphaned server-side operations when Statements are reused.
+- Server-side operations are now closed proactively when `ResultSet.close()` is called, improving resource utilization. The client-side Statement remains open and reusable for re-execution. As a result, `getExecutionResult()` after result consumption returns the cached ResultSet instead of making a server RPC.
 
 ### Fixed
 - Fixed `DatabaseMetaData.getTables()` in Thrift mode returning rows when called with an empty `types` array. Per JDBC spec, empty types means "no types selected" and now correctly returns zero rows (matching SEA mode).

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

@@ -288,13 +288,22 @@ public boolean next() throws SQLException {
 
   @Override
   public void close() throws DatabricksSQLException {
+    // Proactively close server operation when ResultSet is closed explicitly.
+    closeServerOperation();
     isClosed = true;
     this.executionResult.close();
     if (parentStatement != null) {
       parentStatement.handleResultSetClose(this);
     }
   }
 
+  /** Proactively closes the server-side operation via the parent statement. */
+  private void closeServerOperation() {
+    if (parentStatement != null) {
+      parentStatement.closeServerOperation();
+    }
+  }
+
   private static TelemetryCollector resolveTelemetryCollector(
       IDatabricksStatementInternal parentStatement) {
     try {

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

@@ -59,6 +59,9 @@ public class DatabricksStatement implements IDatabricksStatement, IDatabricksSta
   // closed the operation — no further RPCs for this statement ID are possible. The JDBC Statement
   // itself remains open for re-execution. Reset on each new execution. Volatile because cancel()
   // can be called from a different thread (JDBC spec requirement).
+  private volatile boolean serverOperationClosed = false; // Client proactively closed the server
+  // operation after all results were consumed or ResultSet was closed. Prevents duplicate
+  // closeStatement RPC when Statement.close() is called later. Reset on each new execution.
   protected Boolean shouldReturnResultSet =
       null; // Cached result of shouldReturnResultSetWithConfig()
 
@@ -149,20 +152,26 @@ public void close(boolean removeFromSession) throws DatabricksSQLException {
         LOGGER.warn(warningMsg);
         warnings = WarningUtil.addWarning(warnings, warningMsg);
       } else {
-        // Skip server-side close if the server already closed the operation (direct results).
-        // The operation handle is gone on the server side, so closeStatement would fail.
-        if (!directResultsReceived) {
-          this.connection.getSession().getDatabricksClient().closeStatement(statementId);
-        } else {
-          LOGGER.debug(
-              "Statement {} closed locally (direct results — server operation already closed, "
-                  + "skipping closeStatement RPC)",
-              statementId);
-        }
+        // Close ResultSet first — this triggers proactive server close via
+        // closeServerOperation() and sets serverOperationClosed=true, preventing
+        // a duplicate closeStatement RPC below.
         if (resultSet != null) {
           this.resultSet.close();
           this.resultSet = null;
         }
+        // Skip server-side close if operation was already closed:
+        // - directResultsReceived: server closed it (inline results)
+        // - serverOperationClosed: client proactively closed it (results consumed or RS closed)
+        if (!directResultsReceived && !serverOperationClosed) {
+          this.connection.getSession().getDatabricksClient().closeStatement(statementId);
+        } else {
+          LOGGER.debug(
+              "Statement {} closed locally (server operation already closed — "
+                  + "directResults={}, proactivelyClosed={}, skipping closeStatement RPC)",
+              statementId,
+              directResultsReceived,
+              serverOperationClosed);
+        }
       }
     } finally {
       // Always run cleanup even if resultSet.close() or closeStatement() throws.
@@ -246,12 +255,11 @@ public void cancel() throws SQLException {
     LOGGER.debug("public void cancel()");
     checkIfClosed();
 
-    if (statementId != null && !directResultsReceived) {
+    if (statementId != null && !directResultsReceived && !serverOperationClosed) {
       this.connection.getSession().getDatabricksClient().cancelStatement(statementId);
       DatabricksThreadContextHolder.clearStatementInfo();
-    } else if (directResultsReceived) {
-      String warningMsg =
-          "Statement's server operation was already closed (direct results); cancel has no effect.";
+    } else if (directResultsReceived || serverOperationClosed) {
+      String warningMsg = "Statement's server operation was already closed; cancel has no effect.";
       LOGGER.debug(warningMsg);
       warnings = WarningUtil.addWarning(warnings, warningMsg);
     } else {
@@ -694,17 +702,18 @@ public ResultSet getExecutionResult() throws SQLException {
           "No execution available for statement", DatabricksDriverErrorCode.INPUT_VALIDATION_ERROR);
     }
 
-    // For direct results, the server already closed the operation — making an RPC
-    // would return "not found". Return the cached result set instead.
-    if (directResultsReceived) {
+    // For direct results or proactively closed operations, the server operation is gone —
+    // making an RPC would return "not found". Return the cached result set instead.
+    if (directResultsReceived || serverOperationClosed) {
       if (resultSet != null) {
         LOGGER.debug(
-            "Returning cached result for statement {} (direct results received)", statementId);
+            "Returning cached result for statement {} (server operation already closed)",
+            statementId);
         return resultSet;
       }
       throw new DatabricksSQLException(
-          "Direct results were received but no result set is available. "
-              + "The server closed the operation and no further results can be fetched.",
+          "Server operation was already closed and no result set is available. "
+              + "No further results can be fetched.",
           DatabricksDriverErrorCode.INVALID_STATE);
     }
 
@@ -953,6 +962,38 @@ public void markDirectResultsReceived() {
     this.directResultsReceived = true;
   }
 
+  /**
+   * Proactively closes the server-side operation to release server resources while keeping the
+   * client-side Statement open for reuse. Called when all results have been consumed (next()
+   * returns false) or when ResultSet.close() is called.
+   *
+   * <p>After this call, {@link #close(boolean)} will skip the closeStatement RPC since the server
+   * operation is already closed. The Statement can still be re-executed.
+   */
+  @Override
+  public void closeServerOperation() {
+    if (serverOperationClosed || directResultsReceived || statementId == null || isClosed) {
+      return;
+    }
+    try {
+      this.connection.getSession().getDatabricksClient().closeStatement(statementId);
+      // Only mark closed on success — if the RPC fails (transient network error),
+      // Statement.close() should retry the closeStatement RPC rather than skip it,
+      // to avoid leaving the server operation alive until session timeout.
+      this.serverOperationClosed = true;
+      LOGGER.debug(
+          "Proactively closed server operation for statement {} (results consumed)", statementId);
+    } catch (SQLException | RuntimeException e) {
+      // Best-effort — don't fail the user's close for a server cleanup failure.
+      // serverOperationClosed stays false so Statement.close() will retry the RPC.
+      LOGGER.warn(
+          "Failed to proactively close server operation for statement {}: {}",
+          statementId,
+          e.getMessage(),
+          e);
+    }
+  }
+
   /**
    * Resets statement state before a new execution (sync or async). Closes the previous server-side
    * operation handle (if still open) and the local ResultSet, clears flags, and nulls the
@@ -962,39 +1003,36 @@ private void resetForNewExecution() {
     noMoreResults = false;
     updateCount = -1;
 
-    // Close the previous server-side operation if it exists. This prevents resource
-    // leaks when a Statement is re-executed (e.g., PreparedStatement in a loop).
-    // This matches the behavior of pgJDBC, MySQL Connector/J, Trino JDBC, and
-    // Databricks Python SQL Connector — all close the previous operation on re-execute.
-    //
-    // Note on directResultsReceived: we check the flag value from the PREVIOUS execution
-    // here. The flag is reset to false below, after this close attempt.
-    //
-    // Note on latency: this close is synchronous (adds one RPC round-trip before the next
-    // execution). This is consistent with pgJDBC's closeForNextExecution() which is also
-    // synchronous. The correctness benefit (no orphaned server operations) outweighs the
-    // latency cost for typical usage patterns.
-    //
-    // Skip if: (1) no previous execution (statementId==null), or
-    //          (2) server already closed the operation (direct results).
-    if (statementId != null && !directResultsReceived) {
-      try {
-        connection.getSession().getDatabricksClient().closeStatement(statementId);
-      } catch (Exception e) {
-        // Don't block re-execution if closing the previous operation fails.
-        // This covers: network errors, operation already expired/evicted on server,
-        // and transport-level errors (e.g., unexpected server responses).
-        // The new execution will create a fresh operation with a new statementId.
-        LOGGER.debug(
-            "Failed to close previous server operation {} during re-execution: {}",
-            statementId,
-            e.getMessage());
-      }
+    // Close the previous server-side operation if still open. Fire-and-forget on a
+    // daemon thread so the new execution is not blocked by the close RPC latency.
+    // If the user closed the ResultSet before re-executing (best practice), the
+    // proactive close already set serverOperationClosed=true and this is a no-op.
+    if (statementId != null && !directResultsReceived && !serverOperationClosed) {
+      final StatementId prevStatementId = statementId;
+      final IDatabricksClient prevClient = connection.getSession().getDatabricksClient();
+      Thread closeThread =
+          new Thread(
+              () -> {
+                try {
+                  prevClient.closeStatement(prevStatementId);
+                  LOGGER.debug(
+                      "Closed previous server operation {} during re-execution", prevStatementId);
+                } catch (SQLException | RuntimeException e) {
+                  LOGGER.warn(
+                      "Failed to close previous server operation {} during re-execution: {}",
+                      prevStatementId,
+                      e.getMessage(),
+                      e);
+                }
+              });
+      closeThread.setDaemon(true);
+      closeThread.setName("close-stmt-" + prevStatementId);
+      closeThread.start();
+      serverOperationClosed = true;
     }
 
-    directResultsReceived = false;
-
-    // Per JDBC spec, re-executing a Statement implicitly closes the current ResultSet.
+    // Close the previous ResultSet. closeServerOperation() inside resultSet.close()
+    // is a no-op since serverOperationClosed was set above.
     if (resultSet != null) {
       try {
         resultSet.close();
@@ -1004,6 +1042,10 @@ private void resetForNewExecution() {
       resultSet = null;
     }
 
+    // Reset flags AFTER closing old ResultSet
+    directResultsReceived = false;
+    serverOperationClosed = false;
+
     // Null out statementId so that if the new execution fails before setStatementId(),
     // close() takes the statementId==null branch instead of sending closeStatement(stale-id)
     statementId = null;

src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java

@@ -37,4 +37,15 @@ public interface IDatabricksStatementInternal {
   default void markDirectResultsReceived() {
     // no-op by default
   }
+
+  /**
+   * Proactively closes the server-side operation to release server resources while keeping the
+   * client-side Statement open for reuse. Default no-op for implementations that don't support
+   * proactive close.
+   */
+  default void closeServerOperation() {
+    // no-op by default
+  }
+
+  long getLargeMaxRows() throws DatabricksSQLException;
 }