diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index c9388cda73..099e60f73a 100644 --- a/NEXT_CHANGELOG.md +++ b/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. diff --git a/src/main/java/com/databricks/jdbc/api/IDatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/IDatabricksStatement.java index e774581b0f..ea24a2dfb1 100644 --- a/src/main/java/com/databricks/jdbc/api/IDatabricksStatement.java +++ b/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()}. * + *

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 diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index bb33fa1189..3af4250bd4 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/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 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; } } diff --git a/src/main/java/com/databricks/jdbc/api/impl/SessionInfo.java b/src/main/java/com/databricks/jdbc/api/impl/SessionInfo.java index b09909b6ed..799ab83e57 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/SessionInfo.java +++ b/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 } diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java index 0229444674..296878bb17 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksSession.java +++ b/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. * diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index 0d57e0dec0..4e66f541e6 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/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. diff --git a/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionRequest.java b/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionRequest.java index 0897ede538..dbf1888d51 100644 --- a/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionRequest.java +++ b/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 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 sessionConfigs public Map getSessionConfigs() { return sessionConfigs; } + + public CreateSessionRequest setExecutionMode(SessionExecutionMode executionMode) { + this.executionMode = executionMode; + return this; + } + + public SessionExecutionMode getExecutionMode() { + return executionMode; + } } diff --git a/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionResponse.java b/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionResponse.java index 7ef7ef0ff4..724d25e009 100644 --- a/src/main/java/com/databricks/jdbc/model/client/sqlexec/CreateSessionResponse.java +++ b/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; + } } diff --git a/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java b/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java index 24968d2cf4..a1c73831c4 100644 --- a/src/main/java/com/databricks/jdbc/model/client/sqlexec/ExecuteStatementRequest.java +++ b/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 session.updateSessionVersion(new SessionVersion().setVersionId(version))); + session.updateSessionVersion(new SessionVersion().setVersionId(500L)); + + assertEquals(1000L, session.getSessionVersion().getVersionId()); + } + + @Test + public void testSessionVersionIgnoresIncompleteUpdates() throws SQLException { + DatabricksSession session = openSessionWithVersion(10L); + + session.updateSessionVersion(new SessionVersion()); + session.updateSessionVersion(null); + + assertEquals(10L, session.getSessionVersion().getVersionId()); + } + + @Test + public void testSessionVersionClearsOnCloseAndIgnoresLaterUpdates() throws SQLException { + DatabricksSession session = openSessionWithVersion(10L); + + session.close(); + assertNull(session.getSessionVersion()); + session.updateSessionVersion(new SessionVersion().setVersionId(1001L)); + assertNull(session.getSessionVersion()); + } + @Test public void testOpenRedirectedThriftSession() throws SQLException { setupWarehouse(false /* useThrift */); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index fdfb3ed81a..7e6580531f 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -33,6 +33,8 @@ import com.databricks.jdbc.model.core.ResultData; import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.core.ResultSchema; +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.core.ApiClient; @@ -64,6 +66,8 @@ public class DatabricksSdkClientTest { // Reference to MetadataOperationType to ensure import is not removed private static final MetadataOperationType SAMPLE_OP_TYPE = MetadataOperationType.GET_CATALOGS; private static final String SESSION_ID = "session_id"; + private static final long INITIAL_SESSION_VERSION = 10L; + private static final long UPDATED_SESSION_VERSION = 12L; private static final StatementId STATEMENT_ID = new StatementId("statementId"); private static final String STATEMENT = "SELECT * FROM orders WHERE user_id = ? AND shard = ? AND region_code = ? AND namespace = ?"; @@ -81,14 +85,24 @@ public class DatabricksSdkClientTest { } }; + private static SessionVersion sessionVersion(long versionId) { + return new SessionVersion().setVersionId(versionId); + } + + private static CreateSessionResponse createSessionResponse() { + return new CreateSessionResponse() + .setSessionId(SESSION_ID) + .setSessionVersion(sessionVersion(INITIAL_SESSION_VERSION)); + } + private void setupSessionMocks() throws IOException { - CreateSessionResponse response = new CreateSessionResponse().setSessionId(SESSION_ID); + CreateSessionResponse response = createSessionResponse(); when(apiClient.execute(any(Request.class), eq(CreateSessionResponse.class))) .thenReturn(response); } private void setupClientMocks(boolean includeResults, boolean async) throws IOException { - setupClientMocks(includeResults, async, new ArrayList<>(), null); + setupClientMocks(includeResults, async, new ArrayList<>(), null, false, false); } private void setupClientMocks( @@ -97,6 +111,32 @@ private void setupClientMocks( List manifestColumns, Boolean isNativeMetadataResult) throws IOException { + setupClientMocks(includeResults, async, manifestColumns, isNativeMetadataResult, false, false); + } + + private void setupClientMocks( + boolean includeResults, + boolean async, + boolean includeInitialSessionVersion, + boolean includeStatementSessionVersion) + throws IOException { + setupClientMocks( + includeResults, + async, + new ArrayList<>(), + null, + includeInitialSessionVersion, + includeStatementSessionVersion); + } + + private void setupClientMocks( + boolean includeResults, + boolean async, + List manifestColumns, + Boolean isNativeMetadataResult, + boolean includeInitialSessionVersion, + boolean includeStatementSessionVersion) + throws IOException { List params = new ArrayList<>() { { @@ -108,6 +148,9 @@ private void setupClientMocks( }; StatementStatus statementStatus = new StatementStatus().setState(StatementState.SUCCEEDED); + if (includeStatementSessionVersion) { + statementStatus.setSessionVersion(sessionVersion(UPDATED_SESSION_VERSION)); + } ExecuteStatementRequest executeStatementRequest = new ExecuteStatementRequest() .setSessionId(SESSION_ID) @@ -149,7 +192,9 @@ private void setupClientMocks( if (req.getUrl().equals(STATEMENT_PATH)) { return response; } else if (req.getUrl().equals(SESSION_PATH)) { - return new CreateSessionResponse().setSessionId(SESSION_ID); + return includeInitialSessionVersion + ? createSessionResponse() + : new CreateSessionResponse().setSessionId(SESSION_ID); } return null; }); @@ -166,6 +211,30 @@ public void testCreateSession() throws DatabricksSQLException, IOException { databricksSdkClient.createSession(warehouse, null, null, null); assertEquals(sessionInfo.sessionId(), SESSION_ID); assertEquals(sessionInfo.computeResource(), warehouse); + assertEquals(INITIAL_SESSION_VERSION, sessionInfo.sessionVersion()); + verify(apiClient) + .serialize( + argThat( + request -> + request instanceof CreateSessionRequest + && ((CreateSessionRequest) request).getExecutionMode() + == SessionExecutionMode.FAST)); + } + + @Test + public void testCreateSessionWithoutInitialSessionVersion() throws Exception { + when(apiClient.execute(any(Request.class), eq(CreateSessionResponse.class))) + .thenReturn(new CreateSessionResponse().setSessionId(SESSION_ID)); + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + + ImmutableSessionInfo sessionInfo = + databricksSdkClient.createSession(warehouse, null, null, null); + + assertEquals(SESSION_ID, sessionInfo.sessionId()); + assertNull(sessionInfo.sessionVersion()); } @Test @@ -213,7 +282,7 @@ public void testDeleteSession() throws DatabricksSQLException, IOException { @Test public void testExecuteStatement() throws Exception { - setupClientMocks(true, false); + setupClientMocks(true, false, true, true); IDatabricksConnectionContext connectionContext = DatabricksConnectionContext.parse(JDBC_URL, new Properties()); DatabricksSdkClient databricksSdkClient = @@ -235,9 +304,17 @@ public void testExecuteStatement() throws Exception { null); assertEquals(STATEMENT_ID, statement.getStatementId()); assertNotNull(resultSet.getMetaData()); + assertEquals( + UPDATED_SESSION_VERSION, connection.getSession().getSessionVersion().getVersionId()); // Verify a Request with POST method is created and executed - verify(apiClient, atLeastOnce()).serialize(any(ExecuteStatementRequest.class)); + verify(apiClient, atLeastOnce()) + .serialize( + argThat( + request -> + request instanceof ExecuteStatementRequest + && sessionVersion(INITIAL_SESSION_VERSION) + .equals(((ExecuteStatementRequest) request).getSessionVersion()))); verify(apiClient, atLeastOnce()) .execute( argThat( @@ -245,9 +322,59 @@ public void testExecuteStatement() throws Exception { eq(ExecuteStatementResponse.class)); } + @Test + public void testExecuteStatementWithoutInitialSessionVersion() throws Exception { + setupClientMocks(true, false, false, true); + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + DatabricksConnection connection = + new DatabricksConnection(connectionContext, databricksSdkClient); + connection.open(); + DatabricksStatement statement = new DatabricksStatement(connection); + + databricksSdkClient.executeStatement( + STATEMENT, + warehouse, + sqlParams, + StatementType.QUERY, + connection.getSession(), + statement, + null); + + verify(apiClient, atLeastOnce()) + .serialize( + argThat( + request -> + request instanceof ExecuteStatementRequest + && ((ExecuteStatementRequest) request).getSessionVersion() == null)); + assertEquals( + UPDATED_SESSION_VERSION, connection.getSession().getSessionVersion().getVersionId()); + + clearInvocations(apiClient); + DatabricksStatement nextStatement = new DatabricksStatement(connection); + databricksSdkClient.executeStatement( + STATEMENT, + warehouse, + sqlParams, + StatementType.QUERY, + connection.getSession(), + nextStatement, + null); + + verify(apiClient, atLeastOnce()) + .serialize( + argThat( + request -> + request instanceof ExecuteStatementRequest + && sessionVersion(UPDATED_SESSION_VERSION) + .equals(((ExecuteStatementRequest) request).getSessionVersion()))); + } + @Test public void testExecuteStatementAsync() throws Exception { - setupClientMocks(false, true); + setupClientMocks(false, true, true, true); IDatabricksConnectionContext connectionContext = DatabricksConnectionContext.parse(JDBC_URL, new Properties()); DatabricksSdkClient databricksSdkClient = @@ -263,9 +390,17 @@ public void testExecuteStatementAsync() throws Exception { STATEMENT, warehouse, sqlParams, connection.getSession(), statement); assertEquals(STATEMENT_ID, statement.getStatementId()); assertNull(resultSet.getMetaData()); + assertEquals( + UPDATED_SESSION_VERSION, connection.getSession().getSessionVersion().getVersionId()); // Verify a Request with POST method is created and executed - verify(apiClient).serialize(any(ExecuteStatementRequest.class)); + verify(apiClient) + .serialize( + argThat( + request -> + request instanceof ExecuteStatementRequest + && sessionVersion(INITIAL_SESSION_VERSION) + .equals(((ExecuteStatementRequest) request).getSessionVersion()))); verify(apiClient) .execute( argThat( @@ -418,6 +553,34 @@ public void testGetStatementResult_CancelledState_ThrowsWithHY008() throws Excep assertEquals(1008, exception.getErrorCode()); // EXECUTE_STATEMENT_CANCELLED stable code } + @Test + public void testGetStatementResultDoesNotUpdateSessionVersion() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + DatabricksConnection connection = + new DatabricksConnection(connectionContext, databricksSdkClient); + when(apiClient.execute(any(Request.class), eq(CreateSessionResponse.class))) + .thenReturn(createSessionResponse()); + connection.open(); + + GetStatementResponse response = + new GetStatementResponse() + .setStatementId(STATEMENT_ID.toSQLExecStatementId()) + .setStatus( + new StatementStatus() + .setState(StatementState.SUCCEEDED) + .setSessionVersion(sessionVersion(UPDATED_SESSION_VERSION))); + when(apiClient.execute(any(Request.class), eq(GetStatementResponse.class))) + .thenReturn(response); + + databricksSdkClient.getStatementResult(STATEMENT_ID, connection.getSession(), null); + + assertEquals( + INITIAL_SESSION_VERSION, connection.getSession().getSessionVersion().getVersionId()); + } + @Test public void testDisposition_arrowAndCloudFetchEnabled_usesExternalLinks() throws Exception { setupClientMocks(true, false); @@ -485,7 +648,7 @@ public void testExecuteStatementWithTimeout() throws Exception { new DatabricksConnection(connectionContext, databricksSdkClient); // Mock session creation - CreateSessionResponse sessionResponse = new CreateSessionResponse().setSessionId(SESSION_ID); + CreateSessionResponse sessionResponse = createSessionResponse(); when(apiClient.execute(any(Request.class), eq(CreateSessionResponse.class))) .thenReturn(sessionResponse); connection.open(); @@ -505,7 +668,10 @@ public void testExecuteStatementWithTimeout() throws Exception { .setStatus(new StatementStatus().setState(StatementState.RUNNING)); GetStatementResponse successStatementResponse = new GetStatementResponse() - .setStatus(new StatementStatus().setState(StatementState.SUCCEEDED)); + .setStatus( + new StatementStatus() + .setState(StatementState.SUCCEEDED) + .setSessionVersion(sessionVersion(UPDATED_SESSION_VERSION))); // Set up response sequence for execute() calls when(apiClient.execute( @@ -522,6 +688,13 @@ public void testExecuteStatementWithTimeout() throws Exception { .thenReturn(runningStatementResponse) .thenReturn(runningStatementResponse) .thenReturn(successStatementResponse); + String getStatementBody = + String.format("{\"statement_id\":\"%s\"}", STATEMENT_ID.toSQLExecStatementId()); + doAnswer( + invocation -> + invocation.getArgument(0) instanceof GetStatementRequest ? getStatementBody : null) + .when(apiClient) + .serialize(any()); assertDoesNotThrow( () -> @@ -533,6 +706,15 @@ public void testExecuteStatementWithTimeout() throws Exception { connection.getSession(), statement, null)); + assertEquals( + UPDATED_SESSION_VERSION, connection.getSession().getSessionVersion().getVersionId()); + + ArgumentCaptor pollRequestCaptor = ArgumentCaptor.forClass(Request.class); + verify(apiClient, times(3)) + .execute(pollRequestCaptor.capture(), eq(GetStatementResponse.class)); + assertTrue( + pollRequestCaptor.getAllValues().stream() + .allMatch(request -> getStatementBody.equals(request.getBodyString()))); // Verify no cancellation occurred due to timeout verify(apiClient, atLeastOnce()) diff --git a/src/test/java/com/databricks/jdbc/model/core/SessionVersionSerializationTest.java b/src/test/java/com/databricks/jdbc/model/core/SessionVersionSerializationTest.java new file mode 100644 index 0000000000..e09327ad35 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/model/core/SessionVersionSerializationTest.java @@ -0,0 +1,45 @@ +package com.databricks.jdbc.model.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.databricks.jdbc.model.client.sqlexec.CreateSessionRequest; +import com.databricks.jdbc.model.client.sqlexec.CreateSessionResponse; +import com.databricks.jdbc.model.client.sqlexec.ExecuteStatementRequest; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +public class SessionVersionSerializationTest { + private final ObjectMapper objectMapper = new ObjectMapper(); + + @Test + public void testSessionExecutionModeAndVersionUseProtoJsonFieldNames() throws Exception { + JsonNode createRequest = + objectMapper.valueToTree( + new CreateSessionRequest() + .setWarehouseId("warehouse") + .setExecutionMode(SessionExecutionMode.FAST)); + JsonNode executeRequest = + objectMapper.valueToTree( + new ExecuteStatementRequest() + .setSessionVersion(new SessionVersion().setVersionId(42L))); + + assertEquals("FAST", createRequest.get("execution_mode").asText()); + assertEquals(42L, executeRequest.get("session_version").get("version_id").asLong()); + } + + @Test + public void testSessionVersionsDeserializeFromCreateAndStatusResponses() throws Exception { + CreateSessionResponse createResponse = + objectMapper.readValue( + "{\"session_id\":\"session\",\"session_version\":{\"version_id\":7}}", + CreateSessionResponse.class); + StatementStatus status = + objectMapper.readValue( + "{\"state\":\"SUCCEEDED\",\"session_version\":{\"version_id\":9}}", + StatementStatus.class); + + assertEquals(7L, createResponse.getSessionVersion().getVersionId()); + assertEquals(9L, status.getSessionVersion().getVersionId()); + } +}