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

sfc-gh-mmishchenko · web-flow · commit a6450a5426f5 · 2025-09-01T18:43:42.000+02:00
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -7,9 +7,11 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
-- v3.18(TBD)
+- v3.17.3(TBD)
   - Enhanced configuration file permission warning messages.
     - Improved warning messages for readable permission issues to include clear instructions on how to skip warnings using the `SF_SKIP_WARNING_FOR_READ_PERMISSIONS_ON_CONFIG_FILE` environment variable.
+  - Fixed the bug with staging pandas dataframes on AWS - the regional endpoint is used when required
+    - This addresses the issue with `create_dataframe` call on Snowpark
 
 - v3.17.2(August 23,2025)
   - Fixed a bug where platform_detection was retrying failed requests with warnings to non-existent endpoints.
diff --git a/src/snowflake/connector/cursor.py b/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()
@@ -1786,8 +1778,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()
 
@@ -1801,14 +1791,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())
@@ -1829,7 +1814,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()
@@ -1844,15 +1828,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())
@@ -1899,7 +1878,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()
@@ -1915,19 +1893,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):
diff --git a/test/unit/test_cursor.py b/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()
