SNOW-996732: improve clean up logic of connection (#1898)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -22,8 +22,8 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
       - NOTE: this has not been tested extensively, but has been shown to improve the experience when using WSL
   - Bumped platformdirs from >=2.6.0,<4.0.0 to >=2.6.0,<5.0.0
   - Updated diagnostics to use system$allowlist instead of system$whitelist.
-  - Updated diagnostics to use system$allowlist instead of system$whitelist.
-  - Update `write_pandas` to skip TABLE IF NOT EXISTS in truncate mode
+  - Update `write_pandas` to skip TABLE IF NOT EXISTS in truncate mode.
+  - Improved cleanup logic for connection to rely on interpreter shutdown instead of the `__del__` method.
 
 - v3.7.1(February 21, 2024)
 

src/snowflake/connector/connection.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import atexit
 import logging
 import os
 import pathlib
@@ -16,6 +17,7 @@
 import weakref
 from concurrent.futures import as_completed
 from concurrent.futures.thread import ThreadPoolExecutor
+from contextlib import suppress
 from difflib import get_close_matches
 from functools import partial
 from io import StringIO
@@ -439,12 +441,8 @@ def __init__(
 
         # get the imported modules from sys.modules
         self._log_telemetry_imported_packages()
-
-    def __del__(self) -> None:  # pragma: no cover
-        try:
-            self.close(retry=False)
-        except Exception:
-            pass
+        # check SNOW-1218851 for long term improvement plan to refactor ocsp code
+        atexit.register(self._close_at_exit)
 
     @property
     def insecure_mode(self) -> bool:
@@ -744,6 +742,8 @@ def connect(self, **kwargs) -> None:
 
     def close(self, retry: bool = True) -> None:
         """Closes the connection."""
+        # unregister to dereference connection object as it's already closed after the execution
+        atexit.unregister(self._close_at_exit)
         try:
             if not self.rest:
                 logger.debug("Rest object has been destroyed, cannot close session")
@@ -1771,6 +1771,10 @@ def _cache_query_status(self, sf_qid: str, status_ret: QueryStatus) -> None:
             )  # Prevent KeyError when multiple threads try to remove the same query id
             self._done_async_sfqids[sf_qid] = None
 
+    def _close_at_exit(self):
+        with suppress(Exception):
+            self.close(retry=False)
+
     def get_query_status(self, sf_qid: str) -> QueryStatus:
         """Retrieves the status of query with sf_qid.
 

test/integ/test_connection.py

@@ -1342,3 +1342,11 @@ def test_ocsp_mode_insecure(conn_cnx, is_public_test, caplog):
             assert "snowflake.connector.ocsp_snowflake" in caplog.text
         else:
             assert "snowflake.connector.ocsp_snowflake" not in caplog.text
+
+
+@pytest.mark.skipolddriver
+def test_connection_atexit_close(conn_cnx):
+    """Basic Connection test without schema."""
+    with conn_cnx() as conn:
+        conn._close_at_exit()
+        assert conn.is_closed()