[async] Applied #2513 to async code

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_cursor.py

@@ -24,6 +24,7 @@
 )
 from snowflake.connector._sql_util import get_file_transfer_type
 from snowflake.connector.aio._bind_upload_agent import BindUploadAgent
+from snowflake.connector.aio._file_transfer_agent import SnowflakeFileTransferAgent
 from snowflake.connector.aio._result_batch import (
     ResultBatch,
     create_batches_from_response,
@@ -664,11 +665,8 @@ async 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,
@@ -684,9 +682,6 @@ async 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,
-                    unsafe_file_write=self._connection.unsafe_file_write,
-                    reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
                 )
                 await sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()
@@ -1082,8 +1077,6 @@ async 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()
 
@@ -1097,11 +1090,9 @@ async 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,
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
         )
         await file_transfer_agent.execute()
         await self._init_result_and_meta(file_transfer_agent.result())
@@ -1122,8 +1113,6 @@ async 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()
 
@@ -1137,12 +1126,10 @@ async 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
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
         )
         await file_transfer_agent.execute()
         await self._init_result_and_meta(file_transfer_agent.result())
@@ -1191,8 +1178,6 @@ async 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()
 
@@ -1207,13 +1192,11 @@ async 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 should respect user decision on overwriting
-            reraise_error_in_file_transfer_work_function=self.connection._reraise_error_in_file_transfer_work_function,
         )
         await file_transfer_agent.execute()
         await self._init_result_and_meta(file_transfer_agent.result())
@@ -1321,6 +1304,24 @@ async def query_result(self, qid: str) -> SnowflakeCursor:
             )
         return self
 
+    def _create_file_transfer_agent(
+        self,
+        command: str,
+        ret: dict[str, Any],
+        /,
+        **kwargs,
+    ) -> SnowflakeFileTransferAgent:
+
+        return SnowflakeFileTransferAgent(
+            self,
+            command,
+            ret,
+            use_s3_regional_url=self._connection.enable_stage_s3_privatelink_for_us_east_1,
+            unsafe_file_write=self._connection.unsafe_file_write,
+            reraise_error_in_file_transfer_work_function=self._connection._reraise_error_in_file_transfer_work_function,
+            **kwargs,
+        )
+
 
 class DictCursor(DictCursorSync, SnowflakeCursor):
     pass

test/unit/aio/test_cursor_async_unit.py

@@ -29,6 +29,8 @@ 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._unsafe_file_write = False
 
 
 @pytest.mark.parametrize(
@@ -109,7 +111,7 @@ async def mock_cmd_query(*args, **kwargs):
 class TestUploadDownloadMethods(IsolatedAsyncioTestCase):
     """Test the _upload/_download/_upload_stream/_download_stream methods."""
 
-    @patch("snowflake.connector.aio._file_transfer_agent.SnowflakeFileTransferAgent")
+    @patch("snowflake.connector.aio._cursor.SnowflakeFileTransferAgent")
     async def test_download(self, MockFileTransferAgent):
         cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
             MockFileTransferAgent
@@ -125,9 +127,11 @@ async 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.aio._file_transfer_agent.SnowflakeFileTransferAgent")
+    @patch("snowflake.connector.aio._cursor.SnowflakeFileTransferAgent")
     async def test_upload(self, MockFileTransferAgent):
         cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
             MockFileTransferAgent
@@ -143,9 +147,11 @@ async 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.aio._file_transfer_agent.SnowflakeFileTransferAgent")
+    @patch("snowflake.connector.aio._cursor.SnowflakeFileTransferAgent")
     async def test_download_stream(self, MockFileTransferAgent):
         cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
             MockFileTransferAgent
@@ -161,9 +167,10 @@ async 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.aio._file_transfer_agent.SnowflakeFileTransferAgent")
+    @patch("snowflake.connector.aio._cursor.SnowflakeFileTransferAgent")
     async def test_upload_stream(self, MockFileTransferAgent):
         cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
             MockFileTransferAgent
@@ -180,6 +187,8 @@ async 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):
@@ -191,6 +200,9 @@ def _setup_mocks(self, MockFileTransferAgent):
         fake_conn._file_operation_parser.parse_file_operation = AsyncMock()
         fake_conn._stream_downloader = MagicMock()
         fake_conn._stream_downloader.download_as_stream = AsyncMock()
+        # 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._unsafe_file_write = False
 
         cursor = SnowflakeCursor(fake_conn)
         cursor.reset = MagicMock()