Add session version exchange to SQL Exec API client

NEXT_CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Added
 - Added `EnableThriftNativeMetadata` to request and consume supported Thrift-native SEA metadata results.
+- Added session-version exchange for SQL Exec API connections. On Lakehouse Real-Time, use
+  synchronous execution when subsequent statements depend on session changes; asynchronous
+  execution does not guarantee their visibility.
 
 ### Updated
 - `UseBoundedSeaApi` and `EnableThriftNativeMetadata` now default to `1`; when unset, activation is controlled by the server-side `enableSqlExecForJdbc` rollout flag.

src/main/java/com/databricks/jdbc/api/IDatabricksStatement.java

@@ -18,6 +18,9 @@ public interface IDatabricksStatement extends Statement {
    * long-running queries. The actual results can be retrieved later using {@link
    * #getExecutionResult()}.
    *
+   * <p>On Lakehouse Real-Time, use synchronous execution when subsequent statements depend on
+   * session changes; asynchronous execution does not guarantee their visibility.
+   *
    * @param sql The SQL command to be executed
    * @return A {@link ResultSet} handle that can be used to track and retrieve the results
    * @throws SQLException if a database access error occurs, this method is called on a closed

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,32 @@ 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 SessionVersion newSessionVersion) {
+    if (newSessionVersion == null || newSessionVersion.getVersionId() == null) {
+      return;
+    }
+    synchronized (this) {
+      if (!isSessionOpen) {
+        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 +246,7 @@ public void open() throws SQLException {
             throw e;
           }
         }
+        this.sessionVersion.set(sessionInfo == null ? null : sessionInfo.sessionVersion());
         this.isSessionOpen = true;
       }
     }
@@ -240,6 +270,7 @@ public void close() throws SQLException {
         } finally {
           // Always clean up local state
           this.sessionInfo = null;
+          this.sessionVersion.set(null);
           this.isSessionOpen = false;
         }
       }
@@ -406,6 +437,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/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,13 @@ public interface IDatabricksSession {
   @Nullable
   ImmutableSessionInfo getSessionInfo();
 
+  @Nullable
+  default SessionVersion getSessionVersion() {
+    return null;
+  }
+
+  default void updateSessionVersion(@Nullable SessionVersion sessionVersion) {}
+
   /**
    * Get the warehouse associated with the session.
    *

src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java

@@ -37,6 +37,8 @@
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.model.core.ResultManifest;
+import com.databricks.jdbc.model.core.SessionExecutionMode;
+import com.databricks.jdbc.model.core.SessionVersion;
 import com.databricks.jdbc.model.core.StatementStatus;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.sdk.WorkspaceClient;
@@ -118,7 +120,9 @@ public ImmutableSessionInfo createSession(
         schema,
         sessionConf);
     CreateSessionRequest request =
-        new CreateSessionRequest().setWarehouseId(((Warehouse) warehouse).getWarehouseId());
+        new CreateSessionRequest()
+            .setWarehouseId(((Warehouse) warehouse).getWarehouseId())
+            .setExecutionMode(SessionExecutionMode.FAST);
     if (catalog != null) {
       request.setCatalog(catalog);
     }
@@ -155,11 +159,15 @@ public ImmutableSessionInfo createSession(
       LOGGER.error(errorMessage, e);
       throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.SDK_CLIENT_ERROR);
     }
-    DatabricksThreadContextHolder.setSessionId(createSessionResponse.getSessionId());
-    return ImmutableSessionInfo.builder()
-        .computeResource(warehouse)
-        .sessionId(createSessionResponse.getSessionId())
-        .build();
+    String sessionId = createSessionResponse.getSessionId();
+    DatabricksThreadContextHolder.setSessionId(sessionId);
+    ImmutableSessionInfo.Builder sessionInfo =
+        ImmutableSessionInfo.builder().computeResource(warehouse).sessionId(sessionId);
+    SessionVersion initialVersion = createSessionResponse.getSessionVersion();
+    if (initialVersion != null && initialVersion.getVersionId() != null) {
+      sessionInfo.sessionVersion(initialVersion.getVersionId());
+    }
+    return sessionInfo.build();
   }
 
   @Override
@@ -229,6 +237,7 @@ public DatabricksResultSet executeStatement(
       }
       req.withHeaders(getHeaders("executeStatement", statementType, false, additionalHeaders));
       response = apiClient.execute(req, ExecuteStatementResponse.class);
+      updateSessionVersion(session, response.getStatus());
     } catch (IOException e) {
       String errorMessage = "Error while processing the execute statement request";
       LOGGER.error(errorMessage, e);
@@ -274,6 +283,7 @@ public DatabricksResultSet executeStatement(
         TimeoutHandler.forStatement(timeoutInSeconds, typedStatementId, this, timeoutErrorCode);
 
     StatementState responseState = response.getStatus().getState();
+    GetStatementRequest getStatementRequest = new GetStatementRequest().setStatementId(statementId);
     while (responseState == StatementState.PENDING || responseState == StatementState.RUNNING) {
       // Check for timeout
       timeoutHandler.checkTimeout();
@@ -292,9 +302,11 @@ public DatabricksResultSet executeStatement(
       }
       String getStatusPath = String.format(STATEMENT_PATH_WITH_ID, statementId);
       try {
-        Request req = new Request(Request.GET, getStatusPath, apiClient.serialize(request));
+        Request req =
+            new Request(Request.GET, getStatusPath, apiClient.serialize(getStatementRequest));
         req.withHeaders(getHeaders("getStatement"));
         response = wrapGetStatementResponse(apiClient.execute(req, GetStatementResponse.class));
+        updateSessionVersion(session, response.getStatus());
       } catch (IOException e) {
         String errorMessage = "Error while processing the get statement response";
         LOGGER.error(errorMessage, e);
@@ -391,6 +403,7 @@ public DatabricksResultSet executeStatementAsync(
       Request req = new Request(Request.POST, STATEMENT_PATH, apiClient.serialize(request));
       req.withHeaders(getHeaders("executeStatement", statementType, true));
       response = apiClient.execute(req, ExecuteStatementResponse.class);
+      updateSessionVersion(session, response.getStatus());
     } catch (IOException e) {
       String errorMessage = "Error while processing the execute statement async request";
       LOGGER.error(errorMessage, e);
@@ -747,6 +760,10 @@ private ExecuteStatementRequest getRequest(
             .setFormat(format)
             .setResultCompression(compressionCodec)
             .setParameters(parameterListItems);
+    SessionVersion sessionVersion = session.getSessionVersion();
+    if (sessionVersion != null) {
+      request.setSessionVersion(sessionVersion);
+    }
     if (executeAsync) {
       request.setWaitTimeout(ASYNC_TIMEOUT_VALUE);
     } else {
@@ -829,6 +846,12 @@ private ExecuteStatementResponse wrapGetStatementResponse(
         .setResult(getStatementResponse.getResult());
   }
 
+  private void updateSessionVersion(IDatabricksSession session, StatementStatus status) {
+    if (session != null && status != null) {
+      session.updateSessionVersion(status.getSessionVersion());
+    }
+  }
+
   /**
    * Builds actionable error messages for SSL handshake failures. Returns a generic message if the
    * error is not SSL-related.

src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionRequest.java

@@ -1,5 +1,6 @@
 package com.databricks.jdbc.model.client.sqlexec;
 
+import com.databricks.jdbc.model.core.SessionExecutionMode;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.util.Map;
 
@@ -19,6 +20,9 @@ public class CreateSessionRequest {
   @JsonProperty("session_confs")
   private Map<String, String> sessionConfigs;
 
+  @JsonProperty("execution_mode")
+  private SessionExecutionMode executionMode;
+
   public CreateSessionRequest setWarehouseId(String warehouseId) {
     this.warehouseId = warehouseId;
     return this;
@@ -54,4 +58,13 @@ public CreateSessionRequest setSessionConfigs(Map<String, String> sessionConfigs
   public Map<String, String> getSessionConfigs() {
     return sessionConfigs;
   }
+
+  public CreateSessionRequest setExecutionMode(SessionExecutionMode executionMode) {
+    this.executionMode = executionMode;
+    return this;
+  }
+
+  public SessionExecutionMode getExecutionMode() {
+    return executionMode;
+  }
 }

src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionResponse.java

@@ -1,5 +1,6 @@
 package com.databricks.jdbc.model.client.sqlexec;
 
+import com.databricks.jdbc.model.core.SessionVersion;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 /**
@@ -13,6 +14,9 @@ public class CreateSessionResponse {
   @JsonProperty("session_id")
   private String sessionId;
 
+  @JsonProperty("session_version")
+  private SessionVersion sessionVersion;
+
   public CreateSessionResponse setSessionId(String sessionId) {
     this.sessionId = sessionId;
     return this;
@@ -21,4 +25,13 @@ public CreateSessionResponse setSessionId(String sessionId) {
   public String getSessionId() {
     return sessionId;
   }
+
+  public CreateSessionResponse setSessionVersion(SessionVersion sessionVersion) {
+    this.sessionVersion = sessionVersion;
+    return this;
+  }
+
+  public SessionVersion getSessionVersion() {
+    return sessionVersion;
+  }
 }

src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java

@@ -2,6 +2,7 @@
 
 import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.model.core.Disposition;
+import com.databricks.jdbc.model.core.SessionVersion;
 import com.databricks.sdk.service.sql.ExecuteStatementRequestOnWaitTimeout;
 import com.databricks.sdk.service.sql.Format;
 import com.databricks.sdk.service.sql.StatementParameterListItem;
@@ -46,6 +47,9 @@ public class ExecuteStatementRequest {
   @JsonProperty("result_compression")
   private CompressionCodec resultCompression;
 
+  @JsonProperty("session_version")
+  private SessionVersion sessionVersion;
+
   public String getStatement() {
     return statement;
   }
@@ -86,6 +90,10 @@ public CompressionCodec getResultCompression() {
     return resultCompression;
   }
 
+  public SessionVersion getSessionVersion() {
+    return sessionVersion;
+  }
+
   // Setters
   public ExecuteStatementRequest setStatement(String statement) {
     this.statement = statement;
@@ -138,6 +146,11 @@ public ExecuteStatementRequest setParameters(Collection<StatementParameterListIt
     return this;
   }
 
+  public ExecuteStatementRequest setSessionVersion(SessionVersion sessionVersion) {
+    this.sessionVersion = sessionVersion;
+    return this;
+  }
+
   @Override
   public String toString() {
     return new ToStringer(ExecuteStatementRequest.class)
@@ -147,6 +160,7 @@ public String toString() {
         .add("parameters", parameters)
         .add("statement", statement)
         .add("sessionId", sessionId)
+        .add("sessionVersion", sessionVersion)
         .add("waitTimeout", waitTimeout)
         .add("warehouseId", warehouseId)
         .add("rowLimit", rowLimit)
@@ -161,6 +175,8 @@ public int hashCode() {
         onWaitTimeout,
         parameters,
         rowLimit,
+        sessionId,
+        sessionVersion,
         statement,
         waitTimeout,
         warehouseId);
@@ -181,6 +197,7 @@ public boolean equals(Object o) {
         && Objects.equals(statement, that.statement)
         && Objects.equals(waitTimeout, that.waitTimeout)
         && Objects.equals(sessionId, that.sessionId)
+        && Objects.equals(sessionVersion, that.sessionVersion)
         && Objects.equals(warehouseId, that.warehouseId);
   }
 }

src/main/java/com/databricks/jdbc/model/core/SessionExecutionMode.java

@@ -0,0 +1,7 @@
+package com.databricks.jdbc.model.core;
+
+public enum SessionExecutionMode {
+  SESSION_EXECUTION_MODE_UNSPECIFIED,
+  DEFAULT,
+  FAST
+}