[PECOBLR-2321] Result Set Heartbeat / Keep-Alive for Ongoing Query Executions

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### 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).
 - Added `CallableStatement` support with IN parameters. `Connection.prepareCall()` now returns a working `DatabricksCallableStatement` that supports positional parameter binding and execution via `{call proc(?)}` JDBC escape syntax. OUT/INOUT parameters and named parameters throw `SQLFeatureNotSupportedException`.
 - Added AI coding agent detection to the User-Agent header. When the driver is invoked by a known AI coding agent (e.g. Claude Code, Cursor, Gemini CLI), `agent/<product>` is appended to the User-Agent string.
 

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

@@ -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) {
+    if (connectionContext instanceof DatabricksConnectionContext) {
+      DatabricksConnectionContext ctx = (DatabricksConnectionContext) connectionContext;
+      if (ctx.isHeartbeatEnabled()) {
+        return new ResultHeartbeatManager(ctx.getHeartbeatIntervalSeconds());
+      }
+    }
+    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();
@@ -420,6 +439,9 @@ public void close() throws SQLException {
       statement.close(false);
       statementSet.remove(statement);
     }
+    if (heartbeatManager != null) {
+      heartbeatManager.shutdown();
+    }
     this.session.close();
     TelemetryClientFactory.getInstance().closeTelemetryClient(connectionContext);
     DatabricksClientConfiguratorManager.getInstance().removeInstance(connectionContext);

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

@@ -908,6 +908,26 @@ 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 =
+        Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HEARTBEAT_INTERVAL_SECONDS));
+    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/DatabricksResultSet.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;
@@ -123,6 +124,7 @@ public DatabricksResultSet(
     this.cachedTelemetryCollector = resolveTelemetryCollector(parentStatement);
     this.isClosed = false;
     this.wasNull = false;
+    startHeartbeatIfEnabled();
   }
 
   @VisibleForTesting
@@ -283,18 +285,167 @@ public boolean next() throws SQLException {
       cachedTelemetryCollector.recordResultSetIteration(
           statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext);
     }
+    if (!hasNext) {
+      stopHeartbeat();
+    }
     return hasNext;
   }
 
   @Override
   public void close() throws DatabricksSQLException {
+    stopHeartbeat();
     isClosed = true;
     this.executionResult.close();
     if (parentStatement != null) {
       parentStatement.handleResultSetClose(this);
     }
   }
 
+  /** 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 {
+      DatabricksConnection conn =
+          (DatabricksConnection) parentStatement.getStatement().getConnection();
+      ResultHeartbeatManager mgr = conn.getHeartbeatManager();
+      if (mgr == null) {
+        return; // heartbeat not enabled
+      }
+
+      IDatabricksClient client = conn.getSession().getDatabricksClient();
+      final int maxConsecutiveFailures = 10;
+      final java.util.concurrent.atomic.AtomicInteger consecutiveFailures =
+          new java.util.concurrent.atomic.AtomicInteger(0);
+      // Get the stopped flag from the manager — shared between the heartbeat task and
+      // stopHeartbeat(). Prevents RPC on a just-closed client/session: stopHeartbeat sets
+      // the flag before cancel(false), so an in-flight tick sees it and skips the RPC.
+      final java.util.concurrent.atomic.AtomicBoolean stopped = mgr.getStoppedFlag(statementId);
+
+      Runnable heartbeatTask =
+          () -> {
+            if (stopped.get()) {
+              return; // client/session may be closed, skip RPC
+            }
+            try {
+              boolean alive = client.checkStatementAlive(statementId);
+              consecutiveFailures.set(0); // reset on success
+              if (!alive) {
+                LOGGER.info(
+                    "Heartbeat detected terminal state for statement {}, stopping", statementId);
+                stopped.set(true);
+                stopHeartbeat();
+              }
+            } catch (Exception e) {
+              // If stopped was set during the RPC (connection closing), don't count as failure
+              if (stopped.get()) {
+                return;
+              }
+              int failures = consecutiveFailures.incrementAndGet();
+              if (failures == 1) {
+                // First failure — log at INFO so users see the initial problem
+                LOGGER.info(
+                    "Heartbeat failed for statement {} (first failure): {}",
+                    statementId,
+                    e.getMessage());
+              } else {
+                LOGGER.debug(
+                    "Heartbeat failed for statement {} (failure {}/{}): {}",
+                    statementId,
+                    failures,
+                    maxConsecutiveFailures,
+                    e.getMessage());
+              }
+              if (failures >= maxConsecutiveFailures) {
+                // Terminal failure — log at WARN so it's visible in default log config
+                LOGGER.warn(
+                    "Heartbeat stopped for statement {} after {} consecutive failures. "
+                        + "Server-side results may expire. Last error: {}",
+                    statementId,
+                    failures,
+                    e.getMessage());
+                stopped.set(true);
+                stopHeartbeat();
+              }
+            }
+          };
+
+      mgr.startHeartbeat(statementId, heartbeatTask);
+      LOGGER.debug(
+          "Heartbeat started for statement {} (resultType={}, interval={}s)",
+          statementId,
+          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 {
+      DatabricksConnection conn =
+          (DatabricksConnection) parentStatement.getStatement().getConnection();
+      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 {

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

@@ -172,6 +172,13 @@ public void close(boolean removeFromSession) throws DatabricksSQLException {
         this.connection.closeStatement(this);
       }
       DatabricksThreadContextHolder.clearStatementInfo();
+      // Safety net: stop any heartbeat for this statement
+      if (statementId != null) {
+        ResultHeartbeatManager mgr = connection.getHeartbeatManager();
+        if (mgr != null) {
+          mgr.stopHeartbeat(statementId);
+        }
+      }
       shutDownExecutor();
       this.updateCount = -1;
       this.isClosed = true;
@@ -672,6 +679,8 @@ public ResultSet executeAsync(String sql) throws SQLException {
     LOGGER.debug("ResultSet executeAsync() for statement {%s}", sql);
     checkIfClosed();
 
+    // No heartbeat during async wait — the user controls polling via getExecutionResult().
+    // Heartbeat starts later when the ResultSet is constructed (after getExecutionResult()).
     resetForNewExecution();
 
     IDatabricksClient client = connection.getSession().getDatabricksClient();
@@ -969,6 +978,16 @@ private void resetForNewExecution() {
     // when the server returns unexpected responses (e.g., WireMock 404 in tests).
     // For direct results, the server already closed the handle.
 
+    // Stop heartbeat for the previous execution before clearing state.
+    // Without this, the old heartbeat (keyed by old statementId) would fail and self-terminate
+    // after 10 consecutive failures — wasteful and noisy in logs.
+    if (statementId != null) {
+      ResultHeartbeatManager mgr = connection.getHeartbeatManager();
+      if (mgr != null) {
+        mgr.stopHeartbeat(statementId);
+      }
+    }
+
     directResultsReceived = false;
 
     // Per JDBC spec, re-executing a Statement implicitly closes the current ResultSet.