SNOW-1617451: async telemetry support and client name change (#2077)

Author: sfc-gh-yuwang

src/snowflake/connector/aio/_connection.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import atexit
+import copy
 import logging
 import os
 import pathlib
@@ -29,7 +30,7 @@
 from .._query_context_cache import QueryContextCache
 from ..compat import IS_LINUX, quote, urlencode
 from ..config_manager import CONFIG_MANAGER, _get_default_connection_params
-from ..connection import DEFAULT_CONFIGURATION
+from ..connection import DEFAULT_CONFIGURATION as DEFAULT_CONFIGURATION_SYNC
 from ..connection import SnowflakeConnection as SnowflakeConnectionSync
 from ..connection import _get_private_bytes_from_file
 from ..connection_diagnostic import ConnectionDiagnostic
@@ -69,7 +70,9 @@
 from ..time_util import get_time_millis
 from ..util_text import split_statements
 from ._cursor import SnowflakeCursor
+from ._description import CLIENT_NAME
 from ._network import SnowflakeRestful
+from ._telemetry import TelemetryClient
 from ._time_util import HeartBeatTimer
 from .auth import (
     FIRST_PARTY_AUTHENTICATORS,
@@ -86,6 +89,10 @@
 
 logger = getLogger(__name__)
 
+# deep copy to avoid pollute sync config
+DEFAULT_CONFIGURATION = copy.deepcopy(DEFAULT_CONFIGURATION_SYNC)
+DEFAULT_CONFIGURATION["application"] = (CLIENT_NAME, (type(None), str))
+
 
 class SnowflakeConnection(SnowflakeConnectionSync):
     OCSP_ENV_LOCK = asyncio.Lock()
@@ -103,11 +110,7 @@ def __init__(
             kwargs, connection_name, connections_file_path
         )
         self._connected = False
-        # TODO: async telemetry support
-        self._telemetry = None
         self.expired = False
-        # get the imported modules from sys.modules
-        # self._log_telemetry_imported_packages() # TODO: async telemetry support
         # check SNOW-1218851 for long term improvement plan to refactor ocsp code
         atexit.register(self._close_at_exit)
 
@@ -569,7 +572,8 @@ async def _get_query_status(
         return status_ret, status_resp
 
     async def _log_telemetry(self, telemetry_data) -> None:
-        raise NotImplementedError("asyncio telemetry is not supported")
+        if self.telemetry_enabled:
+            await self._telemetry.try_add_log_to_batch(telemetry_data)
 
     async def _log_telemetry_imported_packages(self) -> None:
         if self._log_imported_packages_in_telemetry:
@@ -722,8 +726,9 @@ async def close(self, retry: bool = True) -> None:
             # close telemetry first, since it needs rest to send remaining data
             logger.info("closed")
 
-            # TODO: async telemetry support
-            # self._telemetry.close(send_on_close=bool(retry and self.telemetry_enabled))
+            await self._telemetry.close(
+                send_on_close=bool(retry and self.telemetry_enabled)
+            )
             if (
                 await self._all_async_queries_finished()
                 and not self._server_session_keep_alive
@@ -889,6 +894,8 @@ async def connect(self, **kwargs) -> None:
                     raise Exception(str(exceptions_dict))
         else:
             await self.__open_connection()
+        self._telemetry = TelemetryClient(self._rest)
+        await self._log_telemetry_imported_packages()
 
     def cursor(
         self, cursor_class: type[SnowflakeCursor] = SnowflakeCursor

src/snowflake/connector/aio/_cursor.py

@@ -55,7 +55,7 @@
 )
 from snowflake.connector.errors import BindUploadError, DatabaseError
 from snowflake.connector.file_transfer_agent import SnowflakeProgressPercentage
-from snowflake.connector.telemetry import TelemetryField
+from snowflake.connector.telemetry import TelemetryData, TelemetryField
 from snowflake.connector.time_util import get_time_millis
 
 if TYPE_CHECKING:
@@ -361,8 +361,9 @@ async def _init_result_and_meta(self, data: dict[Any, Any]) -> None:
                 self._total_rowcount += updated_rows
 
     async def _init_multi_statement_results(self, data: dict) -> None:
-        # TODO: async telemetry SNOW-1572217
-        # self._log_telemetry_job_data(TelemetryField.MULTI_STATEMENT, TelemetryData.TRUE)
+        await self._log_telemetry_job_data(
+            TelemetryField.MULTI_STATEMENT, TelemetryData.TRUE
+        )
         self.multi_statement_savedIds = data["resultIds"].split(",")
         self._multi_statement_resultIds = collections.deque(
             self.multi_statement_savedIds
@@ -382,8 +383,24 @@ async def _init_multi_statement_results(self, data: dict) -> None:
     async def _log_telemetry_job_data(
         self, telemetry_field: TelemetryField, value: Any
     ) -> None:
-        # TODO: async telemetry SNOW-1572217
-        pass
+        ts = get_time_millis()
+        try:
+            await self._connection._log_telemetry(
+                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:
+            logger.warning(
+                "Cursor failed to log to telemetry. Connection object may be None.",
+                exc_info=True,
+            )
 
     async def _preprocess_pyformat_query(
         self,
@@ -394,16 +411,15 @@ async def _preprocess_pyformat_query(
         # client side binding
         processed_params = self._connection._process_params_pyformat(params, self)
         # SNOW-513061 collect telemetry for empty sequence usage before we make the breaking change announcement
-        # TODO: async telemetry support
-        # if params is not None and len(params) == 0:
-        #     await self._log_telemetry_job_data(
-        #         TelemetryField.EMPTY_SEQ_INTERPOLATION,
-        #         (
-        #             TelemetryData.TRUE
-        #             if self.connection._interpolate_empty_sequences
-        #             else TelemetryData.FALSE
-        #         ),
-        #     )
+        if params is not None and len(params) == 0:
+            await self._log_telemetry_job_data(
+                TelemetryField.EMPTY_SEQ_INTERPOLATION,
+                (
+                    TelemetryData.TRUE
+                    if self.connection._interpolate_empty_sequences
+                    else TelemetryData.FALSE
+                ),
+            )
         if logger.getEffectiveLevel() <= logging.DEBUG:
             logger.debug(
                 f"binding: [{self._format_query_for_log(command)}] "
@@ -585,14 +601,13 @@ async def execute(
         self._first_chunk_time = get_time_millis()
 
         # if server gives a send time, log the time it took to arrive
-        # TODO: telemetry support in asyncio
-        # if "data" in ret and "sendResultTime" in ret["data"]:
-        #     time_consume_first_result = (
-        #         self._first_chunk_time - ret["data"]["sendResultTime"]
-        #     )
-        #     self._log_telemetry_job_data(
-        #         TelemetryField.TIME_CONSUME_FIRST_RESULT, time_consume_first_result
-        #     )
+        if "data" in ret and "sendResultTime" in ret["data"]:
+            time_consume_first_result = (
+                self._first_chunk_time - ret["data"]["sendResultTime"]
+            )
+            await self._log_telemetry_job_data(
+                TelemetryField.TIME_CONSUME_FIRST_RESULT, time_consume_first_result
+            )
 
         if ret["success"]:
             logger.debug("SUCCESS")
@@ -893,10 +908,9 @@ async def fetch_arrow_batches(self) -> AsyncIterator[Table]:
             await self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
-        # TODO: async telemetry SNOW-1572217
-        # self._log_telemetry_job_data(
-        #     TelemetryField.ARROW_FETCH_BATCHES, TelemetryData.TRUE
-        # )
+        await self._log_telemetry_job_data(
+            TelemetryField.ARROW_FETCH_BATCHES, TelemetryData.TRUE
+        )
         return await self._result_set._fetch_arrow_batches()
 
     @overload
@@ -920,8 +934,9 @@ async def fetch_arrow_all(self, force_return_table: bool = False) -> Table | Non
             await self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
-        # TODO: async telemetry SNOW-1572217
-        # self._log_telemetry_job_data(TelemetryField.ARROW_FETCH_ALL, TelemetryData.TRUE)
+        await self._log_telemetry_job_data(
+            TelemetryField.ARROW_FETCH_ALL, TelemetryData.TRUE
+        )
         return await self._result_set._fetch_arrow_all(
             force_return_table=force_return_table
         )
@@ -933,10 +948,9 @@ async def fetch_pandas_batches(self, **kwargs: Any) -> AsyncIterator[DataFrame]:
             await self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
-        # TODO: async telemetry
-        # self._log_telemetry_job_data(
-        #     TelemetryField.PANDAS_FETCH_BATCHES, TelemetryData.TRUE
-        # )
+        await self._log_telemetry_job_data(
+            TelemetryField.PANDAS_FETCH_BATCHES, TelemetryData.TRUE
+        )
         return await self._result_set._fetch_pandas_batches(**kwargs)
 
     async def fetch_pandas_all(self, **kwargs: Any) -> DataFrame:
@@ -946,9 +960,9 @@ async def fetch_pandas_all(self, **kwargs: Any) -> DataFrame:
         if self._query_result_format != "arrow":
             raise NotSupportedError
         # # TODO: async telemetry
-        # self._log_telemetry_job_data(
-        #     TelemetryField.PANDAS_FETCH_ALL, TelemetryData.TRUE
-        # )
+        await self._log_telemetry_job_data(
+            TelemetryField.PANDAS_FETCH_ALL, TelemetryData.TRUE
+        )
         return await self._result_set._fetch_pandas_all(**kwargs)
 
     async def nextset(self) -> SnowflakeCursor | None:
@@ -981,9 +995,9 @@ async def get_result_batches(self) -> list[ResultBatch] | None:
         if self._result_set is None:
             return None
         # TODO: async telemetry SNOW-1572217
-        # self._log_telemetry_job_data(
-        #     TelemetryField.GET_PARTITIONS_USED, TelemetryData.TRUE
-        # )
+        await self._log_telemetry_job_data(
+            TelemetryField.GET_PARTITIONS_USED, TelemetryData.TRUE
+        )
         return self._result_set.batches
 
     async def get_results_from_sfqid(self, sfqid: str) -> None:

src/snowflake/connector/aio/_description.py

@@ -0,0 +1,9 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+"""Various constants."""
+
+from __future__ import annotations
+
+CLIENT_NAME = "AsyncioPythonConnector"  # don't change!

src/snowflake/connector/aio/_network.py

@@ -60,16 +60,19 @@
     HEADER_AUTHORIZATION_KEY,
     HEADER_SNOWFLAKE_TOKEN,
     ID_TOKEN_EXPIRED_GS_CODE,
+    IMPLEMENTATION,
     MASTER_TOKEN_EXPIRED_GS_CODE,
     MASTER_TOKEN_INVALD_GS_CODE,
     MASTER_TOKEN_NOTFOUND_GS_CODE,
     NO_TOKEN,
-    PYTHON_CONNECTOR_USER_AGENT,
+    PLATFORM,
+    PYTHON_VERSION,
     QUERY_IN_PROGRESS_ASYNC_CODE,
     QUERY_IN_PROGRESS_CODE,
     REQUEST_ID,
     REQUEST_TYPE_RENEW,
     SESSION_EXPIRED_GS_CODE,
+    SNOWFLAKE_CONNECTOR_VERSION,
     ReauthenticationRequest,
     RetryRequest,
 )
@@ -83,13 +86,16 @@
     SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
 )
 from ..time_util import TimeoutBackoffCtx
+from ._description import CLIENT_NAME
 from ._ssl_connector import SnowflakeSSLConnector
 
 if TYPE_CHECKING:
     from snowflake.connector.aio import SnowflakeConnection
 
 logger = logging.getLogger(__name__)
 
+PYTHON_CONNECTOR_USER_AGENT = f"{CLIENT_NAME}/{SNOWFLAKE_CONNECTOR_VERSION} ({PLATFORM}) {IMPLEMENTATION}/{PYTHON_VERSION}"
+
 try:
     import aiohttp
 except ImportError:

src/snowflake/connector/aio/_telemetry.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import logging
+from asyncio import Lock
+from typing import TYPE_CHECKING
+
+from ..secret_detector import SecretDetector
+from ..telemetry import TelemetryClient as TelemetryClientSync
+from ..telemetry import TelemetryData
+from ..test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
+
+if TYPE_CHECKING:
+    from ._network import SnowflakeRestful
+
+logger = logging.getLogger(__name__)
+
+
+class TelemetryClient(TelemetryClientSync):
+    """Client to enqueue and send metrics to the telemetry endpoint in batch."""
+
+    def __init__(self, rest: SnowflakeRestful, flush_size=None) -> None:
+        super().__init__(rest, flush_size)
+        self._lock = Lock()
+
+    async def add_log_to_batch(self, telemetry_data: TelemetryData) -> None:
+        if self.is_closed:
+            raise Exception("Attempted to add log when TelemetryClient is closed")
+        elif not self._enabled:
+            logger.debug("TelemetryClient disabled. Ignoring log.")
+            return
+
+        async with self._lock:
+            self._log_batch.append(telemetry_data)
+
+        if len(self._log_batch) >= self._flush_size:
+            await self.send_batch()
+
+    async def send_batch(self) -> None:
+        if self.is_closed:
+            raise Exception("Attempted to send batch when TelemetryClient is closed")
+        elif not self._enabled:
+            logger.debug("TelemetryClient disabled. Not sending logs.")
+            return
+
+        async with self._lock:
+            to_send = self._log_batch
+            self._log_batch = []
+
+        if not to_send:
+            logger.debug("Nothing to send to telemetry.")
+            return
+
+        body = {"logs": [x.to_dict() for x in to_send]}
+        logger.debug(
+            "Sending %d logs to telemetry. Data is %s.",
+            len(body),
+            SecretDetector.mask_secrets(str(body))[1],
+        )
+        if ENABLE_TELEMETRY_LOG:
+            # This logger guarantees the payload won't be masked. Testing purpose.
+            rt_plain_logger.debug(f"Inband telemetry data being sent is {body}")
+        try:
+            ret = await self._rest.request(
+                TelemetryClient.SF_PATH_TELEMETRY,
+                body=body,
+                method="post",
+                client=None,
+                timeout=5,
+            )
+            if not ret["success"]:
+                logger.info(
+                    "Non-success response from telemetry server: %s. "
+                    "Disabling telemetry.",
+                    str(ret),
+                )
+                self._enabled = False
+            else:
+                logger.debug("Successfully uploading metrics to telemetry.")
+        except Exception:
+            self._enabled = False
+            logger.debug("Failed to upload metrics to telemetry.", exc_info=True)
+
+    async def try_add_log_to_batch(self, telemetry_data: TelemetryData) -> None:
+        try:
+            await self.add_log_to_batch(telemetry_data)
+        except Exception:
+            logger.warning("Failed to add log to telemetry.", exc_info=True)
+
+    async def close(self, send_on_close: bool = True) -> None:
+        if not self.is_closed:
+            logger.debug("Closing telemetry client.")
+            if send_on_close:
+                await self.send_batch()
+            self._rest = None