diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 6c27fbc1ac..5ddd50b89c 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -10,6 +10,7 @@ - Updated bundled Jackson, lz4-java, Netty, and Apache HttpComponents Client and Core dependencies to patched versions to address security findings. ### Fixed +- Added SEA statement status-poll telemetry and fixed successful timeout cancellation failing to flush buffered polling details on SEA and Thrift. - Invalid or incomplete Databricks JDBC URLs now fail with a descriptive `DatabricksSQLException` instead of leaking a `NullPointerException` when required connection parameters are missing. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java index edc85e454d..9ac88d8d57 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java @@ -5,6 +5,7 @@ import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import com.databricks.jdbc.telemetry.TelemetryHelper; import java.util.concurrent.TimeUnit; /** Utility class to handle statement execution timeouts. */ @@ -93,7 +94,15 @@ public static TimeoutHandler forStatement( "Statement ID: " + statementId, () -> { try { + long cancelStartTime = System.nanoTime(); client.cancelStatement(statementId); + long cancelLatencyMillis = + TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - cancelStartTime); + TelemetryHelper.recordOperationLatency( + client.getConnectionContext(), + statementId.toSQLExecStatementId(), + cancelLatencyMillis, + "cancelStatement"); } catch (Exception e) { LOGGER.warn("Cancel statement on timeout failed: " + e.getMessage()); } 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..5421ddbc86 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 @@ -39,6 +39,7 @@ import com.databricks.jdbc.model.core.ResultManifest; import com.databricks.jdbc.model.core.StatementStatus; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.WorkspaceClient; import com.databricks.sdk.core.ApiClient; import com.databricks.sdk.core.DatabricksConfig; @@ -53,6 +54,7 @@ import java.sql.SQLException; import java.time.Instant; import java.util.*; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.Stream; import javax.net.ssl.SSLHandshakeException; @@ -294,7 +296,12 @@ public DatabricksResultSet executeStatement( try { Request req = new Request(Request.GET, getStatusPath, apiClient.serialize(request)); req.withHeaders(getHeaders("getStatement")); + long operationStatusStartTime = System.nanoTime(); response = wrapGetStatementResponse(apiClient.execute(req, GetStatementResponse.class)); + long operationStatusLatencyMillis = + TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - operationStatusStartTime); + TelemetryHelper.recordGetOperationStatus( + connectionContext, statementId, operationStatusLatencyMillis); } catch (IOException e) { String errorMessage = "Error while processing the get statement response"; LOGGER.error(errorMessage, e); diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 0518c61be9..3be11f4b46 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -330,7 +330,10 @@ private TGetOperationStatusResp pollTillOperationFinished( TimeoutHandler timeoutHandler = getTimeoutHandler( - response, timeoutInSeconds, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + response, + statementId, + timeoutInSeconds, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); // Polling until query operation state is finished long pollingStartTime = System.nanoTime(); @@ -1019,6 +1022,7 @@ void setServerProtocolVersion(TProtocolVersion protocolVersion) { private TimeoutHandler getTimeoutHandler( TExecuteStatementResp response, + StatementId statementId, int timeoutInSeconds, DatabricksDriverErrorCode internalErrorCode) { final TOperationHandle operationHandle = response.getOperationHandle(); @@ -1029,7 +1033,15 @@ private TimeoutHandler getTimeoutHandler( () -> { try { LOGGER.debug("Canceling operation due to timeout: {}", operationHandle); + long cancelStartTime = System.nanoTime(); cancelOperation(new TCancelOperationReq().setOperationHandle(operationHandle)); + long cancelLatencyMillis = + TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - cancelStartTime); + TelemetryHelper.recordOperationLatency( + connectionContext, + statementId.toSQLExecStatementId(), + cancelLatencyMillis, + "cancelStatement"); } catch (Exception e) { LOGGER.warn("Failed to cancel operation on timeout: {}", e.getMessage()); } diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 4e45cfe131..3fdf6b8219 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -505,6 +505,23 @@ public static void recordGetOperationStatus( } } + /** Records operation latency for a statement. Silently ignores errors. */ + public static void recordOperationLatency( + IDatabricksConnectionContext connectionContext, + String statementId, + long latencyMillis, + String methodName) { + try { + if (connectionContext != null) { + TelemetryCollectorManager.getInstance() + .getOrCreateCollector(connectionContext) + .recordOperationLatency(statementId, latencyMillis, methodName); + } + } catch (Exception e) { + LOGGER.trace("Error recording operation latency telemetry: {}", e.getMessage()); + } + } + /** * Records chunk download latency. Silently ignores errors. * diff --git a/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java index 2b825a410b..093267bda3 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java +++ b/src/main/java/com/databricks/jdbc/telemetry/latency/TelemetryCollector.java @@ -73,7 +73,11 @@ public void recordTotalChunks(StatementId statementId, long totalChunks) { public void recordOperationLatency(long latencyMillis, String methodName) { // It is possible that statement ID is not present in case of openSession. In which case, we // send telemetry latency log without the statement ID - String statementId = DatabricksThreadContextHolder.getStatementId(); + recordOperationLatency( + DatabricksThreadContextHolder.getStatementId(), latencyMillis, methodName); + } + + public void recordOperationLatency(String statementId, long latencyMillis, String methodName) { OperationType operationType = TelemetryHelper.mapMethodToOperationType(methodName); if (isTelemetryCollected(statementId) && isCloseOperation(operationType)) { // This is terminal state, we will have to export all data corresponding to the statementID diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java index ae75478fc3..7a55ca2fa1 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java @@ -3,14 +3,17 @@ import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import com.databricks.jdbc.telemetry.TelemetryHelper; import java.lang.reflect.Field; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) @@ -22,6 +25,8 @@ class TimeoutHandlerTest { @Mock private StatementId mockStatementId; + @Mock private IDatabricksConnectionContext mockConnectionContext; + @Test void testNoTimeout() { // Create handler with no timeout (0 seconds) @@ -138,6 +143,8 @@ void testNullTimeoutAction() throws Exception { @Test void testForStatementFactory() throws Exception { when(mockStatementId.toString()).thenReturn("test-statement-id"); + when(mockStatementId.toSQLExecStatementId()).thenReturn("test-statement-id"); + when(mockClient.getConnectionContext()).thenReturn(mockConnectionContext); // Create handler with factory method TimeoutHandler handler = @@ -159,11 +166,20 @@ void testForStatementFactory() throws Exception { long currentTime = System.currentTimeMillis(); startTimeField.set(handler, currentTime - TimeUnit.SECONDS.toMillis(6)); // 6 seconds ago - // This should throw a DatabricksTimeoutException - assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); - - // Verify client.cancelStatement was called - verify(mockClient, times(1)).cancelStatement(mockStatementId); + try (MockedStatic telemetryHelper = mockStatic(TelemetryHelper.class)) { + // This should throw a DatabricksTimeoutException + assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); + + // Verify client.cancelStatement was called and its terminal telemetry was recorded. + verify(mockClient, times(1)).cancelStatement(mockStatementId); + telemetryHelper.verify( + () -> + TelemetryHelper.recordOperationLatency( + eq(mockConnectionContext), + eq("test-statement-id"), + anyLong(), + eq("cancelStatement"))); + } } @Test 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..68cf931a7a 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 @@ -35,6 +35,7 @@ import com.databricks.jdbc.model.core.ResultSchema; import com.databricks.jdbc.model.core.StatementStatus; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.ApiClient; import com.databricks.sdk.core.DatabricksError; import com.databricks.sdk.core.http.Request; @@ -52,6 +53,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) @@ -602,23 +604,37 @@ public void testExecuteStatementWithTimeoutExpired() throws Exception { // Verify that the timeout exception (1 second) is thrown due to repeated polling, where each // poll occurs at an interval of 1 second - DatabricksTimeoutException exception = - assertThrows( - DatabricksTimeoutException.class, - () -> - databricksSdkClient.executeStatement( - STATEMENT, - warehouse, - sqlParams, - StatementType.QUERY, - connection.getSession(), - statement, - null)); - - assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); - - // Verify cancel was called - verify(databricksSdkClient).cancelStatement(eq(STATEMENT_ID)); + try (MockedStatic telemetryHelper = mockStatic(TelemetryHelper.class)) { + DatabricksTimeoutException exception = + assertThrows( + DatabricksTimeoutException.class, + () -> + databricksSdkClient.executeStatement( + STATEMENT, + warehouse, + sqlParams, + StatementType.QUERY, + connection.getSession(), + statement, + null)); + + assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); + + // Verify cancel was called and its terminal telemetry was recorded. + verify(databricksSdkClient).cancelStatement(eq(STATEMENT_ID)); + telemetryHelper.verify( + () -> + TelemetryHelper.recordOperationLatency( + eq(connectionContext), + eq(STATEMENT_ID.toSQLExecStatementId()), + anyLong(), + eq("cancelStatement"))); + telemetryHelper.verify( + () -> + TelemetryHelper.recordGetOperationStatus( + eq(connectionContext), eq(STATEMENT_ID.toSQLExecStatementId()), anyLong()), + atLeastOnce()); + } } @Test diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 7f58270c15..8d31bf6b8f 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -21,6 +21,7 @@ import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.exception.DatabricksValidationException; import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; import java.sql.SQLException; @@ -948,15 +949,24 @@ void testExecuteWithTimeoutExpired() throws TException, SQLException { // The execute method should throw a timeout exception since the operation does not complete // within 1 second. The polling interval is 1 second, and multiple polling attempts are made - DatabricksTimeoutException exception = - assertThrows( - DatabricksTimeoutException.class, - () -> accessor.execute(request, parentStatement, session, StatementType.SQL)); - - assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); - - // Verify that cancel was called - verify(thriftClient).CancelOperation(any(TCancelOperationReq.class)); + try (MockedStatic telemetryHelper = mockStatic(TelemetryHelper.class)) { + DatabricksTimeoutException exception = + assertThrows( + DatabricksTimeoutException.class, + () -> accessor.execute(request, parentStatement, session, StatementType.SQL)); + + assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); + + // Verify that cancel was called and its terminal telemetry was recorded. + verify(thriftClient).CancelOperation(any(TCancelOperationReq.class)); + telemetryHelper.verify( + () -> + TelemetryHelper.recordOperationLatency( + eq(connectionContext), + eq(StatementId.deserialize(TEST_STMT_ID).toSQLExecStatementId()), + anyLong(), + eq("cancelStatement"))); + } } @Test diff --git a/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java index b071b90eff..8e0725b798 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/latency/TelemetryCollectorTest.java @@ -10,6 +10,8 @@ import com.databricks.jdbc.model.telemetry.latency.ChunkDetails; import com.databricks.jdbc.model.telemetry.latency.OperationType; import com.databricks.jdbc.telemetry.TelemetryHelper; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -90,6 +92,24 @@ void testRecordOperationLatency_WithCloseOperation() { } } + @Test + void testCancelExportsAccumulatedPollingDetailsAndClearsTracker() { + handler.recordGetOperationStatus(TEST_STATEMENT_ID, 1000L); + handler.recordGetOperationStatus(TEST_STATEMENT_ID, 250L); + StatementTelemetryDetails pendingDetails = + handler.getOrCreateTelemetryDetails(TEST_STATEMENT_ID); + DatabricksThreadContextHolder.setStatementId("different-statement-id"); + + handler.recordOperationLatency(TEST_STATEMENT_ID, 100L, "cancelStatement"); + + JsonNode operationDetail = new ObjectMapper().valueToTree(pendingDetails.getOperationDetail()); + assertEquals(2L, operationDetail.get("n_operation_status_calls").asLong()); + assertEquals(1250L, operationDetail.get("operation_status_latency_millis").asLong()); + assertEquals("CANCEL_STATEMENT", operationDetail.get("operation_type").asText()); + assertEquals(100L, pendingDetails.getOperationLatencyMillis()); + assertFalse(handler.isTelemetryCollected(TEST_STATEMENT_ID)); + } + @Test void testCollectorStoresConnectionContext() { assertSame(mockContext, handler.getConnectionContext());