SNOW-654808, SNOW-652000: Add field source to basic telemetry data and refactor Telemetry data (#1243)

Author: sfc-gh-aling

src/snowflake/connector/connection.py

@@ -83,12 +83,7 @@
     SnowflakeRestful,
 )
 from .sqlstate import SQLSTATE_CONNECTION_NOT_EXISTS, SQLSTATE_FEATURE_NOT_SUPPORTED
-from .telemetry import (
-    TelemetryClient,
-    TelemetryData,
-    TelemetryField,
-    generate_telemetry_data,
-)
+from .telemetry import TelemetryClient, TelemetryData, TelemetryField
 from .telemetry_oob import TelemetryService
 from .time_util import HeartBeatTimer, get_time_millis
 from .util_text import construct_hostname, parse_account, split_statements
@@ -1565,16 +1560,12 @@ def _log_telemetry_imported_packages(self) -> None:
             }
             ts = get_time_millis()
             self._log_telemetry(
-                TelemetryData(
-                    generate_telemetry_data(
-                        from_dict={
-                            TelemetryField.KEY_TYPE.value: TelemetryField.IMPORTED_PACKAGES.value,
-                            TelemetryField.KEY_SOURCE.value: self.application
-                            if self.application
-                            else CLIENT_NAME,
-                            TelemetryField.KEY_VALUE.value: str(imported_modules),
-                        }
-                    ),
-                    ts,
+                TelemetryData.from_telemetry_data_dict(
+                    from_dict={
+                        TelemetryField.KEY_TYPE.value: TelemetryField.IMPORTED_PACKAGES.value,
+                        TelemetryField.KEY_VALUE.value: str(imported_modules),
+                    },
+                    timestamp=ts,
+                    connection=self,
                 )
             )

src/snowflake/connector/cursor.py

@@ -39,7 +39,6 @@
     FileTransferType,
     QueryStatus,
 )
-from .description import CLIENT_NAME
 from .errorcode import (
     ER_CURSOR_IS_CLOSED,
     ER_FAILED_PROCESSING_PYFORMAT,
@@ -62,7 +61,7 @@
 from .file_transfer_agent import SnowflakeFileTransferAgent
 from .options import installed_pandas, pandas
 from .sqlstate import SQLSTATE_FEATURE_NOT_SUPPORTED
-from .telemetry import TelemetryData, TelemetryField, generate_telemetry_data
+from .telemetry import TelemetryData, TelemetryField
 from .time_util import get_time_millis
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -1214,18 +1213,14 @@ def _log_telemetry_job_data(
         ts = get_time_millis()
         try:
             self._connection._log_telemetry(
-                TelemetryData(
-                    generate_telemetry_data(
-                        from_dict={
-                            TelemetryField.KEY_TYPE.value: telemetry_field.value,
-                            TelemetryField.KEY_SOURCE.value: self._connection.application
-                            if self._connection
-                            else CLIENT_NAME,
-                            TelemetryField.KEY_SFQID.value: self._sfqid,
-                            TelemetryField.KEY_VALUE.value: value,
-                        }
-                    ),
-                    ts,
+                TelemetryData.from_telemetry_data_dict(
+                    from_dict={
+                        TelemetryField.KEY_TYPE.value: telemetry_field.value,
+                        TelemetryField.KEY_SFQID.value: self._sfqid,
+                        TelemetryField.KEY_VALUE.value: value,
+                    },
+                    timestamp=ts,
+                    connection=self._connection,
                 )
             )
         except AttributeError:

src/snowflake/connector/errors.py

@@ -14,7 +14,7 @@
 
 from .compat import BASE_EXCEPTION_CLASS
 from .secret_detector import SecretDetector
-from .telemetry import TelemetryData, TelemetryField, generate_telemetry_data
+from .telemetry import TelemetryData, TelemetryField
 from .telemetry_oob import TelemetryService
 from .time_util import get_time_millis
 
@@ -119,25 +119,22 @@ def telemetry_msg(self) -> str | None:
 
     def generate_telemetry_exception_data(self) -> dict[str, str]:
         """Generate the data to send through telemetry."""
-
-        telemetry_data = generate_telemetry_data(
-            from_dict={
-                TelemetryField.KEY_STACKTRACE.value: SecretDetector.mask_secrets(
-                    self.telemetry_traceback
-                )
-            }
-        )
+        telemetry_data_dict = {
+            TelemetryField.KEY_STACKTRACE.value: SecretDetector.mask_secrets(
+                self.telemetry_traceback
+            )
+        }
         telemetry_msg = self.telemetry_msg()
         if self.sfqid:
-            telemetry_data[TelemetryField.KEY_SFQID.value] = self.sfqid
+            telemetry_data_dict[TelemetryField.KEY_SFQID.value] = self.sfqid
         if self.sqlstate:
-            telemetry_data[TelemetryField.KEY_SQLSTATE.value] = self.sqlstate
+            telemetry_data_dict[TelemetryField.KEY_SQLSTATE.value] = self.sqlstate
         if telemetry_msg:
-            telemetry_data[TelemetryField.KEY_REASON.value] = telemetry_msg
+            telemetry_data_dict[TelemetryField.KEY_REASON.value] = telemetry_msg
         if self.errno:
-            telemetry_data[TelemetryField.KEY_ERROR_NUMBER.value] = str(self.errno)
+            telemetry_data_dict[TelemetryField.KEY_ERROR_NUMBER.value] = str(self.errno)
 
-        return telemetry_data
+        return telemetry_data_dict
 
     def send_exception_telemetry(
         self,
@@ -159,7 +156,11 @@ def send_exception_telemetry(
             telemetry_data[TelemetryField.KEY_EXCEPTION.value] = self.__class__.__name__
             ts = get_time_millis()
             try:
-                connection._log_telemetry(TelemetryData(telemetry_data, ts))
+                connection._log_telemetry(
+                    TelemetryData.from_telemetry_data_dict(
+                        from_dict=telemetry_data, timestamp=ts, connection=connection
+                    )
+                )
             except AttributeError:
                 logger.debug("Cursor failed to log to telemetry.", exc_info=True)
         elif connection is None:
@@ -176,13 +177,13 @@ def exception_telemetry(
     ) -> None:
         """Main method to generate and send telemetry data for exceptions."""
         try:
-            telemetry_data = self.generate_telemetry_exception_data()
+            telemetry_data_dict = self.generate_telemetry_exception_data()
             if cursor is not None:
-                self.send_exception_telemetry(cursor.connection, telemetry_data)
+                self.send_exception_telemetry(cursor.connection, telemetry_data_dict)
             elif connection is not None:
-                self.send_exception_telemetry(connection, telemetry_data)
+                self.send_exception_telemetry(connection, telemetry_data_dict)
             else:
-                self.send_exception_telemetry(None, telemetry_data)
+                self.send_exception_telemetry(None, telemetry_data_dict)
         except Exception:
             # Do nothing but log if sending telemetry fails
             logger.debug("Sending exception telemetry failed")

src/snowflake/connector/ocsp_snowflake.py

@@ -70,8 +70,7 @@
 
 from . import constants
 from .cache import SFDictCache, SFDictFileCache
-from .telemetry import TelemetryField
-from .telemetry import generate_telemetry_data as generate_telemetry_data_base
+from .telemetry import TelemetryField, generate_telemetry_data_dict
 
 try:
     OCSP_CACHE: SFDictFileCache[
@@ -225,7 +224,7 @@ def generate_telemetry_data(
         self, event_type: str, urgent: bool = False
     ) -> dict[str, Any]:
         _, exception, _ = sys.exc_info()
-        telemetry_data = generate_telemetry_data_base(
+        telemetry_data = generate_telemetry_data_dict(
             from_dict={
                 TelemetryField.KEY_OOB_EVENT_TYPE.value: event_type,
                 TelemetryField.KEY_OOB_EVENT_SUB_TYPE.value: self.event_sub_type,

src/snowflake/connector/telemetry.py

@@ -15,6 +15,7 @@
 from .test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
 
 if TYPE_CHECKING:
+    from .connection import SnowflakeConnection
     from .network import SnowflakeRestful
 
 logger = logging.getLogger(__name__)
@@ -90,6 +91,29 @@ def __init__(self, message, timestamp):
         self.message = message
         self.timestamp = timestamp
 
+    @classmethod
+    def from_telemetry_data_dict(
+        cls,
+        from_dict: dict,
+        timestamp: int,
+        connection: SnowflakeConnection | None = None,
+        is_oob_telemetry: bool = False,
+    ):
+        """
+        Generate telemetry data with driver info from given dict and timestamp.
+        It takes an optional connection object to read data from.
+        It also takes a boolean is_oob_telemetry to indicate whether it's for out-of-band telemetry, as
+        naming of keys for driver and version is different from the ones of in-band telemetry.
+        """
+        return cls(
+            generate_telemetry_data_dict(
+                from_dict=(from_dict or {}),
+                connection=connection,
+                is_oob_telemetry=is_oob_telemetry,
+            ),
+            timestamp,
+        )
+
     def to_dict(self):
         return {"message": self.message, "timestamp": str(self.timestamp)}
 
@@ -195,17 +219,25 @@ def buffer_size(self):
         return len(self._log_batch)
 
 
-def generate_telemetry_data(
-    from_dict: dict | None = None, is_oob_telemetry: bool = False
+def generate_telemetry_data_dict(
+    from_dict: dict | None = None,
+    connection: SnowflakeConnection | None = None,
+    is_oob_telemetry: bool = False,
 ) -> dict[str, Any]:
     """
-    Generate telemetry data with driver info. The method also takes an optional dict to update from.
+    Generate telemetry data with driver info.
+    The method also takes an optional dict to update from and optional connection object to read data from.
+    It also takes a boolean is_oob_telemetry to indicate whether it's for out-of-band telemetry, as
+    naming of keys for driver and version is different from the ones of in-band telemetry.
     """
     from_dict = from_dict or {}
     return (
         {
             TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
             TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            TelemetryField.KEY_SOURCE.value: connection.application
+            if connection
+            else CLIENT_NAME,
             **from_dict,
         }
         if not is_oob_telemetry

src/snowflake/connector/telemetry_oob.py

@@ -15,7 +15,7 @@
 from .compat import OK
 from .description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
 from .secret_detector import SecretDetector
-from .telemetry import TelemetryField, generate_telemetry_data
+from .telemetry import TelemetryField, generate_telemetry_data_dict
 from .test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
 from .vendored import requests
 
@@ -388,7 +388,7 @@ def log_http_request_error(
             if self.enabled:
                 response_status_code = -1
                 # This mimics the output of HttpRequestBase.toString() from JBDC
-                telemetry_data = generate_telemetry_data(
+                telemetry_data = generate_telemetry_data_dict(
                     from_dict={
                         TelemetryField.KEY_OOB_REQUEST.value: f"{method} {url}",
                         TelemetryField.KEY_OOB_SQL_STATE.value: sqlstate,

test/unit/test_telemetry.py

@@ -135,26 +135,29 @@ def test_telemetry_send_batch_disabled():
     assert rest_call.call_count == 0
 
 
-def test_generate_telemetry_with_driver_info():
-    assert snowflake.connector.telemetry.generate_telemetry_data() == {
+def test_generate_telemetry_data_dict_with_basic_info():
+    assert snowflake.connector.telemetry.generate_telemetry_data_dict() == {
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+        snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: CLIENT_NAME,
     }
 
-    assert snowflake.connector.telemetry.generate_telemetry_data(from_dict={}) == {
+    assert snowflake.connector.telemetry.generate_telemetry_data_dict(from_dict={}) == {
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+        snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: CLIENT_NAME,
     }
 
-    assert snowflake.connector.telemetry.generate_telemetry_data(
+    assert snowflake.connector.telemetry.generate_telemetry_data_dict(
         from_dict={"key": "value"}
     ) == {
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+        snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: CLIENT_NAME,
         "key": "value",
     }
 
-    assert snowflake.connector.telemetry.generate_telemetry_data(
+    assert snowflake.connector.telemetry.generate_telemetry_data_dict(
         from_dict={
             snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: "CUSTOM_CLIENT_NAME",
             snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: "1.2.3",
@@ -163,5 +166,76 @@ def test_generate_telemetry_with_driver_info():
     ) == {
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: "CUSTOM_CLIENT_NAME",
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: "1.2.3",
+        snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: CLIENT_NAME,
         "key": "value",
     }
+
+    mock_connection = Mock()
+    mock_connection.application = "test_application"
+    assert snowflake.connector.telemetry.generate_telemetry_data_dict(
+        from_dict={
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: "CUSTOM_CLIENT_NAME",
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: "1.2.3",
+            "key": "value",
+        },
+        connection=mock_connection,
+    ) == {
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: "CUSTOM_CLIENT_NAME",
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: "1.2.3",
+        snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: mock_connection.application,
+        "key": "value",
+    }
+
+
+def test_generate_telemetry_data():
+    telemetry_data = (
+        snowflake.connector.telemetry.TelemetryData.from_telemetry_data_dict(
+            from_dict={}, timestamp=123
+        )
+    )
+    assert (
+        telemetry_data.message
+        == {
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: CLIENT_NAME,
+        }
+        and telemetry_data.timestamp == 123
+    )
+
+    mock_connection = Mock()
+    mock_connection.application = "test_application"
+    telemetry_data = (
+        snowflake.connector.telemetry.TelemetryData.from_telemetry_data_dict(
+            from_dict={},
+            timestamp=123,
+            connection=mock_connection,
+        )
+    )
+    assert (
+        telemetry_data.message
+        == {
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: mock_connection.application,
+        }
+        and telemetry_data.timestamp == 123
+    )
+
+    telemetry_data = (
+        snowflake.connector.telemetry.TelemetryData.from_telemetry_data_dict(
+            from_dict={"key": "value"},
+            timestamp=123,
+            connection=mock_connection,
+        )
+    )
+    assert (
+        telemetry_data.message
+        == {
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+            snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            snowflake.connector.telemetry.TelemetryField.KEY_SOURCE.value: mock_connection.application,
+            "key": "value",
+        }
+        and telemetry_data.timestamp == 123
+    )