Fix missing timeout cancellation telemetry

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.
 

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());
           }

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);

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());
           }

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.
    *

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

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> 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

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> 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

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> 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

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());