Add session version exchange to SQL Exec API client

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### Added
+- Added session-version exchange for SQL Exec API connections.
 
 ### Updated
 - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior.

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

@@ -429,6 +429,8 @@ private void startHeartbeatIfEnabled() {
       // 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 IDatabricksSession capturedSession = conn.getSession();
+      final String originatingSessionId = parentStatement.getOriginatingSessionId();
       final StatementId capturedStatementId = this.statementId;
       final int maxConsecutiveFailures = 10;
       final java.util.concurrent.atomic.AtomicInteger consecutiveFailures =
@@ -449,7 +451,9 @@ private void startHeartbeatIfEnabled() {
               return; // client/session may be closed, skip RPC
             }
             try {
-              boolean alive = client.checkStatementAlive(capturedStatementId);
+              boolean alive =
+                  client.checkStatementAlive(
+                      capturedStatementId, capturedSession, originatingSessionId);
               consecutiveFailures.set(0); // reset on success
               if (!alive) {
                 LOGGER.info(

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

@@ -21,6 +21,7 @@
 import com.databricks.jdbc.exception.DatabricksTemporaryRedirectException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.jdbc.model.core.SessionVersion;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.jdbc.telemetry.TelemetryHelper;
 import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimedProcessor;
@@ -29,6 +30,7 @@
 import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 
 /**
@@ -43,6 +45,7 @@ public class DatabricksSession implements IDatabricksSession {
   private final IDatabricksComputeResource computeResource;
   private boolean isSessionOpen;
   private ImmutableSessionInfo sessionInfo;
+  private final AtomicReference<Long> sessionVersion = new AtomicReference<>();
 
   /** For context based commands */
   private String catalog;
@@ -111,6 +114,37 @@ public ImmutableSessionInfo getSessionInfo() {
     return sessionInfo;
   }
 
+  @Nullable
+  @Override
+  public SessionVersion getSessionVersion() {
+    Long versionId = sessionVersion.get();
+    return versionId == null ? null : new SessionVersion().setVersionId(versionId);
+  }
+
+  @Override
+  public void updateSessionVersion(
+      @Nullable String expectedSessionId, @Nullable SessionVersion newSessionVersion) {
+    if (expectedSessionId == null
+        || newSessionVersion == null
+        || newSessionVersion.getVersionId() == null) {
+      return;
+    }
+    synchronized (this) {
+      if (!isSessionOpen
+          || sessionInfo == null
+          || !expectedSessionId.equals(sessionInfo.sessionId())) {
+        return;
+      }
+      Long newVersionId = newSessionVersion.getVersionId();
+      sessionVersion.accumulateAndGet(
+          newVersionId,
+          (currentVersion, candidateVersion) ->
+              currentVersion == null || candidateVersion > currentVersion
+                  ? candidateVersion
+                  : currentVersion);
+    }
+  }
+
   @Override
   public IDatabricksComputeResource getComputeResource() {
     LOGGER.debug("public String getComputeResource()");
@@ -217,6 +251,7 @@ public void open() throws SQLException {
             throw e;
           }
         }
+        this.sessionVersion.set(sessionInfo == null ? null : sessionInfo.sessionVersion());
         this.isSessionOpen = true;
       }
     }
@@ -240,6 +275,7 @@ public void close() throws SQLException {
         } finally {
           // Always clean up local state
           this.sessionInfo = null;
+          this.sessionVersion.set(null);
           this.isSessionOpen = false;
         }
       }
@@ -406,6 +442,7 @@ public void forceClose() {
     } catch (SQLException e) {
       LOGGER.error("Error closing session resources, but marking the session as closed.");
     } finally {
+      this.sessionVersion.set(null);
       this.isSessionOpen = false;
     }
   }

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

@@ -25,6 +25,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.*;
+import javax.annotation.Nullable;
 import org.apache.http.entity.InputStreamEntity;
 
 public class DatabricksStatement implements IDatabricksStatement, IDatabricksStatementInternal {
@@ -44,6 +45,7 @@ public class DatabricksStatement implements IDatabricksStatement, IDatabricksSta
   protected final DatabricksConnection connection;
   DatabricksResultSet resultSet;
   private volatile StatementId statementId; // volatile: cancel() reads from a different thread
+  private volatile String originatingSessionId;
   private boolean isClosed;
   private boolean closeOnCompletion;
   private SQLWarning warnings = null;
@@ -69,6 +71,7 @@ public DatabricksStatement(DatabricksConnection connection) throws DatabricksVal
     this.connection = connection;
     this.resultSet = null;
     this.statementId = null;
+    this.originatingSessionId = null;
     this.isClosed = false;
     this.timeoutInSeconds = DEFAULT_STATEMENT_TIMEOUT_SECONDS;
     this.databricksBatchExecutor =
@@ -79,6 +82,7 @@ public DatabricksStatement(DatabricksConnection connection, StatementId statemen
       throws DatabricksValidationException {
     this.connection = connection;
     this.statementId = statementId;
+    this.originatingSessionId = null;
     this.resultSet = null;
     this.isClosed = false;
     this.timeoutInSeconds = DEFAULT_STATEMENT_TIMEOUT_SECONDS;
@@ -644,8 +648,14 @@ public void handleResultSetClose(IDatabricksResultSet resultSet) throws Databric
 
   @Override
   public void setStatementId(StatementId statementId) {
+    setStatementId(statementId, null);
+  }
+
+  @Override
+  public void setStatementId(StatementId statementId, @Nullable String originatingSessionId) {
     LOGGER.debug("void setStatementId(Statement statementId = {})", statementId);
     this.statementId = statementId;
+    this.originatingSessionId = originatingSessionId;
   }
 
   @Override
@@ -658,6 +668,12 @@ public Statement getStatement() {
     return this;
   }
 
+  @Override
+  @Nullable
+  public String getOriginatingSessionId() {
+    return originatingSessionId;
+  }
+
   @Override
   public void allowInputStreamForVolumeOperation(boolean allowInputStream)
       throws DatabricksSQLException {
@@ -1097,6 +1113,7 @@ private void resetForNewExecution() {
     // 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;
+    originatingSessionId = null;
   }
 
   /**

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

@@ -12,6 +12,9 @@ public interface SessionInfo {
 
   IDatabricksComputeResource computeResource();
 
+  @Nullable
+  Long sessionVersion();
+
   @Nullable
   TSessionHandle sessionHandle(); // This field is set only for all-purpose cluster compute
 }

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

@@ -6,6 +6,7 @@
 import com.databricks.jdbc.dbclient.IDatabricksClient;
 import com.databricks.jdbc.dbclient.IDatabricksMetadataClient;
 import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.model.core.SessionVersion;
 import java.sql.SQLException;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -24,6 +25,14 @@ public interface IDatabricksSession {
   @Nullable
   ImmutableSessionInfo getSessionInfo();
 
+  @Nullable
+  default SessionVersion getSessionVersion() {
+    return null;
+  }
+
+  default void updateSessionVersion(
+      @Nullable String expectedSessionId, @Nullable SessionVersion sessionVersion) {}
+
   /**
    * Get the warehouse associated with the session.
    *

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

@@ -4,6 +4,7 @@
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import java.sql.Statement;
+import javax.annotation.Nullable;
 import org.apache.http.entity.InputStreamEntity;
 
 /** Extended callback handle for java.sql.Statement interface */
@@ -19,10 +20,19 @@ public interface IDatabricksStatementInternal {
 
   void setStatementId(StatementId statementId);
 
+  default void setStatementId(StatementId statementId, @Nullable String originatingSessionId) {
+    setStatementId(statementId);
+  }
+
   StatementId getStatementId();
 
   Statement getStatement();
 
+  @Nullable
+  default String getOriginatingSessionId() {
+    return null;
+  }
+
   void allowInputStreamForVolumeOperation(boolean allowedInputStream) throws DatabricksSQLException;
 
   boolean isAllowedInputStreamForVolumeOperation() throws DatabricksSQLException;

src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java

@@ -16,6 +16,7 @@
 import com.databricks.sdk.core.DatabricksConfig;
 import java.sql.SQLException;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 /** Interface for Databricks client which abstracts the integration with Databricks server. */
 public interface IDatabricksClient {
@@ -120,6 +121,12 @@ default boolean checkStatementAlive(StatementId statementId) throws SQLException
     throw new java.sql.SQLFeatureNotSupportedException("Heartbeat not supported by this client");
   }
 
+  default boolean checkStatementAlive(
+      StatementId statementId, IDatabricksSession session, @Nullable String originatingSessionId)
+      throws SQLException {
+    return checkStatementAlive(statementId);
+  }
+
   /**
    * Fetches result for underlying statement-Id
    *