[PECOBLR-2321] Result Set Heartbeat / Keep-Alive for Ongoing Query Executions (#1415)

## Summary
Design + implementation for PECOBLR-2321: periodic heartbeat polling to
keep server-side result state alive while the client consumes results
slowly.

## Problem
When users read query results slowly (pausing between `next()` calls),
the warehouse can auto-stop after its idle timeout. For inline results
(data only on cluster, not uploaded to cloud storage), this means
permanent data loss. The user gets errors like `INVALID_HANDLE_STATUS`
or "operation not found".

## Solution
A `ResultHeartbeatManager` that periodically calls `GetStatementStatus`
(SEA) or `GetOperationStatus` (Thrift) to signal the server that results
are still being consumed. Opt-in via `EnableHeartbeat=1` connection
parameter (default false due to cost implications).

### Design doc
`docs/design/HEARTBEAT_KEEP_ALIVE.md` — includes cross-driver survey,
Mermaid diagrams (sequence flows, state machine, class diagram), and
detailed lifecycle analysis.

### Heartbeat eligibility (skipped when not needed)
| Scenario | Heartbeat? | Reason |
|----------|-----------|--------|
| SEA cloud fetch (Arrow) | **Yes** | Statement must stay alive for URL
refresh |
| Thrift inline (columnar) | **Yes** | Data fetched on-demand; server
can evict |
| Thrift cloud fetch | **Yes** | Operation handle must stay alive |
| SEA inline (JSON) | No | All data loaded into memory at construction |
| Direct results (CLOSED state) | No | Server already closed; data fully
delivered |
| Update count (DML) | No | No result rows; execution polling already
kept it alive |
| Async execution wait | No | User controls polling via
`getExecutionResult()` — heartbeat starts only when ResultSet is
constructed |

### Error resilience
- 10 consecutive failures before self-stop
- Terminal states (CLOSED/ERROR/CANCELED/TIMEDOUT) auto-stop
- Single success resets failure counter

### Zero-leak guarantee
Heartbeat stops in 4 places: `next()` returns false,
`ResultSet.close()`, `Statement.close()`, `Connection.close()`

### Implementation
- **New**: `ResultHeartbeatManager` — per-connection manager with
`ScheduledExecutorService` (daemon thread)
- **New**: `ResultHeartbeatManagerTest` — 7 unit tests
- **Modified**: `DatabricksJdbcUrlParams` — `EnableHeartbeat` (default
0), `HeartbeatIntervalSeconds` (default 60)
- **Modified**: `DatabricksConnectionContext` — getter methods
- **Modified**: `DatabricksConnection` — creates/shuts down manager
- **Modified**: `DatabricksResultSet` — starts heartbeat in constructor,
stops on close/next-false
- **Modified**: `DatabricksStatement` — safety net stop in close()
- **Modified**: `IDatabricksClient` — `checkStatementAlive()` default
method
- **Modified**: `DatabricksSdkClient` — SEA heartbeat via GET
`/sql/statements/{id}`
- **Modified**: `DatabricksThriftServiceClient` — Thrift heartbeat via
`GetOperationStatus`

---------

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

Author: gopalldb

NEXT_CHANGELOG.md

@@ -53,6 +53,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters.
   compatibility issues.
 
 ### Added
+- Added result set heartbeat / keep-alive to prevent server-side result expiry during slow consumption. When enabled via `EnableHeartbeat=1`, the driver periodically polls `GetStatementStatus` (SEA) or `GetOperationStatus` (Thrift) to keep the operation alive while the client reads results. Configurable interval via `HeartbeatIntervalSeconds` (default 60s). Heartbeat automatically stops when results are fully consumed, ResultSet is closed, or the server returns a terminal state. Disabled by default due to cost implications (heartbeats keep the warehouse running).
 - Metadata operations now use SQL SHOW commands for both Thrift and SEA backends,
   ensuring consistent behavior for SQL warehouses regardless of underlying
   protocol. To revert to native Thrift metadata RPCs, set `UseQueryForMetadata=0`.

docs/design/HEARTBEAT_KEEP_ALIVE.md

@@ -285,6 +285,7 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration>
+          <excludedGroups>e2e</excludedGroups>
           <excludes>
             <exclude>**/DatabricksDriverExamples.java</exclude>
             <exclude>**/integration/**/*.java</exclude>

jdbc-core/pom.xml

@@ -38,6 +38,7 @@ public class DatabricksConnection implements IDatabricksConnection, IDatabricksC
   private final Set<IDatabricksStatementInternal> statementSet = ConcurrentHashMap.newKeySet();
   private SQLWarning warnings = null;
   private final IDatabricksConnectionContext connectionContext;
+  private final ResultHeartbeatManager heartbeatManager;
 
   /**
    * Creates an instance of Databricks connection for given connection context.
@@ -49,6 +50,7 @@ public DatabricksConnection(IDatabricksConnectionContext connectionContext)
     this.connectionContext = connectionContext;
     DatabricksThreadContextHolder.setConnectionContext(connectionContext);
     this.session = new DatabricksSession(connectionContext);
+    this.heartbeatManager = createHeartbeatManager(connectionContext);
   }
 
   @VisibleForTesting
@@ -58,10 +60,27 @@ public DatabricksConnection(
     this.connectionContext = connectionContext;
     DatabricksThreadContextHolder.setConnectionContext(connectionContext);
     this.session = new DatabricksSession(connectionContext, testDatabricksClient);
+    this.heartbeatManager = createHeartbeatManager(connectionContext);
     UserAgentManager.setUserAgent(connectionContext);
     TelemetryHelper.updateTelemetryAppName(connectionContext, null);
   }
 
+  private static ResultHeartbeatManager createHeartbeatManager(
+      IDatabricksConnectionContext connectionContext) {
+    // Use interface methods instead of instanceof check so mocks and
+    // alternate implementations can also enable heartbeat
+    if (connectionContext.isHeartbeatEnabled()) {
+      return new ResultHeartbeatManager(
+          connectionContext.getHeartbeatIntervalSeconds(), connectionContext.getConnectionUuid());
+    }
+    return null;
+  }
+
+  /** Returns the heartbeat manager, or null if heartbeat is disabled. */
+  ResultHeartbeatManager getHeartbeatManager() {
+    return heartbeatManager;
+  }
+
   @Override
   public void open() throws SQLException {
     this.session.open();
@@ -416,6 +435,11 @@ public void rollback() throws SQLException {
   @Override
   public void close() throws SQLException {
     LOGGER.debug("public void close()");
+    // Shutdown heartbeat FIRST — prevents RPCs on closing connections and
+    // ensures shutdown runs even if statement.close() throws
+    if (heartbeatManager != null) {
+      heartbeatManager.shutdown();
+    }
     for (IDatabricksStatementInternal statement : statementSet) {
       statement.close(false);
       statementSet.remove(statement);

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

@@ -973,6 +973,33 @@ public boolean isTelemetryEnabled() {
     return getParameter(DatabricksJdbcUrlParams.ENABLE_TELEMETRY).equals("1");
   }
 
+  public boolean isHeartbeatEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_HEARTBEAT).equals("1");
+  }
+
+  public int getHeartbeatIntervalSeconds() {
+    int interval;
+    try {
+      interval = Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HEARTBEAT_INTERVAL_SECONDS));
+    } catch (NumberFormatException e) {
+      LOGGER.warn(
+          "Invalid HeartbeatIntervalSeconds value '{}'. Using default 60.",
+          getParameter(DatabricksJdbcUrlParams.HEARTBEAT_INTERVAL_SECONDS));
+      return 60;
+    }
+    if (interval <= 0) {
+      LOGGER.warn("HeartbeatIntervalSeconds must be positive, got {}. Using default 60.", interval);
+      return 60;
+    }
+    if (interval > 3600) {
+      LOGGER.warn(
+          "HeartbeatIntervalSeconds {} is very large (> 1 hour). "
+              + "Heartbeat may not keep the operation alive.",
+          interval);
+    }
+    return interval;
+  }
+
   @Override
   public String getVolumeOperationAllowedPaths() {
     return getParameter(

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

@@ -21,6 +21,7 @@
 import com.databricks.jdbc.common.Nullable;
 import com.databricks.jdbc.common.StatementType;
 import com.databricks.jdbc.common.util.WarningUtil;
+import com.databricks.jdbc.dbclient.IDatabricksClient;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
@@ -134,6 +135,7 @@ public DatabricksResultSet(
     this.maxRowsLimit = resolveMaxRowsLimit(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
+    startHeartbeatIfEnabled();
   }
 
   @VisibleForTesting
@@ -204,6 +206,7 @@ public DatabricksResultSet(
     this.maxRowsLimit = resolveMaxRowsLimit(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
+    startHeartbeatIfEnabled();
   }
 
   /* Constructing results for getUDTs, getTypeInfo, getProcedures metadata calls */
@@ -305,6 +308,12 @@ public DatabricksResultSet(
   @Override
   public boolean next() throws SQLException {
     checkIfClosed();
+    if (executionResult == null) {
+      throw new DatabricksSQLException(
+          "Cannot iterate: no result data available. "
+              + "For async execution, call getExecutionResult() first.",
+          DatabricksDriverErrorCode.INVALID_STATE);
+    }
     // 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
@@ -317,15 +326,22 @@ public boolean next() throws SQLException {
           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();
+    boolean hasNext;
+    try {
+      hasNext = this.executionResult.next();
+    } catch (Exception e) {
+      // Stop heartbeat on iteration failure — prevents keeping the warehouse alive
+      // for an abandoned ResultSet (up to 10 ticks × interval before self-stop).
+      stopHeartbeat();
+      throw e;
+    }
     // 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).
@@ -336,15 +352,21 @@ public boolean next() throws SQLException {
       cachedTelemetryCollector.recordResultSetIteration(
           statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext);
     }
+    if (!hasNext) {
+      stopHeartbeat();
+    }
     return hasNext;
   }
 
   @Override
   public void close() throws DatabricksSQLException {
+    stopHeartbeat();
     // Proactively close server operation when ResultSet is closed explicitly.
     closeServerOperation();
     isClosed = true;
-    this.executionResult.close();
+    if (executionResult != null) {
+      executionResult.close();
+    }
     if (parentStatement != null) {
       parentStatement.handleResultSetClose(this);
     }
@@ -357,6 +379,187 @@ private void closeServerOperation() {
     }
   }
 
+  /** Starts heartbeat polling if enabled on the connection and this result set is eligible. */
+  private void startHeartbeatIfEnabled() {
+    if (parentStatement == null || statementId == null) {
+      return;
+    }
+    if (!isHeartbeatEligible()) {
+      return;
+    }
+
+    try {
+      // Use JDBC unwrap() to handle pooled connection wrappers (HikariCP, DBCP)
+      java.sql.Connection rawConn = parentStatement.getStatement().getConnection();
+      DatabricksConnection conn;
+      if (rawConn instanceof DatabricksConnection) {
+        conn = (DatabricksConnection) rawConn;
+      } else if (rawConn.isWrapperFor(DatabricksConnection.class)) {
+        conn = rawConn.unwrap(DatabricksConnection.class);
+      } else {
+        LOGGER.debug("Cannot start heartbeat: connection is not a DatabricksConnection");
+        return;
+      }
+
+      ResultHeartbeatManager mgr = conn.getHeartbeatManager();
+      if (mgr == null) {
+        return; // heartbeat not enabled
+      }
+
+      // Capture only what the lambda needs — avoid capturing 'this' to prevent
+      // abandoned ResultSets from keeping the warehouse alive via heartbeat.
+      // Note: capturing 'client' retains a reference to the session/connection. If the
+      // connection is GC'd without close(), heartbeat RPCs will fail and self-stop after
+      // maxConsecutiveFailures (10 ticks, ~10 min at 60s interval). Acceptable tradeoff.
+      final IDatabricksClient client = conn.getSession().getDatabricksClient();
+      final StatementId capturedStatementId = this.statementId;
+      final int maxConsecutiveFailures = 10;
+      final java.util.concurrent.atomic.AtomicInteger consecutiveFailures =
+          new java.util.concurrent.atomic.AtomicInteger(0);
+      // Read the stopped flag from the manager on each tick instead of pre-capturing.
+      // Pre-capturing caused an orphan-flag bug: startHeartbeat() internally calls
+      // stopHeartbeat() which removes and replaces the flag, leaving the lambda with a
+      // permanently-true reference. Reading from the manager each tick always gets the
+      // current flag.
+      final ResultHeartbeatManager capturedMgr = mgr;
+
+      Runnable heartbeatTask =
+          () -> {
+            // Read current flag each tick — avoids orphan-flag issue
+            java.util.concurrent.atomic.AtomicBoolean stopped =
+                capturedMgr.getStoppedFlag(capturedStatementId);
+            if (stopped.get()) {
+              return; // client/session may be closed, skip RPC
+            }
+            try {
+              boolean alive = client.checkStatementAlive(capturedStatementId);
+              consecutiveFailures.set(0); // reset on success
+              if (!alive) {
+                LOGGER.info(
+                    "Heartbeat detected terminal state for statement {}, stopping",
+                    capturedStatementId.toSQLExecStatementId());
+                capturedMgr.stopHeartbeat(capturedStatementId);
+              }
+            } catch (Throwable e) {
+              if (e instanceof VirtualMachineError) {
+                capturedMgr.stopHeartbeat(capturedStatementId);
+                throw (VirtualMachineError) e;
+              }
+              if (capturedMgr.getStoppedFlag(capturedStatementId).get()) {
+                return;
+              }
+              if (e instanceof java.sql.SQLFeatureNotSupportedException) {
+                LOGGER.debug(
+                    "Heartbeat not supported by client for statement {}, stopping",
+                    capturedStatementId.toSQLExecStatementId());
+                capturedMgr.stopHeartbeat(capturedStatementId);
+                return;
+              }
+              int failures = consecutiveFailures.incrementAndGet();
+              if (failures == 1) {
+                LOGGER.info(
+                    "Heartbeat failed for statement {} (first failure): {}",
+                    capturedStatementId.toSQLExecStatementId(),
+                    e.getMessage());
+              } else {
+                LOGGER.debug(
+                    "Heartbeat failed for statement {} (failure {}/{}): {}",
+                    capturedStatementId.toSQLExecStatementId(),
+                    failures,
+                    maxConsecutiveFailures,
+                    e.getMessage());
+              }
+              if (failures >= maxConsecutiveFailures) {
+                LOGGER.warn(
+                    "Heartbeat stopped for statement {} after {} consecutive failures. "
+                        + "Server-side results may expire. Last error: {}",
+                    capturedStatementId.toSQLExecStatementId(),
+                    failures,
+                    e.getMessage());
+                capturedMgr.stopHeartbeat(capturedStatementId);
+              }
+            }
+          };
+
+      mgr.startHeartbeat(capturedStatementId, heartbeatTask);
+      LOGGER.debug(
+          "Heartbeat started for statement {} (resultType={}, interval={}s)",
+          capturedStatementId.toSQLExecStatementId(),
+          resultSetType,
+          mgr.getIntervalSeconds());
+    } catch (Exception e) {
+      LOGGER.debug("Failed to start heartbeat: {}", e.getMessage());
+    }
+  }
+
+  /** Stops the heartbeat for this result set's statement. Idempotent. */
+  private void stopHeartbeat() {
+    if (parentStatement == null || statementId == null) {
+      return;
+    }
+    try {
+      // Use same unwrap pattern as startHeartbeatIfEnabled() for pooled connections
+      java.sql.Connection rawConn = parentStatement.getStatement().getConnection();
+      DatabricksConnection conn;
+      if (rawConn instanceof DatabricksConnection) {
+        conn = (DatabricksConnection) rawConn;
+      } else if (rawConn.isWrapperFor(DatabricksConnection.class)) {
+        conn = rawConn.unwrap(DatabricksConnection.class);
+      } else {
+        return;
+      }
+      ResultHeartbeatManager mgr = conn.getHeartbeatManager();
+      if (mgr != null) {
+        mgr.stopHeartbeat(statementId);
+      }
+    } catch (Exception e) {
+      LOGGER.debug("Failed to stop heartbeat: {}", e.getMessage());
+    }
+  }
+
+  /**
+   * Determines whether this result set is eligible for heartbeat polling. Package-visible for
+   * testing.
+   *
+   * <p>Heartbeat is NOT needed when:
+   *
+   * <ul>
+   *   <li>No execution result (nothing to fetch, also covers async PENDING/RUNNING with no data)
+   *   <li>SEA inline (InlineJsonResult): all rows loaded in memory at construction
+   *   <li>Update count (DML): no result rows to keep alive
+   *   <li>Direct results (CLOSED state): server already closed, data fully delivered
+   *   <li>Async execution (PENDING/RUNNING): user controls polling via getExecutionResult()
+   * </ul>
+   */
+  boolean isHeartbeatEligible() {
+    // No execution result — nothing to fetch
+    if (executionResult == null) {
+      return false;
+    }
+    // SEA inline — all data loaded in memory at construction
+    if (resultSetType == ResultSetType.SEA_INLINE) {
+      return false;
+    }
+    // Update count — no result rows
+    if (statementType == StatementType.UPDATE) {
+      return false;
+    }
+    // Check execution state
+    if (executionStatus != null) {
+      com.databricks.jdbc.api.ExecutionState state = executionStatus.getExecutionState();
+      // Direct results — server already closed
+      if (state == com.databricks.jdbc.api.ExecutionState.CLOSED) {
+        return false;
+      }
+      // Async execution — user controls polling
+      if (state == com.databricks.jdbc.api.ExecutionState.PENDING
+          || state == com.databricks.jdbc.api.ExecutionState.RUNNING) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   private static TelemetryCollector resolveTelemetryCollector(
       IDatabricksStatementInternal parentStatement) {
     try {