SNOW-1570328: fix disabling client telemetry (#2013)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -10,6 +10,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.12.1(TBD)
   - Fixed a bug that session token is logged when renewing session.
+  - Fixed a bug that disabling client telemetry does not work.
 
 - v3.12.0(July 24,2024)
   - Set default connection timeout of 10 seconds and socket read timeout of 10 minutes for HTTP calls in file transfer.

src/snowflake/connector/connection.py

@@ -407,7 +407,8 @@ def __init__(
         self.messages = []
         self._async_sfqids: dict[str, None] = {}
         self._done_async_sfqids: dict[str, None] = {}
-        self.telemetry_enabled = False
+        self._client_param_telemetry_enabled = True
+        self._server_param_telemetry_enabled = False
         self._session_parameters: dict[str, str | int | bool] = {}
         logger.info(
             "Snowflake Connector for Python Version: %s, "
@@ -624,11 +625,22 @@ def consent_cache_id_token(self):
 
     @property
     def telemetry_enabled(self) -> bool:
-        return self._telemetry_enabled
+        return bool(
+            self._client_param_telemetry_enabled
+            and self._server_param_telemetry_enabled
+        )
 
     @telemetry_enabled.setter
     def telemetry_enabled(self, value) -> None:
-        self._telemetry_enabled = True if value else False
+        self._client_param_telemetry_enabled = True if value else False
+        if (
+            self._client_param_telemetry_enabled
+            and not self._server_param_telemetry_enabled
+        ):
+            logger.info(
+                "Telemetry has been disabled by the session parameter CLIENT_TELEMETRY_ENABLED."
+                " Set session parameter CLIENT_TELEMETRY_ENABLED to true to enable telemetry."
+            )
 
     @property
     def service_name(self) -> str | None:
@@ -774,7 +786,7 @@ def close(self, retry: bool = True) -> None:
 
             # close telemetry first, since it needs rest to send remaining data
             logger.info("closed")
-            self._telemetry.close(send_on_close=retry)
+            self._telemetry.close(send_on_close=bool(retry and self.telemetry_enabled))
             if (
                 self._all_async_queries_finished()
                 and not self._server_session_keep_alive
@@ -1714,7 +1726,7 @@ def _update_parameters(
         for name, value in parameters.items():
             self._session_parameters[name] = value
             if PARAMETER_CLIENT_TELEMETRY_ENABLED == name:
-                self.telemetry_enabled = value
+                self._server_param_telemetry_enabled = value
             elif PARAMETER_CLIENT_SESSION_KEEP_ALIVE == name:
                 # Only set if the local config is None.
                 # Always give preference to user config.

test/integ/test_connection.py

@@ -1404,3 +1404,42 @@ def test_mock_non_existing_server(conn_cnx, caplog):
                 "writing OCSP response cache file to",
             ]
         )
+
+
+@pytest.mark.skipolddriver
+def test_disable_telemetry(conn_cnx, caplog):
+    # default behavior, closing connection, it will send telemetry
+    with caplog.at_level(logging.DEBUG):
+        with conn_cnx() as conn:
+            with conn.cursor() as cur:
+                cur.execute("select 1").fetchall()
+            assert (
+                len(conn._telemetry._log_batch) == 3
+            )  # 3 events are import package, fetch first, fetch last
+    assert "POST /telemetry/send" in caplog.text
+    caplog.clear()
+
+    # set session parameters to false
+    with caplog.at_level(logging.DEBUG):
+        with conn_cnx(
+            session_parameters={"CLIENT_TELEMETRY_ENABLED": False}
+        ) as conn, conn.cursor() as cur:
+            cur.execute("select 1").fetchall()
+            assert not conn.telemetry_enabled and not conn._telemetry._log_batch
+            # this enable won't work as the session parameter is set to false
+            conn.telemetry_enabled = True
+            cur.execute("select 1").fetchall()
+            assert not conn.telemetry_enabled and not conn._telemetry._log_batch
+
+    assert "POST /telemetry/send" not in caplog.text
+    caplog.clear()
+
+    # test disable telemetry in the client
+    with caplog.at_level(logging.DEBUG):
+        with conn_cnx() as conn:
+            assert conn.telemetry_enabled and len(conn._telemetry._log_batch) == 1
+            conn.telemetry_enabled = False
+            with conn.cursor() as cur:
+                cur.execute("select 1").fetchall()
+            assert not conn.telemetry_enabled
+    assert "POST /telemetry/send" not in caplog.text