NO-SNOW: Send telemetry when closing connection (#2607)

Author: sfc-gh-turbaszek

src/snowflake/connector/connection.py

@@ -1158,7 +1158,8 @@ def close(self, retry: bool = True) -> None:
 
             # close telemetry first, since it needs rest to send remaining data
             logger.debug("closed")
-            self._telemetry.close(send_on_close=bool(retry and self.telemetry_enabled))
+            if self.telemetry_enabled:
+                self._telemetry.close(retry=retry)
             if (
                 self._all_async_queries_finished()
                 and not self._server_session_keep_alive

src/snowflake/connector/telemetry.py

@@ -155,7 +155,7 @@ def try_add_log_to_batch(self, telemetry_data: TelemetryData) -> None:
         except Exception:
             logger.warning("Failed to add log to telemetry.", exc_info=True)
 
-    def send_batch(self) -> None:
+    def send_batch(self, retry: bool = False) -> None:
         if self.is_closed:
             raise Exception("Attempted to send batch when TelemetryClient is closed")
         elif not self._enabled:
@@ -186,6 +186,7 @@ def send_batch(self) -> None:
                 method="post",
                 client=None,
                 timeout=5,
+                _no_retry=not retry,
             )
             if not ret["success"]:
                 logger.info(
@@ -204,11 +205,10 @@ def send_batch(self) -> None:
     def is_closed(self) -> bool:
         return self._rest is None
 
-    def close(self, send_on_close: bool = True) -> None:
+    def close(self, retry: bool = False) -> None:
         if not self.is_closed:
             logger.debug("Closing telemetry client.")
-            if send_on_close:
-                self.send_batch()
+            self.send_batch(retry=retry)
             self._rest = None
 
     def disable(self) -> None:

test/unit/test_telemetry.py

@@ -157,6 +157,30 @@ def test_telemetry_send_batch_disabled():
     assert rest_call.call_count == 0
 
 
+def test_telemetry_send_batch_with_retry_flag():
+    """Tests that send_batch respects the retry parameter."""
+    client, rest_call = get_client_and_mock()
+
+    client.add_log_to_batch(snowflake.connector.telemetry.TelemetryData({}, 2000))
+
+    # Test with retry=True
+    client.send_batch(retry=True)
+
+    assert rest_call.call_count == 1
+    # Verify _no_retry parameter is False when retry=True
+    call_kwargs = rest_call.call_args[1]
+    assert call_kwargs["_no_retry"] is False
+
+    # Add another log and test with retry=False (default)
+    client.add_log_to_batch(snowflake.connector.telemetry.TelemetryData({}, 3000))
+    client.send_batch(retry=False)
+
+    assert rest_call.call_count == 2
+    # Verify _no_retry parameter is True when retry=False
+    call_kwargs = rest_call.call_args[1]
+    assert call_kwargs["_no_retry"] is True
+
+
 def test_generate_telemetry_data_dict_with_basic_info():
     assert snowflake.connector.telemetry.generate_telemetry_data_dict() == {
         snowflake.connector.telemetry.TelemetryField.KEY_DRIVER_TYPE.value: CLIENT_NAME,