SNOW-2028051 introduce a new client_fetch_threads connection parameter to decouple threads number limitations on fetching and pre-fetching (#2255)

Author: sfc-gh-mmishchenko

src/snowflake/connector/connection.py

@@ -120,6 +120,7 @@
 
 DEFAULT_CLIENT_PREFETCH_THREADS = 4
 MAX_CLIENT_PREFETCH_THREADS = 10
+MAX_CLIENT_FETCH_THREADS = 1024
 DEFAULT_BACKOFF_POLICY = exponential_backoff()
 
 
@@ -222,6 +223,7 @@ def _get_private_bytes_from_file(
         (type(None), int),
     ),  # snowflake
     "client_prefetch_threads": (4, int),  # snowflake
+    "client_fetch_threads": (None, (type(None), int)),
     "numpy": (False, bool),  # snowflake
     "ocsp_response_cache_filename": (None, (type(None), str)),  # snowflake internal
     "converter_class": (DefaultConverterClass(), SnowflakeConverter),
@@ -380,6 +382,7 @@ class SnowflakeConnection:
             See the backoff_policies module for details and implementation examples.
         client_session_keep_alive_heartbeat_frequency: Heartbeat frequency to keep connection alive in seconds.
         client_prefetch_threads: Number of threads to download the result set.
+        client_fetch_threads: Number of threads to fetch staged query results.
         rest: Snowflake REST API object. Internal use only. Maybe removed in a later release.
         application: Application name to communicate with Snowflake as. By default, this is "PythonConnector".
         errorhandler: Handler used with errors. By default, an exception will be raised on error.
@@ -639,6 +642,16 @@ def client_prefetch_threads(self, value) -> None:
         self._client_prefetch_threads = value
         self._validate_client_prefetch_threads()
 
+    @property
+    def client_fetch_threads(self) -> int | None:
+        return self._client_fetch_threads
+
+    @client_fetch_threads.setter
+    def client_fetch_threads(self, value: None | int) -> None:
+        if value is not None:
+            value = min(max(1, value), MAX_CLIENT_FETCH_THREADS)
+        self._client_fetch_threads = value
+
     @property
     def rest(self) -> SnowflakeRestful | None:
         return self._rest

src/snowflake/connector/cursor.py

@@ -1186,7 +1186,8 @@ def _init_result_and_meta(self, data: dict[Any, Any]) -> None:
         self._result_set = ResultSet(
             self,
             result_chunks,
-            self._connection.client_prefetch_threads,
+            self._connection.client_fetch_threads
+            or self._connection.client_prefetch_threads,
         )
         self._rownumber = -1
         self._result_state = ResultState.VALID

test/integ/test_connection.py

@@ -971,6 +971,24 @@ def test_client_prefetch_threads_setting(conn_cnx):
         assert conn.client_prefetch_threads == new_thread_count
 
 
+@pytest.mark.skipolddriver
+def test_client_fetch_threads_setting(conn_cnx):
+    """Tests whether client_fetch_threads is None by default and setting the parameter has effect."""
+    with conn_cnx() as conn:
+        assert conn.client_fetch_threads is None
+        conn.client_fetch_threads = 32
+        assert conn.client_fetch_threads == 32
+
+
+@pytest.mark.external
+def test_client_failover_connection_url(conn_cnx):
+    with conn_cnx("client_failover") as conn:
+        with conn.cursor() as cur:
+            assert cur.execute("select 1;").fetchall() == [
+                (1,),
+            ]
+
+
 def test_connection_gc(conn_cnx):
     """This test makes sure that a heartbeat thread doesn't prevent garbage collection of SnowflakeConnection."""
     conn = conn_cnx(client_session_keep_alive=True).__enter__()