fixup! NO-SNOW: Add thread-safe pool for tracking opened connections

Author: sfc-gh-turbaszek

src/snowflake/connector/connection.py

@@ -1154,15 +1154,15 @@ def connect(self, **kwargs) -> None:
             self.__open_connection()
 
         # Register the connection in the pool after successful connection
-        _connections_pool.add_connection(self)
+        _connections_registry.add_connection(self)
 
     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:
             # Remove connection from the pool
-            _connections_pool.remove_connection(self)
+            _connections_registry.remove_connection(self)
 
             if not self.rest:
                 logger.debug("Rest object has been destroyed, cannot close session")
@@ -2542,20 +2542,20 @@ def _detect_application() -> None | str:
             return "snowflake_notebook"
 
 
-class _ConnectionsPool:
-    """Thread-safe pool for tracking opened SnowflakeConnection instances.
+class _ConnectionsRegistry:
+    """Thread-safe registry for tracking opened SnowflakeConnection instances.
 
     This class maintains a registry of active connections using weak references
     to avoid preventing garbage collection.
     """
 
     def __init__(self):
-        """Initialize the connections pool with an empty registry and a lock."""
+        """Initialize the connections registry with an empty registry and a lock."""
         self._connections: weakref.WeakSet = weakref.WeakSet()
         self._lock = Lock()
 
     def add_connection(self, connection: SnowflakeConnection) -> None:
-        """Add a connection to the pool.
+        """Add a connection to the registry.
 
         Args:
             connection: The SnowflakeConnection instance to register.
@@ -2567,26 +2567,29 @@ def add_connection(self, connection: SnowflakeConnection) -> None:
             )
 
     def remove_connection(self, connection: SnowflakeConnection) -> None:
-        """Remove a connection from the pool.
+        """Remove a connection from the registry.
 
         Args:
             connection: The SnowflakeConnection instance to unregister.
         """
         with self._lock:
             self._connections.discard(connection)
             logger.debug(
-                f"Connection {id(connection)} removed from pool. Total connections: {len(self._connections)}"
+                f"Connection {id(connection)} removed from registry. Total connections: {len(self._connections)}"
             )
 
             if len(self._connections) == 0:
-                # If no connections left then stop CRL background task
-                # to avoid script dangling
-                CRLCacheFactory.stop_periodic_cleanup()
+                self._last_connection_handler()
+
+    def _last_connection_handler(self):
+        # If no connections left then stop CRL background task
+        # to avoid script dangling
+        CRLCacheFactory.stop_periodic_cleanup()
 
     def get_connection_count(self) -> int:
         with self._lock:
             return len(self._connections)
 
 
 # Global instance of the connections pool
-_connections_pool = _ConnectionsPool()
+_connections_registry = _ConnectionsRegistry()

test/unit/test_connection.py

@@ -955,41 +955,42 @@ def test_connect_metadata_preservation():
     # Should have parameters like account, user, password, etc.
 
 
-def test_connections_pool(mock_post_requests):
+def test_connections_registry(mock_post_requests):
     """Test that connections are properly tracked in the _ConnectionsPool."""
-    from snowflake.connector.connection import _connections_pool
+    from snowflake.connector.connection import _connections_registry
 
     # Get initial connection count
-    initial_count = _connections_pool.get_connection_count()
+    initial_count = _connections_registry.get_connection_count()
 
     # Create and connect first connection
     conn1 = fake_connector()
     assert (
-        _connections_pool.get_connection_count() == initial_count + 1
+        _connections_registry.get_connection_count() == initial_count + 1
     ), "Connection count should increase by 1 after creating a connection"
 
     # Create and connect second connection
     conn2 = fake_connector()
     assert (
-        _connections_pool.get_connection_count() == initial_count + 2
+        _connections_registry.get_connection_count() == initial_count + 2
     ), "Connection count should increase by 2 after creating two connections"
 
     # Close first connection
     conn1.close()
     assert (
-        _connections_pool.get_connection_count() == initial_count + 1
+        _connections_registry.get_connection_count() == initial_count + 1
     ), "Connection count should decrease by 1 after closing a connection"
 
     # Close second connection
     conn2.close()
     assert (
-        _connections_pool.get_connection_count() == initial_count
+        _connections_registry.get_connection_count() == initial_count
     ), "Connection count should return to initial count after closing all connections"
 
 
 @mock.patch("snowflake.connector.connection.CRLCacheFactory")
-def test_connections_pool_stops_crl_task_if_empty(crl_mock, mock_post_requests):
+def test_connections_registry_stops_crl_task_if_empty(crl_mock, mock_post_requests):
     """Test the individual methods of _ConnectionsPool."""
+    from snowflake.connector.connection import _connections_registry
 
     # Create a connection
     conn1 = fake_connector()
@@ -1001,4 +1002,5 @@ def test_connections_pool_stops_crl_task_if_empty(crl_mock, mock_post_requests):
 
     # Stop the task if the pool is emptied
     conn2.close()
+    assert _connections_registry.get_connection_count() == 0
     crl_mock.stop_periodic_cleanup.assert_called_once()