Superseded by #1640

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### Added
+- Added FAST explicit-session mode and 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 =
+          parentStatement.shouldTrackSessionVersion() ? conn.getSession() : null;
       final StatementId capturedStatementId = this.statementId;
       final int maxConsecutiveFailures = 10;
       final java.util.concurrent.atomic.AtomicInteger consecutiveFailures =
@@ -449,7 +451,7 @@ private void startHeartbeatIfEnabled() {
               return; // client/session may be closed, skip RPC
             }
             try {
-              boolean alive = client.checkStatementAlive(capturedStatementId);
+              boolean alive = client.checkStatementAlive(capturedStatementId, capturedSession);
               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

@@ -44,6 +44,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 boolean trackSessionVersion;
   private boolean isClosed;
   private boolean closeOnCompletion;
   private SQLWarning warnings = null;
@@ -69,6 +70,7 @@ public DatabricksStatement(DatabricksConnection connection) throws DatabricksVal
     this.connection = connection;
     this.resultSet = null;
     this.statementId = null;
+    this.trackSessionVersion = true;
     this.isClosed = false;
     this.timeoutInSeconds = DEFAULT_STATEMENT_TIMEOUT_SECONDS;
     this.databricksBatchExecutor =
@@ -79,6 +81,7 @@ public DatabricksStatement(DatabricksConnection connection, StatementId statemen
       throws DatabricksValidationException {
     this.connection = connection;
     this.statementId = statementId;
+    this.trackSessionVersion = false;
     this.resultSet = null;
     this.isClosed = false;
     this.timeoutInSeconds = DEFAULT_STATEMENT_TIMEOUT_SECONDS;
@@ -646,6 +649,7 @@ public void handleResultSetClose(IDatabricksResultSet resultSet) throws Databric
   public void setStatementId(StatementId statementId) {
     LOGGER.debug("void setStatementId(Statement statementId = {})", statementId);
     this.statementId = statementId;
+    this.trackSessionVersion = true;
   }
 
   @Override
@@ -658,6 +662,11 @@ public Statement getStatement() {
     return this;
   }
 
+  @Override
+  public boolean shouldTrackSessionVersion() {
+    return trackSessionVersion;
+  }
+
   @Override
   public void allowInputStreamForVolumeOperation(boolean allowInputStream)
       throws DatabricksSQLException {

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,12 @@ public interface IDatabricksSession {
   @Nullable
   ImmutableSessionInfo getSessionInfo();
 
+  @Nullable
+  SessionVersion getSessionVersion();
+
+  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

@@ -23,6 +23,10 @@ public interface IDatabricksStatementInternal {
 
   Statement getStatement();
 
+  default boolean shouldTrackSessionVersion() {
+    return true;
+  }
+
   void allowInputStreamForVolumeOperation(boolean allowedInputStream) throws DatabricksSQLException;
 
   boolean isAllowedInputStreamForVolumeOperation() throws DatabricksSQLException;

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

@@ -120,6 +120,11 @@ 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)
+      throws SQLException {
+    return checkStatementAlive(statementId);
+  }
+
   /**
    * Fetches result for underlying statement-Id
    *