SNOW-1763096: Add async telemetry support (#2585)

Author: sfc-gh-turbaszek

src/snowflake/connector/errors.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+import inspect
 import logging
 import os
 import re
@@ -14,6 +15,8 @@
 from .time_util import get_time_millis
 
 if TYPE_CHECKING:  # pragma: no cover
+    from .aio._connection import SnowflakeConnection as AsyncSnowflakeConnection
+    from .aio._cursor import SnowflakeCursor as AsyncSnowflakeCursor
     from .connection import SnowflakeConnection
     from .cursor import SnowflakeCursor
 
@@ -35,8 +38,8 @@ def __init__(
         sfqid: str | None = None,
         query: str | None = None,
         done_format_msg: bool | None = None,
-        connection: SnowflakeConnection | None = None,
-        cursor: SnowflakeCursor | None = None,
+        connection: SnowflakeConnection | AsyncSnowflakeConnection | None = None,
+        cursor: SnowflakeCursor | AsyncSnowflakeCursor | None = None,
         errtype: TelemetryField = TelemetryField.SQL_EXCEPTION,
         send_telemetry: bool = True,
     ) -> None:
@@ -145,11 +148,10 @@ def generate_telemetry_exception_data(
 
     def send_exception_telemetry(
         self,
-        connection: SnowflakeConnection | None,
+        connection: SnowflakeConnection | AsyncSnowflakeConnection | None,
         telemetry_data: dict[str, Any],
     ) -> None:
         """Send telemetry data by in-band telemetry if it is enabled, otherwise send through out-of-band telemetry."""
-
         if (
             connection is not None
             and connection.telemetry_enabled
@@ -159,21 +161,34 @@ 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:
-                connection._log_telemetry(
+                result = connection._log_telemetry(
                     TelemetryData.from_telemetry_data_dict(
                         from_dict=telemetry_data, timestamp=ts, connection=connection
                     )
                 )
+                if inspect.isawaitable(result):
+                    try:
+                        import asyncio
+
+                        asyncio.get_running_loop().create_task(result)
+                    except Exception:
+                        logger.debug(
+                            "Failed to schedule async telemetry logging.",
+                            exc_info=True,
+                        )
             except AttributeError:
                 logger.debug("Cursor failed to log to telemetry.", exc_info=True)
 
     def exception_telemetry(
         self,
         msg: str,
-        cursor: SnowflakeCursor | None,
-        connection: SnowflakeConnection | None,
+        cursor: SnowflakeCursor | AsyncSnowflakeCursor | None,
+        connection: SnowflakeConnection | AsyncSnowflakeConnection | None,
     ) -> None:
         """Main method to generate and send telemetry data for exceptions."""
         try:
@@ -370,6 +385,18 @@ def errorhandler_make_exception(
             )
         return error_class(error_value)
 
+    @staticmethod
+    def _is_aio_connection(
+        connection: SnowflakeConnection | AsyncSnowflakeConnection,
+    ) -> bool:
+        try:
+            # Try import async connection. The import may fail if aio is not installed.
+            from .aio._connection import SnowflakeConnection as AsyncSnowflakeConnection
+
+            return isinstance(connection, AsyncSnowflakeConnection)
+        except ImportError:
+            return False
+
 
 class _Warning(Exception):
     """Exception for important warnings."""

src/snowflake/connector/telemetry.py

@@ -51,6 +51,7 @@ 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/csp_helpers.py

@@ -446,3 +446,7 @@ def __enter__(self):
     def __exit__(self, *args, **kwargs):
         self.os_environment_patch.__exit__(*args)
         super().__exit__(*args, **kwargs)
+
+
+def is_running_against_gcp():
+    return os.getenv("cloud_provider").lower() == "gcp"

test/integ/aio_it/test_cursor_binding_async.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from test.csp_helpers import is_running_against_gcp
+
 import pytest
 
 from snowflake.connector.errors import ProgrammingError
@@ -46,21 +48,23 @@ async def test_binding_security(conn_cnx, db_parameters):
 
             # SQL injection safe test
             # Good Example
-            with pytest.raises(ProgrammingError):
-                await cnx.cursor().execute(
-                    "SELECT * FROM {name} WHERE aa=%s".format(
-                        name=db_parameters["name"]
-                    ),
-                    ("1 or aa>0",),
-                )
-
-            with pytest.raises(ProgrammingError):
-                await cnx.cursor().execute(
-                    "SELECT * FROM {name} WHERE aa=%(aa)s".format(
-                        name=db_parameters["name"]
-                    ),
-                    {"aa": "1 or aa>0"},
-                )
+            if not is_running_against_gcp():
+                with pytest.raises(ProgrammingError):
+                    r = await cnx.cursor().execute(
+                        "SELECT * FROM {name} WHERE aa=%s".format(
+                            name=db_parameters["name"]
+                        ),
+                        ("1 or aa>0",),
+                    )
+                    await r.fetchall()
+
+                with pytest.raises(ProgrammingError):
+                    await cnx.cursor().execute(
+                        "SELECT * FROM {name} WHERE aa=%(aa)s".format(
+                            name=db_parameters["name"]
+                        ),
+                        {"aa": "1 or aa>0"},
+                    )
 
             # Bad Example in application. DON'T DO THIS
             c = cnx.cursor()

test/integ/test_cursor_binding.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+from test.csp_helpers import is_running_against_gcp
+
 import pytest
 
 from snowflake.connector.errors import ProgrammingError
@@ -42,21 +44,22 @@ def test_binding_security(conn_cnx, db_parameters):
 
             # SQL injection safe test
             # Good Example
-            with pytest.raises(ProgrammingError):
-                cnx.cursor().execute(
-                    "SELECT * FROM {name} WHERE aa=%s".format(
-                        name=db_parameters["name"]
-                    ),
-                    ("1 or aa>0",),
-                )
-
-            with pytest.raises(ProgrammingError):
-                cnx.cursor().execute(
-                    "SELECT * FROM {name} WHERE aa=%(aa)s".format(
-                        name=db_parameters["name"]
-                    ),
-                    {"aa": "1 or aa>0"},
-                )
+            if not is_running_against_gcp():
+                with pytest.raises(ProgrammingError):
+                    cnx.cursor().execute(
+                        "SELECT * FROM {name} WHERE aa=%s".format(
+                            name=db_parameters["name"]
+                        ),
+                        ("1 or aa>0",),
+                    )
+
+                with pytest.raises(ProgrammingError):
+                    cnx.cursor().execute(
+                        "SELECT * FROM {name} WHERE aa=%(aa)s".format(
+                            name=db_parameters["name"]
+                        ),
+                        {"aa": "1 or aa>0"},
+                    )
 
             # Bad Example in application. DON'T DO THIS
             c = cnx.cursor()

test/unit/aio/test_errors_telemetry.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, Mock, patch
+
+from snowflake.connector.aio import SnowflakeConnection
+from snowflake.connector.errors import Error
+from snowflake.connector.telemetry import TelemetryData, TelemetryField
+
+
+def _extract_message_from_log_call(mock_conn: Mock) -> dict:
+    mock_conn._log_telemetry.assert_called_once()
+    td = mock_conn._log_telemetry.call_args[0][0]
+    assert isinstance(td, TelemetryData)
+    return td.message
+
+
+async def test_error_telemetry_async_connection():
+    conn = Mock(SnowflakeConnection)
+    conn.telemetry_enabled = True
+    conn._telemetry = Mock()
+    conn._telemetry.is_closed = False
+    conn.application = "pytest_app_async"
+    conn._log_telemetry = AsyncMock()
+
+    with patch("asyncio.get_running_loop") as loop_mock:
+        Error(msg="kaboom", errno=654321, sqlstate="00000", connection=conn)
+        loop_mock.return_value.create_task.assert_called_once()
+
+    msg = _extract_message_from_log_call(conn)
+    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

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+from snowflake.connector.errors import Error
+from snowflake.connector.telemetry import TelemetryData, TelemetryField
+
+
+def _extract_message_from_log_call(mock_conn: Mock) -> dict:
+    mock_conn._log_telemetry.assert_called_once()
+    td = mock_conn._log_telemetry.call_args[0][0]
+    assert isinstance(td, TelemetryData)
+    return td.message
+
+
+def test_error_telemetry_sync_connection():
+    conn = Mock()
+    conn.telemetry_enabled = True
+    conn._telemetry = Mock()
+    conn._telemetry.is_closed = False
+    conn.application = "pytest_app"
+    conn._log_telemetry = Mock()
+
+    err = Error(msg="boom", errno=123456, sqlstate="00000", connection=conn)
+    assert str(err)
+
+    msg = _extract_message_from_log_call(conn)
+    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