From c4595fe94d707b3a604f809fa96267ff28a9f3c3 Mon Sep 17 00:00:00 2001 From: Ali Abouelazm Date: Fri, 16 Oct 2026 00:15:27 -0700 Subject: [PATCH] fix: include statement IDs in Thrift failure telemetry Signed-off-by: Ali Abouelazm --- CHANGELOG.md | 1 + src/databricks/sql/backend/thrift_backend.py | 66 ++++- src/databricks/sql/exc.py | 6 +- .../sql/telemetry/telemetry_client.py | 14 +- src/databricks/sql/telemetry/utils.py | 4 +- tests/unit/test_thrift_error_telemetry.py | 264 ++++++++++++++++++ 6 files changed, 346 insertions(+), 9 deletions(-) create mode 100644 tests/unit/test_thrift_error_telemetry.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5e33347..396cbe982 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index c2ec4a8c4..f7559d913 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/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( diff --git a/src/databricks/sql/exc.py b/src/databricks/sql/exc.py index ae648fd10..7581168ac 100644 --- a/src/databricks/sql/exc.py +++ b/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): diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index fed2001cb..1dc8adc3d 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/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 diff --git a/src/databricks/sql/telemetry/utils.py b/src/databricks/sql/telemetry/utils.py index b4f74c44f..80fa4a99a 100644 --- a/src/databricks/sql/telemetry/utils.py +++ b/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 diff --git a/tests/unit/test_thrift_error_telemetry.py b/tests/unit/test_thrift_error_telemetry.py new file mode 100644 index 000000000..9d919fec2 --- /dev/null +++ b/tests/unit/test_thrift_error_telemetry.py @@ -0,0 +1,264 @@ +"""Failure events retain Thrift statement identity without exporting SQL context.""" + +import json +from unittest.mock import MagicMock, patch +from uuid import UUID + +import pytest + +from databricks.sql.auth.authenticators import AuthProvider +from databricks.sql.backend.thrift_backend import ThriftDatabricksClient +from databricks.sql.backend.types import CommandId, SessionId +from databricks.sql.exc import ( + DatabaseError, + DataError, + OperationalError, + RequestError, + ServerOperationError, +) +from databricks.sql.telemetry.telemetry_client import ( + NoopTelemetryClient, + TelemetryClient, + TelemetryClientFactory, +) +from databricks.sql.thrift_api.TCLIService import ttypes +from databricks.sql.types import SSLOptions + +STATEMENT_ID = "12345678-1234-5678-1234-567812345678" +SESSION_ID = "87654321-4321-8765-4321-876543218765" +SQL_TEXT = "SELECT 'private-statement-value' FROM missing_table" +OK = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) +ERROR = ttypes.TStatus( + statusCode=ttypes.TStatusCode.ERROR_STATUS, errorMessage="Statement failed" +) +OP_HANDLE = ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier( + guid=UUID(STATEMENT_ID).bytes, secret=b"private-operation-secret" + ), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, +) + + +@pytest.fixture +def failure_events(): + with patch("databricks.sql.telemetry.telemetry_client.UnifiedHttpClient"): + client = TelemetryClient( + telemetry_enabled=True, + session_id_hex="another-session-on-this-host", + auth_provider=AuthProvider(), + host_url="test-host", + executor=MagicMock(), + batch_size=100, + client_context=MagicMock(), + ) + with ( + patch.object( + TelemetryClientFactory, "get_telemetry_client", return_value=client + ), + patch.object(client, "_export_event") as export, + ): + yield export + + +@pytest.fixture +def backend(): + with patch( + "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True + ): + client = ThriftDatabricksClient( + "test-host", + 443, + "/test-path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + client._session_id_hex = SESSION_ID + return client + + +def execute(backend, async_op=False): + session = SessionId.from_thrift_handle( + ttypes.TSessionHandle( + sessionId=ttypes.THandleIdentifier( + guid=UUID(SESSION_ID).bytes, secret=b"secret" + ) + ) + ) + return backend.execute_command( + SQL_TEXT, session, 100, 100, False, MagicMock(), async_op=async_op + ) + + +def assert_failure_event(export, statement_id=STATEMENT_ID, session_id=SESSION_ID): + export.assert_called_once() + log = export.call_args.args[0] + event = log.entry.sql_driver_log + assert event.sql_statement_id == statement_id + assert event.session_id == session_id + serialized = log.to_json() + assert SQL_TEXT not in serialized + assert "private-statement-value" not in serialized + assert "private-operation-secret" not in serialized + assert event.error_info is not None + return json.loads(event.to_json()) + + +@pytest.mark.parametrize("direct", [True, False], ids=["direct", "polled"]) +@pytest.mark.parametrize("display_message", [None, "Statement failed"]) +def test_statement_error_has_statement_id( + backend, failure_events, direct, display_message +): + status = ttypes.TGetOperationStatusResp( + status=OK, + operationState=ttypes.TOperationState.ERROR_STATE, + errorMessage="Statement failed", + displayMessage=display_message, + diagnosticInfo=SQL_TEXT, + ) + backend._client.ExecuteStatement.return_value = ttypes.TExecuteStatementResp( + status=OK, + operationHandle=OP_HANDLE, + directResults=( + ttypes.TSparkDirectResults(operationStatus=status) if direct else None + ), + ) + backend._client.GetOperationStatus.return_value = status + with pytest.raises(ServerOperationError, match="Statement failed"): + execute(backend) + assert_failure_event(failure_events) + + +@pytest.mark.parametrize("async_op", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "field, response_type", + [ + ("operationStatus", ttypes.TGetOperationStatusResp), + ("resultSetMetadata", ttypes.TGetResultSetMetadataResp), + ("resultSet", ttypes.TFetchResultsResp), + ("closeOperation", ttypes.TCloseOperationResp), + ], +) +def test_direct_result_error_has_statement_id( + backend, failure_events, async_op, field, response_type +): + backend._client.ExecuteStatement.return_value = ttypes.TExecuteStatementResp( + status=OK, + operationHandle=OP_HANDLE, + directResults=ttypes.TSparkDirectResults( + **{field: response_type(status=ERROR)} + ), + ) + with pytest.raises(DatabaseError, match="Statement failed"): + execute(backend, async_op=async_op) + assert_failure_event(failure_events) + + +@pytest.mark.parametrize("has_handle", [True, False]) +def test_execute_rpc_error_uses_available_response_handle( + backend, failure_events, has_handle +): + backend._client.ExecuteStatement.return_value = ttypes.TExecuteStatementResp( + status=ERROR, operationHandle=OP_HANDLE if has_handle else None + ) + with pytest.raises(DatabaseError, match="Statement failed"): + execute(backend) + event = assert_failure_event(failure_events, STATEMENT_ID if has_handle else None) + if not has_handle: + assert "sql_statement_id" not in event + + +def test_poll_rpc_error_uses_request_handle(backend, failure_events): + backend._client.GetOperationStatus.return_value = ttypes.TGetOperationStatusResp( + status=ERROR + ) + with pytest.raises(DatabaseError, match="Statement failed"): + backend._poll_for_status(OP_HANDLE) + assert_failure_event(failure_events) + + +def test_poll_request_failure_uses_request_handle(backend, failure_events): + backend._client.GetOperationStatus.side_effect = RuntimeError("Connection failed") + with pytest.raises(RequestError): + backend._poll_for_status(OP_HANDLE) + assert_failure_event(failure_events) + + +def test_open_session_error_has_no_statement_id(backend, failure_events): + backend._session_id_hex = None + backend._client.OpenSession.return_value = ttypes.TOpenSessionResp(status=ERROR) + with pytest.raises(DatabaseError, match="Statement failed"): + backend.open_session({}, None, None) + event = assert_failure_event(failure_events, None, "another-session-on-this-host") + assert "sql_statement_id" not in event + + +def test_statement_error_with_disabled_telemetry(backend): + backend._client.ExecuteStatement.return_value = ttypes.TExecuteStatementResp( + status=ERROR, operationHandle=OP_HANDLE + ) + with ( + patch.object( + TelemetryClientFactory, + "get_telemetry_client", + return_value=NoopTelemetryClient(), + ), + patch.object(TelemetryClient, "_export_event") as export, + ): + with pytest.raises(DatabaseError, match="Statement failed"): + execute(backend) + export.assert_not_called() + + +def test_execute_request_failure_has_no_statement_id(backend, failure_events): + backend._client.ExecuteStatement.side_effect = RuntimeError("Connection failed") + with pytest.raises(RequestError): + execute(backend) + event = assert_failure_event(failure_events, None) + assert "sql_statement_id" not in event + + +def test_closed_statement_has_statement_id(backend, failure_events): + status = ttypes.TGetOperationStatusResp( + status=OK, operationState=ttypes.TOperationState.CLOSED_STATE + ) + with pytest.raises(DatabaseError, match="unexpectedly closed"): + backend._check_command_not_in_error_or_closed_state(OP_HANDLE, status) + assert_failure_event(failure_events) + + +def test_unexpected_result_format_has_statement_id(backend, failure_events): + backend._client.ExecuteStatement.return_value = ttypes.TExecuteStatementResp( + status=OK, + operationHandle=OP_HANDLE, + directResults=ttypes.TSparkDirectResults( + operationStatus=ttypes.TGetOperationStatusResp( + status=OK, operationState=ttypes.TOperationState.FINISHED_STATE + ), + resultSetMetadata=ttypes.TGetResultSetMetadataResp( + status=OK, resultFormat=ttypes.TSparkRowSetType.ROW_BASED_SET + ), + ), + ) + with pytest.raises(OperationalError, match="Expected results to be"): + execute(backend) + assert_failure_event(failure_events) + + +def test_fetch_offset_error_has_statement_id(backend, failure_events): + backend._client.FetchResults.return_value = ttypes.TFetchResultsResp( + status=OK, results=ttypes.TRowSet(startRowOffset=1) + ) + with pytest.raises(DataError, match="Expected results to start from"): + backend.fetch_results( + CommandId.from_thrift_handle(OP_HANDLE), + max_rows=100, + max_bytes=100, + expected_row_start_offset=0, + lz4_compressed=False, + arrow_schema_bytes=None, + description=[], + chunk_id=0, + ) + assert_failure_event(failure_events)