fix: include statement IDs in Thrift failure telemetry

Signed-off-by: Ali Abouelazm <aliazm419@gmail.com>

Author: AliAbouelazm

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 # Unreleased
+- Include the available SQL statement ID in Thrift failure telemetry so errors can be correlated with the failed statement.
 - Add the `escape_catalog_name` connection parameter (default `False`) to opt in to escaping unescaped `_` in metadata catalog names on Thrift and kernel connections. Existing escapes and `%` wildcards are preserved.
 
 # 4.6.0 (2026-09-24)

src/databricks/sql/backend/thrift_backend.py

@@ -303,7 +303,19 @@ def _initialize_retry_args(self, kwargs):
             )
 
     @staticmethod
-    def _check_response_for_error(response, host_url=None, detect_reyden=False):
+    def _get_statement_id(op_handle):
+        operation_id = getattr(op_handle, "operationId", None)
+        guid = getattr(operation_id, "guid", None)
+        return guid_to_hex_id(guid) if guid is not None else None
+
+    @staticmethod
+    def _check_response_for_error(
+        response,
+        host_url=None,
+        detect_reyden=False,
+        session_id_hex=None,
+        sql_statement_id=None,
+    ):
         if response.status and response.status.statusCode in [
             ttypes.TStatusCode.ERROR_STATUS,
             ttypes.TStatusCode.INVALID_HANDLE_STATUS,
@@ -323,6 +335,11 @@ def _check_response_for_error(response, host_url=None, detect_reyden=False):
             raise DatabaseError(
                 response.status.errorMessage,
                 host_url=host_url,
+                session_id_hex=session_id_hex,
+                sql_statement_id=ThriftDatabricksClient._get_statement_id(
+                    getattr(response, "operationHandle", None)
+                )
+                or sql_statement_id,
             )
 
     @staticmethod
@@ -378,6 +395,10 @@ def _handle_request_error(self, error_info, attempt, elapsed):
                 full_error_info_context,
                 self._host,
                 error_info.error,
+                session_id_hex=self._session_id_hex,
+                sql_statement_id=self._get_statement_id(
+                    getattr(error_info.request, "operationHandle", None)
+                ),
             )
             logger.info(network_request_error.message_with_context())
 
@@ -558,6 +579,10 @@ def attempt_request(attempt):
                     response,
                     self._host,
                     detect_reyden=getattr(method, "__name__", None) == "OpenSession",
+                    session_id_hex=self._session_id_hex,
+                    sql_statement_id=self._get_statement_id(
+                        getattr(request, "operationHandle", None)
+                    ),
                 )
                 return response
 
@@ -680,6 +705,8 @@ def _check_command_not_in_error_or_closed_state(
                         "diagnostic-info": get_operations_resp.diagnosticInfo,
                     },
                     host_url=self._host,
+                    session_id_hex=self._session_id_hex,
+                    sql_statement_id=self._get_statement_id(op_handle),
                 )
             else:
                 raise ServerOperationError(
@@ -690,6 +717,8 @@ def _check_command_not_in_error_or_closed_state(
                         "diagnostic-info": None,
                     },
                     host_url=self._host,
+                    session_id_hex=self._session_id_hex,
+                    sql_statement_id=self._get_statement_id(op_handle),
                 )
         elif get_operations_resp.operationState == ttypes.TOperationState.CLOSED_STATE:
             raise DatabaseError(
@@ -701,6 +730,8 @@ def _check_command_not_in_error_or_closed_state(
                     and guid_to_hex_id(op_handle.operationId.guid)
                 },
                 host_url=self._host,
+                session_id_hex=self._session_id_hex,
+                sql_statement_id=self._get_statement_id(op_handle),
             )
 
     def _poll_for_status(self, op_handle):
@@ -862,6 +893,8 @@ def _results_message_to_execute_response(
                     ]
                 ),
                 host_url=self._host,
+                session_id_hex=self._session_id_hex,
+                sql_statement_id=self._get_statement_id(resp.operationHandle),
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
@@ -1019,27 +1052,40 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         return state
 
     @staticmethod
-    def _check_direct_results_for_error(t_spark_direct_results, host_url=None):
+    def _check_direct_results_for_error(
+        t_spark_direct_results,
+        host_url=None,
+        session_id_hex=None,
+        sql_statement_id=None,
+    ):
         if t_spark_direct_results:
             if t_spark_direct_results.operationStatus:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.operationStatus,
                     host_url,
+                    session_id_hex=session_id_hex,
+                    sql_statement_id=sql_statement_id,
                 )
             if t_spark_direct_results.resultSetMetadata:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.resultSetMetadata,
                     host_url,
+                    session_id_hex=session_id_hex,
+                    sql_statement_id=sql_statement_id,
                 )
             if t_spark_direct_results.resultSet:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.resultSet,
                     host_url,
+                    session_id_hex=session_id_hex,
+                    sql_statement_id=sql_statement_id,
                 )
             if t_spark_direct_results.closeOperation:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.closeOperation,
                     host_url,
+                    session_id_hex=session_id_hex,
+                    sql_statement_id=sql_statement_id,
                 )
 
     def execute_command(
@@ -1327,7 +1373,12 @@ def _handle_execute_response(self, resp, cursor):
             raise ValueError(f"Invalid Thrift handle: {resp.operationHandle}")
 
         cursor.active_command_id = command_id
-        self._check_direct_results_for_error(resp.directResults, self._host)
+        self._check_direct_results_for_error(
+            resp.directResults,
+            self._host,
+            session_id_hex=self._session_id_hex,
+            sql_statement_id=command_id.to_hex_guid(),
+        )
 
         final_operation_state, final_status_resp = self._wait_until_command_done(
             resp.operationHandle,
@@ -1351,7 +1402,12 @@ def _handle_execute_response_async(self, resp, cursor):
             raise ValueError(f"Invalid Thrift handle: {resp.operationHandle}")
 
         cursor.active_command_id = command_id
-        self._check_direct_results_for_error(resp.directResults, self._host)
+        self._check_direct_results_for_error(
+            resp.directResults,
+            self._host,
+            session_id_hex=self._session_id_hex,
+            sql_statement_id=command_id.to_hex_guid(),
+        )
 
     def fetch_results(
         self,
@@ -1390,6 +1446,8 @@ def fetch_results(
                     expected_row_start_offset, resp.results.startRowOffset
                 ),
                 host_url=self._host,
+                session_id_hex=self._session_id_hex,
+                sql_statement_id=command_id.to_hex_guid(),
             )
 
         queue = ThriftResultSetQueueFactory.build_queue(

src/databricks/sql/exc.py

@@ -19,6 +19,7 @@ def __init__(
         host_url=None,
         *args,
         session_id_hex=None,
+        sql_statement_id=None,
         **kwargs,
     ):
         super().__init__(message, *args, **kwargs)
@@ -33,7 +34,10 @@ def __init__(
                 host_url=host_url
             )
             telemetry_client.export_failure_log(
-                error_name, self.message, session_id=session_id_hex
+                error_name,
+                self.message,
+                session_id=session_id_hex,
+                sql_statement_id=sql_statement_id,
             )
 
     def __str__(self):

src/databricks/sql/telemetry/telemetry_client.py

@@ -161,7 +161,9 @@ def export_initial_telemetry_log(
     ):
         pass
 
-    def export_failure_log(self, error_name, error_message, session_id=None):
+    def export_failure_log(
+        self, error_name, error_message, session_id=None, sql_statement_id=None
+    ):
         pass
 
     def export_latency_log(
@@ -414,9 +416,15 @@ def export_initial_telemetry_log(
         self._user_agent = user_agent
         self._export_telemetry_log(session_id=session_id)
 
-    def export_failure_log(self, error_name, error_message, session_id=None):
+    def export_failure_log(
+        self, error_name, error_message, session_id=None, sql_statement_id=None
+    ):
         error_info = DriverErrorInfo(error_name=error_name, stack_trace=error_message)
-        self._export_telemetry_log(session_id=session_id, error_info=error_info)
+        self._export_telemetry_log(
+            session_id=session_id,
+            sql_statement_id=sql_statement_id,
+            error_info=error_info,
+        )
 
     def export_latency_log(
         self, latency_ms, sql_execution_event, sql_statement_id, session_id=None

src/databricks/sql/telemetry/utils.py

@@ -19,7 +19,9 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
         pass
 
     @abstractmethod
-    def export_failure_log(self, error_name, error_message):
+    def export_failure_log(
+        self, error_name, error_message, session_id=None, sql_statement_id=None
+    ):
         logger.debug("subclass must implement export_failure_log")
         pass