SNOW-2173966 introduce server DoP cap (#2375)

Author: sfc-gh-zyao

src/snowflake/connector/_utils.py

@@ -32,6 +32,15 @@ class TempObjectType(Enum):
 
 REQUEST_ID_STATEMENT_PARAM_NAME = "requestId"
 
+# Default server side cap on Degree of Parallelism for file transfer
+# This default value is set to 2^30 (~ 10^9), such that it will not
+# throttle regular sessions.
+_DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER = 1 << 30
+# Variable name of server DoP cap for file transfer
+_VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER = (
+    "snowflake_server_dop_cap_for_file_transfer"
+)
+
 
 def generate_random_alphanumeric(length: int = 10) -> str:
     return "".join(choice(ALPHANUMERIC) for _ in range(length))
@@ -60,6 +69,15 @@ def is_uuid4(str_or_uuid: str | UUID) -> bool:
     return uuid_str == str_or_uuid
 
 
+def _snowflake_max_parallelism_for_file_transfer(connection):
+    """Returns the server side cap on max parallelism for file transfer for the given connection."""
+    return getattr(
+        connection,
+        f"_{_VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER}",
+        _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
+    )
+
+
 class _TrackedQueryCancellationTimer(Timer):
     def __init__(self, interval, function, args=None, kwargs=None):
         super().__init__(interval, function, args, kwargs)

src/snowflake/connector/connection.py

@@ -29,6 +29,10 @@
 
 from . import errors, proxy
 from ._query_context_cache import QueryContextCache
+from ._utils import (
+    _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
+    _VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
+)
 from .auth import (
     FIRST_PARTY_AUTHENTICATORS,
     Auth,
@@ -369,6 +373,10 @@ def _get_private_bytes_from_file(
         str,
         # SNOW-2096721: External (Spark) session ID
     ),
+    _VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER: (
+        _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,  # default value
+        int,  # type
+    ),  # snowflake internal
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")

src/snowflake/connector/cursor.py

@@ -37,6 +37,7 @@
 from ._sql_util import get_file_transfer_type
 from ._utils import (
     REQUEST_ID_STATEMENT_PARAM_NAME,
+    _snowflake_max_parallelism_for_file_transfer,
     _TrackedQueryCancellationTimer,
     is_uuid4,
 )
@@ -1086,6 +1087,9 @@ def execute(
                     use_s3_regional_url=self._connection.enable_stage_s3_privatelink_for_us_east_1,
                     iobound_tpe_limit=self._connection.iobound_tpe_limit,
                     unsafe_file_write=self._connection.unsafe_file_write,
+                    snowflake_server_dop_cap_for_file_transfer=_snowflake_max_parallelism_for_file_transfer(
+                        self._connection
+                    ),
                 )
                 sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()
@@ -1800,6 +1804,9 @@ def _download(
             self,
             "",  # empty command because it is triggered by directly calling this util not by a SQL query
             ret,
+            snowflake_server_dop_cap_for_file_transfer=_snowflake_max_parallelism_for_file_transfer(
+                self._connection
+            ),
         )
         file_transfer_agent.execute()
         self._init_result_and_meta(file_transfer_agent.result())
@@ -1840,6 +1847,9 @@ def _upload(
             "",  # empty command because it is triggered by directly calling this util not by a SQL query
             ret,
             force_put_overwrite=False,  # _upload should respect user decision on overwriting
+            snowflake_server_dop_cap_for_file_transfer=_snowflake_max_parallelism_for_file_transfer(
+                self._connection
+            ),
         )
         file_transfer_agent.execute()
         self._init_result_and_meta(file_transfer_agent.result())
@@ -1908,6 +1918,9 @@ def _upload_stream(
             ret,
             source_from_stream=input_stream,
             force_put_overwrite=False,  # _upload_stream should respect user decision on overwriting
+            snowflake_server_dop_cap_for_file_transfer=_snowflake_max_parallelism_for_file_transfer(
+                self._connection
+            ),
         )
         file_transfer_agent.execute()
         self._init_result_and_meta(file_transfer_agent.result())

src/snowflake/connector/file_transfer_agent.py

@@ -15,6 +15,7 @@
 from time import time
 from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar
 
+from ._utils import _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER
 from .azure_storage_client import SnowflakeAzureRestClient
 from .compat import IS_WINDOWS
 from .constants import (
@@ -355,6 +356,7 @@ def __init__(
         use_s3_regional_url: bool = False,
         iobound_tpe_limit: int | None = None,
         unsafe_file_write: bool = False,
+        snowflake_server_dop_cap_for_file_transfer=_DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
     ) -> None:
         self._cursor = cursor
         self._command = command
@@ -387,6 +389,9 @@ def __init__(
         self._credentials: StorageCredential | None = None
         self._iobound_tpe_limit = iobound_tpe_limit
         self._unsafe_file_write = unsafe_file_write
+        self._snowflake_server_dop_cap_for_file_transfer = (
+            snowflake_server_dop_cap_for_file_transfer
+        )
 
     def execute(self) -> None:
         self._parse_command()
@@ -443,12 +448,16 @@ def execute(self) -> None:
             result.result_status = result.result_status.value
 
     def transfer(self, metas: list[SnowflakeFileMeta]) -> None:
-        iobound_tpe_limit = min(len(metas), os.cpu_count())
+        iobound_tpe_limit = min(
+            len(metas), os.cpu_count(), self._snowflake_server_dop_cap_for_file_transfer
+        )
         logger.debug("Decided IO-bound TPE size: %d", iobound_tpe_limit)
         if self._iobound_tpe_limit is not None:
             logger.debug("IO-bound TPE size is limited to: %d", self._iobound_tpe_limit)
             iobound_tpe_limit = min(iobound_tpe_limit, self._iobound_tpe_limit)
-        max_concurrency = self._parallel
+        max_concurrency = min(
+            self._parallel, self._snowflake_server_dop_cap_for_file_transfer
+        )
         network_tpe = ThreadPoolExecutor(max_concurrency)
         preprocess_tpe = ThreadPoolExecutor(iobound_tpe_limit)
         postprocess_tpe = ThreadPoolExecutor(iobound_tpe_limit)

test/unit/test_cursor.py

@@ -189,3 +189,60 @@ def _setup_mocks(self, MockFileTransferAgent):
         cursor.reset = MagicMock()
         cursor._init_result_and_meta = MagicMock()
         return cursor, fake_conn, mock_file_transfer_agent_instance
+
+    def _run_dop_cap_test(self, task, dop_cap):
+        """A helper to run dop cap test.
+
+        It mainly verifies that when performing the specified task, we are using a FileTransferAgent with DoP cap as specified.
+        """
+        from snowflake.connector._utils import (
+            _VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
+        )
+
+        mock_conn = FakeConnection()
+        setattr(
+            mock_conn, f"_{_VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER}", dop_cap
+        )
+
+        class FakeFileOperationParser:
+            def parse_file_operation(
+                self,
+                stage_location,
+                local_file_name,
+                target_directory,
+                command_type,
+                options,
+                has_source_from_stream=False,
+            ):
+                return {}
+
+        mock_cursor = SnowflakeCursor(mock_conn)
+        mock_conn._file_operation_parser = FakeFileOperationParser()
+        with patch.object(
+            mock_cursor, "_init_result_and_meta", return_value=None
+        ), patch(
+            "snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent"
+        ) as MockFileTransferAgent:
+            task(mock_cursor)
+            # Verify that when running the file operation, we are using FileTransferAgent with server DoP cap as 1.
+            _, kwargs = MockFileTransferAgent.call_args
+            assert dop_cap == kwargs["snowflake_server_dop_cap_for_file_transfer"]
+
+    def test_dop_cap_for_upload(self):
+        def task(cursor):
+            cursor._upload("/tmp/test.txt", "@st", {})
+
+        self._run_dop_cap_test(task, dop_cap=1)
+
+    def test_dop_cap_for_upload_stream(self):
+        def task(cursor):
+            mock_input_stream = MagicMock()
+            cursor._upload_stream(mock_input_stream, "@st", {})
+
+        self._run_dop_cap_test(task, dop_cap=1)
+
+    def test_dop_cap_for_download(self):
+        def task(cursor):
+            cursor._download("@st", "/tmp", {})
+
+        self._run_dop_cap_test(task, dop_cap=1)

test/unit/test_put_get.py

@@ -293,3 +293,58 @@ def test_strip_stage_prefix_from_dst_file_name_for_download():
         agent._strip_stage_prefix_from_dst_file_name_for_download.assert_called_with(
             file
         )
+
+
+# The server DoP cap is newly introduced and therefore should not be tested in
+# old drivers.
+@pytest.mark.skipolddriver
+def test_server_dop_cap(tmp_path):
+    file1 = tmp_path / "file1"
+    file2 = tmp_path / "file2"
+    file1.touch()
+    file2.touch()
+    # Positive case
+    rest_client = SnowflakeFileTransferAgent(
+        mock.MagicMock(autospec=SnowflakeCursor),
+        "PUT some_file.txt",
+        {
+            "data": {
+                "command": "UPLOAD",
+                "src_locations": [file1, file2],
+                "sourceCompression": "none",
+                "parallel": 8,
+                "stageInfo": {
+                    "creds": {},
+                    "location": "some_bucket",
+                    "region": "no_region",
+                    "locationType": "AZURE",
+                    "path": "remote_loc",
+                    "endPoint": "",
+                    "storageAccount": "storage_account",
+                },
+            },
+            "success": True,
+        },
+        snowflake_server_dop_cap_for_file_transfer=1,
+    )
+    with mock.patch(
+        "snowflake.connector.file_transfer_agent.ThreadPoolExecutor"
+    ) as tpe:
+        with mock.patch("snowflake.connector.file_transfer_agent.threading.Condition"):
+            with mock.patch(
+                "snowflake.connector.file_transfer_agent.TransferMetadata",
+                return_value=mock.Mock(
+                    num_files_started=0,
+                    num_files_completed=3,
+                ),
+            ):
+                try:
+                    rest_client.execute()
+                except AttributeError:
+                    pass
+
+    # We expect 3 thread pool executors to be created with thread count as 1,
+    # because we will create executors for network, preprocess and postprocess,
+    # and due to the server DoP cap, each of them will have a thread count
+    # of 1.
+    assert len(list(filter(lambda e: e.args == (1,), tpe.call_args_list))) == 3