From 8dddd5a206078f60d65391d324355e26de288d61 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Thu, 6 Nov 2025 10:29:55 +0100 Subject: [PATCH 1/5] NO-SNOW: Stop CRL cache background task when closing conn --- src/snowflake/connector/connection.py | 4 ++++ src/snowflake/connector/crl.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 75f6451777..4ec1e9495a 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -88,6 +88,7 @@ ) from .converter import SnowflakeConverter from .crl import CRLConfig +from .crl_cache import CRLCacheFactory from .cursor import LOG_MAX_QUERY_LENGTH, SnowflakeCursor, SnowflakeCursorBase from .description import ( CLIENT_NAME, @@ -1157,6 +1158,9 @@ def close(self, retry: bool = True) -> None: # unregister to dereference connection object as it's already closed after the execution atexit.unregister(self._close_at_exit) try: + # Stop CRL-related background process + CRLCacheFactory.stop_periodic_cleanup() + if not self.rest: logger.debug("Rest object has been destroyed, cannot close session") return diff --git a/src/snowflake/connector/crl.py b/src/snowflake/connector/crl.py index f3892f5952..e6af12c412 100644 --- a/src/snowflake/connector/crl.py +++ b/src/snowflake/connector/crl.py @@ -63,7 +63,7 @@ class CRLConfig: crl_cache_dir: Path | str | None = None crl_cache_removal_delay_days: int = 7 crl_cache_cleanup_interval_hours: int = 1 - crl_cache_start_cleanup: bool = False + crl_cache_start_cleanup: bool = True @classmethod def from_connection(cls, sf_connection) -> CRLConfig: From a6369fd848d283eb8a84e86f36c269948c576e06 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Thu, 6 Nov 2025 14:09:27 +0100 Subject: [PATCH 2/5] NO-SNOW: Add thread-safe pool for tracking opened connections --- src/snowflake/connector/connection.py | 57 ++++++++++++++++++++++++++- test/unit/test_connection.py | 49 +++++++++++++++++++++++ 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 4ec1e9495a..03d543c6ba 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -1153,13 +1153,16 @@ def connect(self, **kwargs) -> None: else: self.__open_connection() + # Register the connection in the pool after successful connection + _connections_pool.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: - # Stop CRL-related background process - CRLCacheFactory.stop_periodic_cleanup() + # Remove connection from the pool + _connections_pool.remove_connection(self) if not self.rest: logger.debug("Rest object has been destroyed, cannot close session") @@ -2537,3 +2540,53 @@ def _detect_application() -> None | str: return "jupyter_notebook" if "snowbooks" in sys.modules: return "snowflake_notebook" + + +class _ConnectionsPool: + """Thread-safe pool 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.""" + self._connections: weakref.WeakSet = weakref.WeakSet() + self._lock = Lock() + + def add_connection(self, connection: SnowflakeConnection) -> None: + """Add a connection to the pool. + + Args: + connection: The SnowflakeConnection instance to register. + """ + with self._lock: + self._connections.add(connection) + logger.debug( + f"Connection {id(connection)} added to pool. Total connections: {len(self._connections)}" + ) + + def remove_connection(self, connection: SnowflakeConnection) -> None: + """Remove a connection from the pool. + + 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)}" + ) + + if len(self._connections) == 0: + # 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() diff --git a/test/unit/test_connection.py b/test/unit/test_connection.py index d9c29046a0..1e531d83a7 100644 --- a/test/unit/test_connection.py +++ b/test/unit/test_connection.py @@ -953,3 +953,52 @@ def test_connect_metadata_preservation(): len(params) > 0 ), "connect should have parameters from SnowflakeConnection.__init__" # Should have parameters like account, user, password, etc. + + +def test_connections_pool(mock_post_requests): + """Test that connections are properly tracked in the _ConnectionsPool.""" + from snowflake.connector.connection import _connections_pool + + # Get initial connection count + initial_count = _connections_pool.get_connection_count() + + # Create and connect first connection + conn1 = fake_connector() + assert ( + _connections_pool.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 + ), "Connection count should increase by 2 after creating two connections" + + # Close first connection + conn1.close() + assert ( + _connections_pool.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 + ), "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): + """Test the individual methods of _ConnectionsPool.""" + + # Create a connection + conn1 = fake_connector() + conn2 = fake_connector() + + # Don't stop the task if pool is not empty + conn1.close() + crl_mock.stop_periodic_cleanup.assert_not_called() + + # Stop the task if the pool is emptied + conn2.close() + crl_mock.stop_periodic_cleanup.assert_called_once() From 02616cfc8d146d3151ffc6a39934b77c91c0fef6 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Thu, 6 Nov 2025 15:10:52 +0100 Subject: [PATCH 3/5] fixup! NO-SNOW: Add thread-safe pool for tracking opened connections --- src/snowflake/connector/connection.py | 27 +++++++++++++++------------ test/unit/test_connection.py | 18 ++++++++++-------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 03d543c6ba..06c8299823 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -1154,7 +1154,7 @@ 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.""" @@ -1162,7 +1162,7 @@ def close(self, retry: bool = True) -> None: 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,7 +2567,7 @@ 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. @@ -2575,13 +2575,16 @@ def remove_connection(self, connection: SnowflakeConnection) -> None: 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: @@ -2589,4 +2592,4 @@ def get_connection_count(self) -> int: # Global instance of the connections pool -_connections_pool = _ConnectionsPool() +_connections_registry = _ConnectionsRegistry() diff --git a/test/unit/test_connection.py b/test/unit/test_connection.py index 1e531d83a7..1e1d75cea5 100644 --- a/test/unit/test_connection.py +++ b/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() From c39e9899ea80740da495894c579ebf317d1524b8 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Wed, 12 Nov 2025 10:14:23 +0100 Subject: [PATCH 4/5] fixup! fixup! NO-SNOW: Add thread-safe pool for tracking opened connections --- test/unit/test_connection.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/test/unit/test_connection.py b/test/unit/test_connection.py index 1e1d75cea5..086cfa2348 100644 --- a/test/unit/test_connection.py +++ b/test/unit/test_connection.py @@ -990,17 +990,22 @@ def test_connections_registry(mock_post_requests): @mock.patch("snowflake.connector.connection.CRLCacheFactory") 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() - conn2 = fake_connector() + from snowflake.connector.connection import _ConnectionsRegistry - # Don't stop the task if pool is not empty - conn1.close() - crl_mock.stop_periodic_cleanup.assert_not_called() - - # 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() + # Mock the registry to avoid side effects from other tests due to _ConnectionsRegistry being a singleton + with mock.patch( + "snowflake.connector.connection._connections_registry", _ConnectionsRegistry() + ) as mock_registry: + # Create a connection + conn1 = fake_connector() + conn2 = fake_connector() + assert mock_registry.get_connection_count() == 2 + + # Don't stop the task if pool is not empty + conn1.close() + crl_mock.stop_periodic_cleanup.assert_not_called() + + # Stop the task if the pool is emptied + conn2.close() + assert mock_registry.get_connection_count() == 0 + crl_mock.stop_periodic_cleanup.assert_called_once() From cac6251f507e0608b95e1487eae99f1b6cba126f Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Wed, 12 Nov 2025 12:28:15 +0100 Subject: [PATCH 5/5] fixup! fixup! fixup! NO-SNOW: Add thread-safe pool for tracking opened connections --- test/unit/test_connection.py | 35 ++--------------------------------- 1 file changed, 2 insertions(+), 33 deletions(-) diff --git a/test/unit/test_connection.py b/test/unit/test_connection.py index 086cfa2348..36053400a0 100644 --- a/test/unit/test_connection.py +++ b/test/unit/test_connection.py @@ -955,40 +955,8 @@ def test_connect_metadata_preservation(): # Should have parameters like account, user, password, etc. -def test_connections_registry(mock_post_requests): - """Test that connections are properly tracked in the _ConnectionsPool.""" - from snowflake.connector.connection import _connections_registry - - # Get initial connection count - initial_count = _connections_registry.get_connection_count() - - # Create and connect first connection - conn1 = fake_connector() - assert ( - _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_registry.get_connection_count() == initial_count + 2 - ), "Connection count should increase by 2 after creating two connections" - - # Close first connection - conn1.close() - assert ( - _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_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_registry_stops_crl_task_if_empty(crl_mock, mock_post_requests): +def test_connections_registry_lifecycle(crl_mock, mock_post_requests): """Test the individual methods of _ConnectionsPool.""" from snowflake.connector.connection import _ConnectionsRegistry @@ -1004,6 +972,7 @@ def test_connections_registry_stops_crl_task_if_empty(crl_mock, mock_post_reques # Don't stop the task if pool is not empty conn1.close() crl_mock.stop_periodic_cleanup.assert_not_called() + assert mock_registry.get_connection_count() == 1 # Stop the task if the pool is emptied conn2.close()