SNOW-2129434: Add in-band ocsp exception telemetry (snowflakedb#2406)

Author: sfc-gh-pbulawa

DESCRIPTION.md

@@ -7,6 +7,9 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
+- v3.16.1(TBD)
+  - Added in-band OCSP exception telemetry.
+
 - v3.16.0(July 04,2025)
   - Bumped numpy dependency from <2.1.0 to <=2.2.4.
   - Added Windows support for Python 3.13.

src/snowflake/connector/errors.py

@@ -36,6 +36,8 @@ def __init__(
         done_format_msg: bool | None = None,
         connection: SnowflakeConnection | None = None,
         cursor: SnowflakeCursor | None = None,
+        errtype: TelemetryField = TelemetryField.SQL_EXCEPTION,
+        send_telemetry: bool = True,
     ) -> None:
         super().__init__(msg)
         self.msg = msg
@@ -44,6 +46,8 @@ def __init__(
         self.sqlstate = sqlstate or "n/a"
         self.sfqid = sfqid
         self.query = query
+        self.errtype = errtype
+        self.send_telemetry = send_telemetry
 
         if self.msg:
             # TODO: If there's a message then check to see if errno (and maybe sqlstate)
@@ -74,7 +78,9 @@ def __init__(
 
         # We want to skip the last frame/line in the traceback since it is the current frame
         self.telemetry_traceback = self.generate_telemetry_stacktrace()
-        self.exception_telemetry(msg, cursor, connection)
+
+        if self.send_telemetry:
+            self.exception_telemetry(msg, cursor, connection)
 
     def __repr__(self) -> str:
         return self.__str__()
@@ -131,6 +137,8 @@ def generate_telemetry_exception_data(
             telemetry_data_dict[TelemetryField.KEY_REASON.value] = telemetry_msg
         if self.errno:
             telemetry_data_dict[TelemetryField.KEY_ERROR_NUMBER.value] = str(self.errno)
+        if self.msg:
+            telemetry_data_dict[TelemetryField.KEY_ERROR_MESSAGE.value] = self.msg
 
         return telemetry_data_dict
 
@@ -147,9 +155,7 @@ def send_exception_telemetry(
             and not connection._telemetry.is_closed
         ):
             # Send with in-band telemetry
-            telemetry_data[TelemetryField.KEY_TYPE.value] = (
-                TelemetryField.SQL_EXCEPTION.value
-            )
+            telemetry_data[TelemetryField.KEY_TYPE.value] = self.errtype.value
             telemetry_data[TelemetryField.KEY_SOURCE.value] = connection.application
             telemetry_data[TelemetryField.KEY_EXCEPTION.value] = self.__class__.__name__
             ts = get_time_millis()
@@ -415,9 +421,14 @@ def telemetry_msg(self) -> str:
 class RevocationCheckError(OperationalError):
     """Exception for errors during certificate revocation check."""
 
-    # We already send OCSP exception events
-    def exception_telemetry(self, msg, cursor, connection) -> None:
-        pass
+    def __init__(self, **kwargs) -> None:
+        send_telemetry = kwargs.pop("send_telemetry", False)
+        Error.__init__(
+            self,
+            errtype=TelemetryField.OCSP_EXCEPTION,
+            send_telemetry=send_telemetry,
+            **kwargs,
+        )
 
 
 # internal errors

src/snowflake/connector/network.py

@@ -81,6 +81,7 @@
     OtherHTTPRetryableError,
     ProgrammingError,
     RefreshTokenError,
+    RevocationCheckError,
     ServiceUnavailableError,
     TooManyRequests,
 )
@@ -991,6 +992,9 @@ def _request_exec_wrapper(
                 raise RetryRequest(err_msg)
             self._handle_unknown_error(method, full_url, headers, data, conn)
             return {}
+        except RevocationCheckError as rce:
+            rce.exception_telemetry(rce.msg, None, self._connection)
+            raise rce
         except RetryRequest as e:
             cause = e.args[0]
             if no_retry:

src/snowflake/connector/ocsp_snowflake.py

@@ -138,7 +138,7 @@ def deserialize_exception(exception_dict: dict | None) -> Exception | None:
                     f" the original error error class and message are {exc_class} and {exception_dict['msg']}"
                 )
                 return RevocationCheckError(
-                    f"Got error {str(deserialize_exc)} while deserializing ocsp cache, please try "
+                    msg=f"Got error {str(deserialize_exc)} while deserializing ocsp cache, please try "
                     f"cleaning up the "
                     f"OCSP cache under directory {OCSP_RESPONSE_VALIDATION_CACHE.file_path}",
                     errno=ER_OCSP_RESPONSE_LOAD_FAILURE,

src/snowflake/connector/telemetry.py

@@ -25,6 +25,7 @@ class TelemetryField(Enum):
     TIME_DOWNLOADING_CHUNKS = "client_time_downloading_chunks"
     TIME_PARSING_CHUNKS = "client_time_parsing_chunks"
     SQL_EXCEPTION = "client_sql_exception"
+    OCSP_EXCEPTION = "client_ocsp_exception"
     GET_PARTITIONS_USED = "client_get_partitions_used"
     EMPTY_SEQ_INTERPOLATION = "client_pyformat_empty_seq_interpolation"
     # fetch_pandas_* usage

test/unit/test_telemetry.py

@@ -1,10 +1,17 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+from unittest import mock
 from unittest.mock import Mock
 
+import pytest
+
 import snowflake.connector.telemetry
 from snowflake.connector.description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
+from src.snowflake.connector.errorcode import ER_OCSP_RESPONSE_UNAVAILABLE
+from src.snowflake.connector.errors import RevocationCheckError
+from src.snowflake.connector.network import SnowflakeRestful
+from src.snowflake.connector.telemetry import TelemetryData, TelemetryField
 
 
 def test_telemetry_data_to_dict():
@@ -235,3 +242,103 @@ def test_generate_telemetry_data():
         }
         and telemetry_data.timestamp == 123
     )
+
+
+def test_raising_error_generates_telemetry_event_when_connection_is_present():
+    mock_connection = get_mocked_telemetry_connection()
+
+    with pytest.raises(RevocationCheckError):
+        raise RevocationCheckError(
+            msg="Response unavailable",
+            errno=ER_OCSP_RESPONSE_UNAVAILABLE,
+            connection=mock_connection,
+            send_telemetry=True,
+        )
+
+    mock_connection._log_telemetry.assert_called_once()
+    assert_telemetry_data_for_revocation_check_error(
+        mock_connection._log_telemetry.call_args[0][0]
+    )
+
+
+def test_raising_error_with_send_telemetry_off_does_not_generate_telemetry_event_when_connection_is_present():
+    mock_connection = get_mocked_telemetry_connection()
+
+    with pytest.raises(RevocationCheckError):
+        raise RevocationCheckError(
+            msg="Response unavailable",
+            errno=ER_OCSP_RESPONSE_UNAVAILABLE,
+            connection=mock_connection,
+            send_telemetry=False,
+        )
+
+    mock_connection._log_telemetry.assert_not_called()
+
+
+def test_request_throws_revocation_check_error():
+    retry_ctx = Mock()
+    retry_ctx.current_retry_count = 0
+    retry_ctx.timeout = 10
+    retry_ctx.add_retry_params.return_value = "https://example.com"
+
+    mock_connection = get_mocked_telemetry_connection()
+
+    with mock.patch.object(SnowflakeRestful, "_request_exec") as _request_exec_mocked:
+        _request_exec_mocked.side_effect = RevocationCheckError(
+            msg="Response unavailable", errno=ER_OCSP_RESPONSE_UNAVAILABLE
+        )
+        mock_restful = SnowflakeRestful(connection=mock_connection)
+        with pytest.raises(RevocationCheckError):
+            mock_restful._request_exec_wrapper(
+                None,
+                None,
+                None,
+                None,
+                None,
+                retry_ctx,
+            )
+        mock_restful._connection._log_telemetry.assert_called_once()
+        assert_telemetry_data_for_revocation_check_error(
+            mock_connection._log_telemetry.call_args[0][0]
+        )
+
+
+def get_mocked_telemetry_connection(telemetry_enabled: bool = True) -> Mock:
+    mock_connection = Mock()
+    mock_connection.application = "test_application"
+    mock_connection.telemetry_enabled = telemetry_enabled
+    mock_connection.is_closed = False
+
+    mock_connection._log_telemetry = Mock()
+
+    mock_telemetry = Mock()
+    mock_telemetry.is_closed = False
+    mock_connection._telemetry = mock_telemetry
+
+    return mock_connection
+
+
+def assert_telemetry_data_for_revocation_check_error(telemetry_data: TelemetryData):
+    assert telemetry_data.message[TelemetryField.KEY_DRIVER_TYPE.value] == CLIENT_NAME
+    assert (
+        telemetry_data.message[TelemetryField.KEY_DRIVER_VERSION.value]
+        == SNOWFLAKE_CONNECTOR_VERSION
+    )
+    assert telemetry_data.message[TelemetryField.KEY_SOURCE.value] == "test_application"
+    assert (
+        telemetry_data.message[TelemetryField.KEY_TYPE.value]
+        == TelemetryField.OCSP_EXCEPTION.value
+    )
+    assert telemetry_data.message[TelemetryField.KEY_ERROR_NUMBER.value] == str(
+        ER_OCSP_RESPONSE_UNAVAILABLE
+    )
+    assert (
+        telemetry_data.message[TelemetryField.KEY_EXCEPTION.value]
+        == "RevocationCheckError"
+    )
+    assert (
+        "Response unavailable"
+        in telemetry_data.message[TelemetryField.KEY_ERROR_MESSAGE.value]
+    )
+    assert TelemetryField.KEY_STACKTRACE.value in telemetry_data.message
+    assert TelemetryField.KEY_REASON.value in telemetry_data.message

test/unit/test_telemetry_oob.py

@@ -10,7 +10,7 @@
 TEST_RACE_CONDITION_THREAD_COUNT = 2
 TEST_RACE_CONDITION_DELAY_SECONDS = 1
 telemetry_data = {}
-exception = RevocationCheckError("Test OCSP Revocation error")
+exception = RevocationCheckError(msg="Test OCSP Revocation error")
 event_type = "Test OCSP Exception"
 stack_trace = [
     "Traceback (most recent call last):\n",