fixup! NO-SNOW: Send aio usage outside error loop

Author: sfc-gh-turbaszek

src/snowflake/connector/aio/_connection.py

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

src/snowflake/connector/aio/_telemetry.py

@@ -38,7 +38,7 @@ async def add_log_to_batch(self, telemetry_data: TelemetryData) -> None:
         if len(self._log_batch) >= self._flush_size:
             await self.send_batch()
 
-    async def send_batch(self) -> None:
+    async def send_batch(self, retry: bool = True) -> None:
         if self.is_closed:
             raise Exception("Attempted to send batch when TelemetryClient is closed")
         elif not self._enabled:
@@ -69,6 +69,7 @@ async def send_batch(self) -> None:
                 method="post",
                 client=None,
                 timeout=5,
+                _no_retry=not retry,
             )
             if not ret["success"]:
                 logger.info(
@@ -89,9 +90,8 @@ async def try_add_log_to_batch(self, telemetry_data: TelemetryData) -> None:
         except Exception:
             logger.warning("Failed to add log to telemetry.", exc_info=True)
 
-    async def close(self, send_on_close: bool = True) -> None:
+    async def close(self, retry: bool = True) -> None:
         if not self.is_closed:
             logger.debug("Closing telemetry client.")
-            if send_on_close:
-                await self.send_batch()
+            await self.send_batch(retry=retry)
             self._rest = None

src/snowflake/connector/connection.py

@@ -984,7 +984,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

@@ -157,7 +157,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:
@@ -188,6 +188,7 @@ def send_batch(self) -> None:
                 method="post",
                 client=None,
                 timeout=5,
+                _no_retry=not retry,
             )
             if not ret["success"]:
                 logger.info(
@@ -206,11 +207,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: