NO-SNOW: Send aio usage outside error loop

Author: sfc-gh-turbaszek

src/snowflake/connector/aio/_connection.py

@@ -737,6 +737,18 @@ async def _log_telemetry_imported_packages(self) -> None:
                 )
             )
 
+    async def _log_async_usage(self):
+        await self._log_telemetry(
+            TelemetryData.from_telemetry_data_dict(
+                from_dict={
+                    TelemetryField.KEY_TYPE.value: TelemetryField.USES_AIO.value,
+                    TelemetryField.KEY_VALUE.value: str(True),
+                },
+                timestamp=get_time_millis(),
+                connection=self,
+            )
+        )
+
     async def _next_sequence_counter(self) -> int:
         """Gets next sequence counter. Used internally."""
         async with self._lock_sequence_counter:
@@ -1030,6 +1042,7 @@ async def connect(self, **kwargs) -> None:
             await self.__open_connection()
         self._telemetry = TelemetryClient(self._rest)
         await self._log_telemetry_imported_packages()
+        await self._log_async_usage()
 
     def cursor(
         self, cursor_class: type[SnowflakeCursor] = SnowflakeCursor

src/snowflake/connector/connection.py

@@ -591,6 +591,7 @@ def __init__(
 
         # get the imported modules from sys.modules
         self._log_telemetry_imported_packages()
+        self._log_async_usage()
         # check SNOW-1218851 for long term improvement plan to refactor ocsp code
         atexit.register(self._close_at_exit)
 
@@ -2295,6 +2296,18 @@ def _log_telemetry_imported_packages(self) -> None:
                 )
             )
 
+    def _log_async_usage(self):
+        self._log_telemetry(
+            TelemetryData.from_telemetry_data_dict(
+                from_dict={
+                    TelemetryField.KEY_TYPE.value: TelemetryField.USES_AIO.value,
+                    TelemetryField.KEY_VALUE.value: str(False),
+                },
+                timestamp=get_time_millis(),
+                connection=self,
+            )
+        )
+
     def is_valid(self) -> bool:
         """This function tries to answer the question: Is this connection still good for sending queries?
         Attempts to validate the connections both on the TCP/IP and Session levels."""

src/snowflake/connector/errors.py

@@ -161,9 +161,6 @@ def send_exception_telemetry(
             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__
-            telemetry_data[TelemetryField.KEY_USES_AIO.value] = str(
-                self._is_aio_connection(connection)
-            ).lower()
             ts = get_time_millis()
             try:
                 result = connection._log_telemetry(

src/snowflake/connector/telemetry.py

@@ -41,6 +41,8 @@ class TelemetryField(Enum):
     IMPORTED_PACKAGES = "client_imported_packages"
     # multi-statement usage
     MULTI_STATEMENT = "client_multi_statement_query"
+    # aio usage
+    USES_AIO = "client_aio_use"
     # Keys for telemetry data sent through either in-band or out-of-band telemetry
     KEY_TYPE = "type"
     KEY_SOURCE = "source"
@@ -51,7 +53,6 @@ class TelemetryField(Enum):
     KEY_REASON = "reason"
     KEY_VALUE = "value"
     KEY_EXCEPTION = "exception"
-    KEY_USES_AIO = "uses_aio"
     # Reserved UpperCamelName keys
     KEY_ERROR_NUMBER = "ErrorNumber"
     KEY_ERROR_MESSAGE = "ErrorMessage"

test/unit/aio/test_errors_telemetry.py

@@ -30,6 +30,5 @@ async def test_error_telemetry_async_connection():
     assert msg[TelemetryField.KEY_TYPE.value] == TelemetryField.SQL_EXCEPTION.value
     assert msg[TelemetryField.KEY_SOURCE.value] == conn.application
     assert msg[TelemetryField.KEY_EXCEPTION.value] == "Error"
-    assert msg[TelemetryField.KEY_USES_AIO.value] == "true"
     assert TelemetryField.KEY_DRIVER_TYPE.value in msg
     assert TelemetryField.KEY_DRIVER_VERSION.value in msg

test/unit/test_errors_telemetry.py

@@ -28,6 +28,5 @@ def test_error_telemetry_sync_connection():
     assert msg[TelemetryField.KEY_TYPE.value] == TelemetryField.SQL_EXCEPTION.value
     assert msg[TelemetryField.KEY_SOURCE.value] == conn.application
     assert msg[TelemetryField.KEY_EXCEPTION.value] == "Error"
-    assert msg[TelemetryField.KEY_USES_AIO.value] == "false"
     assert TelemetryField.KEY_DRIVER_TYPE.value in msg
     assert TelemetryField.KEY_DRIVER_VERSION.value in msg