SNOW-2283945 use AWS regional endpoints when required for storing pandas frames (#2513)

(cherry picked from commit a6450a5)

Author: sfc-gh-mmishchenko

src/snowflake/connector/cursor.py

@@ -79,7 +79,10 @@
     from pyarrow import Table
 
     from .connection import SnowflakeConnection
-    from .file_transfer_agent import SnowflakeProgressPercentage
+    from .file_transfer_agent import (
+        SnowflakeFileTransferAgent,
+        SnowflakeProgressPercentage,
+    )
     from .result_batch import ResultBatch
 
 T = TypeVar("T", bound=collections.abc.Sequence)
@@ -1064,11 +1067,7 @@ def execute(
             )
             logger.debug("PUT OR GET: %s", self.is_file_transfer)
             if self.is_file_transfer:
-                from .file_transfer_agent import SnowflakeFileTransferAgent
-
-                # Decide whether to use the old, or new code path
-                sf_file_transfer_agent = SnowflakeFileTransferAgent(
-                    self,
+                sf_file_transfer_agent = self._create_file_transfer_agent(
                     query,
                     ret,
                     put_callback=_put_callback,
@@ -1084,13 +1083,6 @@ def execute(
                     skip_upload_on_content_match=_skip_upload_on_content_match,
                     source_from_stream=file_stream,
                     multipart_threshold=data.get("threshold"),
-                    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
-                    ),
-                    reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
                 )
                 sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()
@@ -1785,8 +1777,6 @@ def _download(
             _do_reset (bool, optional): Whether to reset the cursor before
                 downloading, by default we will reset the cursor.
         """
-        from .file_transfer_agent import SnowflakeFileTransferAgent
-
         if _do_reset:
             self.reset()
 
@@ -1800,14 +1790,9 @@ def _download(
         )
 
         # Execute the file operation based on the interpretation above.
-        file_transfer_agent = SnowflakeFileTransferAgent(
-            self,
+        file_transfer_agent = self._create_file_transfer_agent(
             "",  # 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
-            ),
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
         )
         file_transfer_agent.execute()
         self._init_result_and_meta(file_transfer_agent.result())
@@ -1828,7 +1813,6 @@ def _upload(
             _do_reset (bool, optional): Whether to reset the cursor before
                 uploading, by default we will reset the cursor.
         """
-        from .file_transfer_agent import SnowflakeFileTransferAgent
 
         if _do_reset:
             self.reset()
@@ -1843,15 +1827,10 @@ def _upload(
         )
 
         # Execute the file operation based on the interpretation above.
-        file_transfer_agent = SnowflakeFileTransferAgent(
-            self,
+        file_transfer_agent = self._create_file_transfer_agent(
             "",  # 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
-            ),
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
         )
         file_transfer_agent.execute()
         self._init_result_and_meta(file_transfer_agent.result())
@@ -1898,7 +1877,6 @@ def _upload_stream(
             _do_reset (bool, optional): Whether to reset the cursor before
                 uploading, by default we will reset the cursor.
         """
-        from .file_transfer_agent import SnowflakeFileTransferAgent
 
         if _do_reset:
             self.reset()
@@ -1914,19 +1892,37 @@ def _upload_stream(
         )
 
         # Execute the file operation based on the interpretation above.
-        file_transfer_agent = SnowflakeFileTransferAgent(
-            self,
+        file_transfer_agent = self._create_file_transfer_agent(
             "",  # empty command because it is triggered by directly calling this util not by a SQL query
             ret,
             source_from_stream=input_stream,
             force_put_overwrite=False,  # _upload_stream should respect user decision on overwriting
+        )
+        file_transfer_agent.execute()
+        self._init_result_and_meta(file_transfer_agent.result())
+
+    def _create_file_transfer_agent(
+        self,
+        command: str,
+        ret: dict[str, Any],
+        /,
+        **kwargs,
+    ) -> SnowflakeFileTransferAgent:
+        from .file_transfer_agent import SnowflakeFileTransferAgent
+
+        return SnowflakeFileTransferAgent(
+            self,
+            command,
+            ret,
+            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
             ),
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
+            reraise_error_in_file_transfer_work_function=self._connection._reraise_error_in_file_transfer_work_function,
+            **kwargs,
         )
-        file_transfer_agent.execute()
-        self._init_result_and_meta(file_transfer_agent.result())
 
 
 class DictCursor(SnowflakeCursor):

test/unit/test_cursor.py

@@ -25,6 +25,9 @@ def __init__(self):
         self._log_max_query_length = 0
         self._reuse_results = None
         self._reraise_error_in_file_transfer_work_function = False
+        self._enable_stage_s3_privatelink_for_us_east_1 = False
+        self._iobound_tpe_limit = None
+        self._unsafe_file_write = False
 
 
 @pytest.mark.parametrize(
@@ -121,6 +124,8 @@ def test_download(self, MockFileTransferAgent):
         #   - download_as_stream of connection._stream_downloader
         fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
         fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        MockFileTransferAgent.assert_called_once()
+        assert MockFileTransferAgent.call_args.kwargs.get("use_s3_regional_url", False)
         mock_file_transfer_agent_instance.execute.assert_called_once()
 
     @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
@@ -139,6 +144,8 @@ def test_upload(self, MockFileTransferAgent):
         #   - download_as_stream of connection._stream_downloader
         fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
         fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        MockFileTransferAgent.assert_called_once()
+        assert MockFileTransferAgent.call_args.kwargs.get("use_s3_regional_url", False)
         mock_file_transfer_agent_instance.execute.assert_called_once()
 
     @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
@@ -157,6 +164,7 @@ def test_download_stream(self, MockFileTransferAgent):
         #   - execute in SnowflakeFileTransferAgent
         fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
         fake_conn._stream_downloader.download_as_stream.assert_called_once()
+        MockFileTransferAgent.assert_not_called()
         mock_file_transfer_agent_instance.execute.assert_not_called()
 
     @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
@@ -176,6 +184,8 @@ def test_upload_stream(self, MockFileTransferAgent):
         #   - download_as_stream of connection._stream_downloader
         fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
         fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        MockFileTransferAgent.assert_called_once()
+        assert MockFileTransferAgent.call_args.kwargs.get("use_s3_regional_url", False)
         mock_file_transfer_agent_instance.execute.assert_called_once()
 
     def _setup_mocks(self, MockFileTransferAgent):
@@ -185,6 +195,10 @@ def _setup_mocks(self, MockFileTransferAgent):
         fake_conn = FakeConnection()
         fake_conn._file_operation_parser = MagicMock()
         fake_conn._stream_downloader = MagicMock()
+        # this should be true on all new AWS deployments to use regional endpoints for staging operations
+        fake_conn._enable_stage_s3_privatelink_for_us_east_1 = True
+        fake_conn._iobound_tpe_limit = 1
+        fake_conn._unsafe_file_write = False
 
         cursor = SnowflakeCursor(fake_conn)
         cursor.reset = MagicMock()