SNOW-645247: Python connector telemetry is partially broken (#1231)

Author: sfc-gh-aling

src/snowflake/connector/cursor.py

@@ -62,7 +62,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
+from .telemetry import TelemetryData, TelemetryField, generate_telemetry_data
 from .time_util import get_time_millis
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -1211,15 +1211,23 @@ def _log_telemetry_job_data(
         self, telemetry_field: TelemetryField, value: Any
     ) -> None:
         """Builds an instance of TelemetryData with the given field and logs it."""
-        obj = {
-            "type": telemetry_field.value,
-            "source": self._connection.application if self._connection else CLIENT_NAME,
-            "query_id": self._sfqid,
-            "value": int(value),
-        }
         ts = get_time_millis()
         try:
-            self._connection._log_telemetry(TelemetryData(obj, ts))
+            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,
+                )
+            )
         except AttributeError:
             logger.warning(
                 "Cursor failed to log to telemetry. Connection object may be None.",

src/snowflake/connector/errors.py

@@ -13,9 +13,8 @@
 from typing import TYPE_CHECKING
 
 from .compat import BASE_EXCEPTION_CLASS
-from .description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
 from .secret_detector import SecretDetector
-from .telemetry import TelemetryData, TelemetryField
+from .telemetry import TelemetryData, TelemetryField, generate_telemetry_data
 from .telemetry_oob import TelemetryService
 from .time_util import get_time_millis
 
@@ -121,10 +120,13 @@ def telemetry_msg(self) -> str | None:
     def generate_telemetry_exception_data(self) -> dict[str, str]:
         """Generate the data to send through telemetry."""
 
-        telemetry_data = {
-            TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
-            TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
-        }
+        telemetry_data = generate_telemetry_data(
+            from_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
@@ -135,10 +137,6 @@ def generate_telemetry_exception_data(self) -> dict[str, str]:
         if self.errno:
             telemetry_data[TelemetryField.KEY_ERROR_NUMBER.value] = str(self.errno)
 
-        telemetry_data[
-            TelemetryField.KEY_STACKTRACE.value
-        ] = SecretDetector.mask_secrets(self.telemetry_traceback)
-
         return telemetry_data
 
     def send_exception_telemetry(

src/snowflake/connector/ocsp_snowflake.py

@@ -22,6 +22,7 @@
 from os.path import expanduser
 from threading import Lock, RLock
 from time import gmtime, strftime
+from typing import Any
 
 import jwt
 
@@ -69,6 +70,8 @@
 
 from . import constants
 from .cache import SFDictCache, SFDictFileCache
+from .telemetry import TelemetryField
+from .telemetry import generate_telemetry_data as generate_telemetry_data_base
 
 try:
     OCSP_CACHE: SFDictFileCache[
@@ -218,20 +221,26 @@ def set_fail_open(self, fail_open):
     def set_insecure_mode(self, insecure_mode):
         self.insecure_mode = insecure_mode
 
-    def generate_telemetry_data(self, event_type, urgent=False):
+    def generate_telemetry_data(
+        self, event_type: str, urgent: bool = False
+    ) -> dict[str, Any]:
         _, exception, _ = sys.exc_info()
-        telemetry_data = {}
-        telemetry_data.update({"eventType": event_type})
-        telemetry_data.update({"eventSubType": self.event_sub_type})
-        telemetry_data.update({"sfcPeerHost": self.sfc_peer_host})
-        telemetry_data.update({"certId": self.cert_id})
-        telemetry_data.update({"ocspRequestBase64": self.ocsp_req})
-        telemetry_data.update({"ocspResponderURL": self.ocsp_url})
-        telemetry_data.update({"errorMessage": self.error_msg})
-        telemetry_data.update({"insecureMode": self.insecure_mode})
-        telemetry_data.update({"failOpen": self.fail_open})
-        telemetry_data.update({"cacheEnabled": self.cache_enabled})
-        telemetry_data.update({"cacheHit": self.cache_hit})
+        telemetry_data = generate_telemetry_data_base(
+            from_dict={
+                TelemetryField.KEY_OOB_EVENT_TYPE.value: event_type,
+                TelemetryField.KEY_OOB_EVENT_SUB_TYPE.value: self.event_sub_type,
+                TelemetryField.KEY_OOB_SFC_PEER_HOST.value: self.sfc_peer_host,
+                TelemetryField.KEY_OOB_CERT_ID.value: self.cert_id,
+                TelemetryField.KEY_OOB_OCSP_REQUEST_BASE64.value: self.ocsp_req,
+                TelemetryField.KEY_OOB_OCSP_RESPONDER_URL.value: self.ocsp_url,
+                TelemetryField.KEY_OOB_ERROR_MESSAGE.value: self.error_msg,
+                TelemetryField.KEY_OOB_INSECURE_MODE.value: self.insecure_mode,
+                TelemetryField.KEY_OOB_FAIL_OPEN.value: self.fail_open,
+                TelemetryField.KEY_OOB_CACHE_ENABLED.value: self.cache_enabled,
+                TelemetryField.KEY_OOB_CACHE_HIT.value: self.cache_hit,
+            },
+            is_oob_telemetry=True,
+        )
 
         telemetry_client = TelemetryService.get_instance()
         telemetry_client.log_ocsp_exception(

src/snowflake/connector/telemetry.py

@@ -8,8 +8,9 @@
 import logging
 from enum import Enum, unique
 from threading import Lock
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from .description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
 from .secret_detector import SecretDetector
 from .test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
 
@@ -38,14 +39,43 @@ class TelemetryField(Enum):
     # Keys for telemetry data sent through either in-band or out-of-band telemetry
     KEY_TYPE = "type"
     KEY_SOURCE = "source"
-    KEY_SFQID = "QueryID"
-    KEY_SQLSTATE = "SQLState"
-    KEY_DRIVER_TYPE = "DriverType"
-    KEY_DRIVER_VERSION = "DriverVersion"
+    KEY_SFQID = "query_id"
+    KEY_SQLSTATE = "sql_state"
+    KEY_DRIVER_TYPE = "driver_type"
+    KEY_DRIVER_VERSION = "driver_version"
     KEY_REASON = "reason"
+    KEY_VALUE = "value"
+    KEY_EXCEPTION = "exception"
+    # Reserved UpperCamelName keys
     KEY_ERROR_NUMBER = "ErrorNumber"
+    KEY_ERROR_MESSAGE = "ErrorMessage"
     KEY_STACKTRACE = "Stacktrace"
-    KEY_EXCEPTION = "Exception"
+    # OOB camelName keys
+    KEY_OOB_DRIVER = "driver"
+    KEY_OOB_VERSION = "version"
+    KEY_OOB_TELEMETRY_SERVER_DEPLOYMENT = "telemetryServerDeployment"
+    KEY_OOB_CONNECTION_STRING = "connectionString"
+    KEY_OOB_EXCEPTION_MESSAGE = "exceptionMessage"
+    KEY_OOB_ERROR_MESSAGE = "errorMessage"
+    KEY_OOB_EXCEPTION_STACK_TRACE = "exceptionStackTrace"
+    KEY_OOB_EVENT_TYPE = "eventType"
+    KEY_OOB_ERROR_CODE = "errorCode"
+    KEY_OOB_SQL_STATE = "sqlState"
+    KEY_OOB_REQUEST = "request"
+    KEY_OOB_RESPONSE = "response"
+    KEY_OOB_RESPONSE_STATUS_LINE = "responseStatusLine"
+    KEY_OOB_RESPONSE_STATUS_CODE = "responseStatusCode"
+    KEY_OOB_RETRY_TIMEOUT = "retryTimeout"
+    KEY_OOB_RETRY_COUNT = "retryCount"
+    KEY_OOB_EVENT_SUB_TYPE = "eventSubType"
+    KEY_OOB_SFC_PEER_HOST = "sfcPeerHost"
+    KEY_OOB_CERT_ID = "certId"
+    KEY_OOB_OCSP_REQUEST_BASE64 = "ocspRequestBase64"
+    KEY_OOB_OCSP_RESPONDER_URL = "ocspResponderURL"
+    KEY_OOB_INSECURE_MODE = "insecureMode"
+    KEY_OOB_FAIL_OPEN = "failOpen"
+    KEY_OOB_CACHE_ENABLED = "cacheEnabled"
+    KEY_OOB_CACHE_HIT = "cacheHit"
 
 
 class TelemetryData:
@@ -161,3 +191,25 @@ def is_enabled(self):
 
     def buffer_size(self):
         return len(self._log_batch)
+
+
+def generate_telemetry_data(
+    from_dict: dict | 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.
+    """
+    from_dict = from_dict or {}
+    return (
+        {
+            TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+            TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            **from_dict,
+        }
+        if not is_oob_telemetry
+        else {
+            TelemetryField.KEY_OOB_DRIVER.value: CLIENT_NAME,
+            TelemetryField.KEY_OOB_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+            **from_dict,
+        }
+    )

src/snowflake/connector/telemetry_oob.py

@@ -15,6 +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 .test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
 from .vendored import requests
 
@@ -122,10 +123,14 @@ def generate_tags(self):
 
         telemetry = TelemetryService.get_instance()
         # Add telemetry service generated tags
-        tags["driver"] = CLIENT_NAME
-        tags["version"] = str(SNOWFLAKE_CONNECTOR_VERSION)
-        tags["telemetryServerDeployment"] = telemetry.deployment.name
-        tags["connectionString"] = telemetry.get_connection_string()
+        tags[TelemetryField.KEY_OOB_DRIVER] = CLIENT_NAME
+        tags[TelemetryField.KEY_OOB_VERSION] = str(SNOWFLAKE_CONNECTOR_VERSION)
+        tags[
+            TelemetryField.KEY_OOB_TELEMETRY_SERVER_DEPLOYMENT
+        ] = telemetry.deployment.name
+        tags[
+            TelemetryField.KEY_OOB_CONNECTION_STRING
+        ] = telemetry.get_connection_string()
         if telemetry.context and len(telemetry.context) > 0:
             for k, v in telemetry.context.items():
                 if v is not None:
@@ -339,14 +344,18 @@ def log_ocsp_exception(
             if self.enabled:
                 event_name = "OCSPException"
                 if exception is not None:
-                    telemetry_data["exceptionMessage"] = str(exception)
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_MESSAGE.value
+                    ] = str(exception)
                 if stack_trace is not None:
-                    telemetry_data["exceptionStackTrace"] = stack_trace
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_STACK_TRACE.value
+                    ] = stack_trace
 
                 if tags is None:
                     tags = dict()
 
-                tags["eventType"] = event_type
+                tags[TelemetryField.KEY_OOB_EVENT_TYPE.value] = event_type
 
                 log_event = TelemetryLogEvent(
                     name=event_name, tags=tags, urgent=urgent, value=telemetry_data
@@ -377,33 +386,53 @@ def log_http_request_error(
             tags = dict()
         try:
             if self.enabled:
-                telemetry_data = dict()
                 response_status_code = -1
                 # This mimics the output of HttpRequestBase.toString() from JBDC
-                telemetry_data["request"] = f"{method} {url}"
-                telemetry_data["sqlState"] = sqlstate
-                telemetry_data["errorCode"] = errno
+                telemetry_data = generate_telemetry_data(
+                    from_dict={
+                        TelemetryField.KEY_OOB_REQUEST.value: f"{method} {url}",
+                        TelemetryField.KEY_OOB_SQL_STATE.value: sqlstate,
+                        TelemetryField.KEY_OOB_ERROR_CODE.value: errno,
+                    },
+                    is_oob_telemetry=True,
+                )
                 if response:
-                    telemetry_data["response"] = response.json()
-                    telemetry_data["responseStatusLine"] = str(response.reason)
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_RESPONSE.value
+                    ] = response.json()
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_RESPONSE_STATUS_LINE.value
+                    ] = str(response.reason)
                     if response.status_code:
                         response_status_code = str(response.status_code)
-                        telemetry_data["responseStatusCode"] = response_status_code
+                        telemetry_data[
+                            TelemetryField.KEY_OOB_RESPONSE_STATUS_CODE.value
+                        ] = response_status_code
                 if retry_timeout:
-                    telemetry_data["retryTimeout"] = str(retry_timeout)
+                    telemetry_data[TelemetryField.KEY_OOB_RETRY_TIMEOUT.value] = str(
+                        retry_timeout
+                    )
                 if retry_count:
-                    telemetry_data["retryCount"] = str(retry_count)
+                    telemetry_data[TelemetryField.KEY_OOB_RETRY_COUNT.value] = str(
+                        retry_count
+                    )
                 if exception:
-                    telemetry_data["exceptionMessage"] = str(exception)
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_MESSAGE.value
+                    ] = str(exception)
                 if stack_trace:
-                    telemetry_data["exceptionStackTrace"] = stack_trace
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_STACK_TRACE.value
+                    ] = stack_trace
 
                 if tags is None:
                     tags = dict()
 
-                tags["responseStatusCode"] = response_status_code
-                tags["sqlState"] = str(sqlstate)
-                tags["errorCode"] = errno
+                tags[
+                    TelemetryField.KEY_OOB_RESPONSE_STATUS_CODE.value
+                ] = response_status_code
+                tags[TelemetryField.KEY_OOB_SQL_STATE.value] = str(sqlstate)
+                tags[TelemetryField.KEY_OOB_ERROR_CODE.value] = errno
 
                 log_event = TelemetryLogEvent(
                     name=event_name, tags=tags, value=telemetry_data, urgent=urgent

test/unit/test_telemetry.py

@@ -8,6 +8,7 @@
 from unittest.mock import Mock
 
 import snowflake.connector.telemetry
+from snowflake.connector.description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
 
 
 def test_telemetry_data_to_dict():
@@ -132,3 +133,35 @@ def test_telemetry_send_batch_disabled():
     client.send_batch()
     assert client.buffer_size() == 1
     assert rest_call.call_count == 0
+
+
+def test_generate_telemetry_with_driver_info():
+    assert snowflake.connector.telemetry.generate_telemetry_data() == {
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+    }
+
+    assert snowflake.connector.telemetry.generate_telemetry_data(from_dict={}) == {
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: SNOWFLAKE_CONNECTOR_VERSION,
+    }
+
+    assert snowflake.connector.telemetry.generate_telemetry_data(
+        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,
+        "key": "value",
+    }
+
+    assert snowflake.connector.telemetry.generate_telemetry_data(
+        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",
+        }
+    ) == {
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: "CUSTOM_CLIENT_NAME",
+        snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_VERSION.value: "1.2.3",
+        "key": "value",
+    }