SNOW-2216803 allow re-raising error in file transfer work function in main thread (#2443)

Author: sfc-gh-zyao

src/snowflake/connector/connection.py

@@ -389,6 +389,10 @@ def _get_private_bytes_from_file(
         _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,  # default value
         int,  # type
     ),  # snowflake internal
+    "reraise_error_in_file_transfer_work_function": (
+        False,
+        bool,
+    ),
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")

src/snowflake/connector/cursor.py

@@ -1090,6 +1090,7 @@ def execute(
                     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()
@@ -1807,6 +1808,7 @@ def _download(
             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())
@@ -1850,6 +1852,7 @@ def _upload(
             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())
@@ -1921,6 +1924,7 @@ def _upload_stream(
             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())

src/snowflake/connector/file_transfer_agent.py

@@ -357,6 +357,7 @@ def __init__(
         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,
+        reraise_error_in_file_transfer_work_function: bool = False,
     ) -> None:
         self._cursor = cursor
         self._command = command
@@ -392,6 +393,9 @@ def __init__(
         self._snowflake_server_dop_cap_for_file_transfer = (
             snowflake_server_dop_cap_for_file_transfer
         )
+        self._reraise_error_in_file_transfer_work_function = (
+            reraise_error_in_file_transfer_work_function
+        )
 
     def execute(self) -> None:
         self._parse_command()
@@ -471,6 +475,7 @@ def transfer(self, metas: list[SnowflakeFileMeta]) -> None:
         transfer_metadata = TransferMetadata()  # this is protected by cv_chunk_process
         is_upload = self._command_type == CMD_TYPE_UPLOAD
         exception_caught_in_callback: Exception | None = None
+        exception_caught_in_work: Exception | None = None
         logger.debug(
             "Going to %sload %d files", "up" if is_upload else "down", len(metas)
         )
@@ -626,6 +631,17 @@ def function_and_callback_wrapper(
                 logger.error(f"An exception was raised in {repr(work)}", exc_info=True)
                 file_meta.error_details = e
                 result = (False, e)
+                # If the reraise is enabled, notify the main thread of work
+                # function error, with the concrete exception stored aside in
+                # exception_caught_in_work, such that towards the end of
+                # the transfer call, we reraise the error as is immediately
+                # instead of continuing the execution after transfer.
+                if self._reraise_error_in_file_transfer_work_function:
+                    with cv_main_thread:
+                        nonlocal exception_caught_in_work
+                        exception_caught_in_work = e
+                        cv_main_thread.notify()
+
             try:
                 _callback(*result, file_meta)
             except Exception as e:
@@ -670,6 +686,10 @@ def function_and_callback_wrapper(
         with cv_main_thread:
             while transfer_metadata.num_files_completed < num_total_files:
                 cv_main_thread.wait()
+                # If both exception_caught_in_work and exception_caught_in_callback
+                # are present, the former will take precedence.
+                if exception_caught_in_work is not None:
+                    raise exception_caught_in_work
                 if exception_caught_in_callback is not None:
                     raise exception_caught_in_callback
 

test/unit/test_connection.py

@@ -773,3 +773,41 @@ def test_single_use_refresh_tokens_option_is_plumbed_into_authbyauthcode(
             oauth_enable_single_use_refresh_tokens=rtr_enabled,
         )
         assert conn.auth_class._enable_single_use_refresh_tokens == rtr_enabled
+
+
+# Skip for old drivers because the connection config of
+# reraise_error_in_file_transfer_work_function is newly introduced.
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize("reraise_enabled", [True, False, None])
+def test_reraise_error_in_file_transfer_work_function_config(
+    reraise_enabled: bool | None,
+):
+    """Test that reraise_error_in_file_transfer_work_function config is
+    properly set on connection."""
+
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful._post_request",
+        return_value={
+            "data": {
+                "serverVersion": "a.b.c",
+            },
+            "code": None,
+            "message": None,
+            "success": True,
+        },
+    ):
+        if reraise_enabled is not None:
+            # Create a connection with the config set to the value of reraise_enabled.
+            conn = fake_connector(
+                **{"reraise_error_in_file_transfer_work_function": reraise_enabled}
+            )
+        else:
+            # Special test setup: when reraise_enabled is None, create a
+            # connection without setting the config.
+            conn = fake_connector()
+
+        # When reraise_enabled is None, we expect a default value of False,
+        # so taking bool() on it also makes sense.
+        expected_value = bool(reraise_enabled)
+        actual_value = conn._reraise_error_in_file_transfer_work_function
+        assert actual_value == expected_value

test/unit/test_cursor.py

@@ -24,6 +24,7 @@ class FakeConnection(SnowflakeConnection):
     def __init__(self):
         self._log_max_query_length = 0
         self._reuse_results = None
+        self._reraise_error_in_file_transfer_work_function = False
 
 
 @pytest.mark.parametrize(

test/unit/test_put_get.py

@@ -348,3 +348,136 @@ def test_server_dop_cap(tmp_path):
     # 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
+
+
+def _setup_test_for_reraise_file_transfer_work_fn_error(tmp_path, reraise_param_value):
+    """Helper function to set up common test infrastructure for tests related to re-raising file transfer work function error.
+
+    Returns:
+        tuple: (agent, test_exception, mock_client, mock_create_client)
+    """
+
+    file1 = tmp_path / "file1"
+    file1.write_text("test content")
+
+    # Mock cursor with connection attribute
+    mock_cursor = mock.MagicMock(autospec=SnowflakeCursor)
+    mock_cursor.connection._reraise_error_in_file_transfer_work_function = (
+        reraise_param_value
+    )
+
+    # Create file transfer agent
+    agent = SnowflakeFileTransferAgent(
+        mock_cursor,
+        "PUT some_file.txt",
+        {
+            "data": {
+                "command": "UPLOAD",
+                "src_locations": [str(file1)],
+                "sourceCompression": "none",
+                "parallel": 1,
+                "stageInfo": {
+                    "creds": {
+                        "AZURE_SAS_TOKEN": "sas_token",
+                    },
+                    "location": "some_bucket",
+                    "region": "no_region",
+                    "locationType": "AZURE",
+                    "path": "remote_loc",
+                    "endPoint": "",
+                    "storageAccount": "storage_account",
+                },
+            },
+            "success": True,
+        },
+        reraise_error_in_file_transfer_work_function=reraise_param_value,
+    )
+
+    # Quick check to make sure the field _reraise_error_in_file_transfer_work_function is correctly populated
+    assert (
+        agent._reraise_error_in_file_transfer_work_function == reraise_param_value
+    ), f"expected {reraise_param_value}, got {agent._reraise_error_in_file_transfer_work_function}"
+
+    # Parse command and initialize file metadata
+    agent._parse_command()
+    agent._init_file_metadata()
+    agent._process_file_compression_type()
+
+    # Create a custom exception to be raised by the work function
+    test_exception = Exception("Test work function failure")
+
+    def mock_upload_chunk_with_delay(*args, **kwargs):
+        import time
+
+        time.sleep(0.2)
+        raise test_exception
+
+    # Set up mock client patch, which we will activate in each unit test case.
+    mock_create_client = mock.patch.object(agent, "_create_file_transfer_client")
+    mock_client = mock.MagicMock()
+    mock_client.upload_chunk.side_effect = mock_upload_chunk_with_delay
+
+    # Set up mock client attributes needed for the transfer flow
+    mock_client.meta = agent._file_metadata[0]
+    mock_client.num_of_chunks = 1
+    mock_client.successful_transfers = 0
+    mock_client.failed_transfers = 0
+    mock_client.lock = mock.MagicMock()
+    # Mock methods that would be called during cleanup
+    mock_client.finish_upload = mock.MagicMock()
+    mock_client.delete_client_data = mock.MagicMock()
+
+    return agent, test_exception, mock_client, mock_create_client
+
+
+# Skip for old drivers because the connection config of
+# reraise_error_in_file_transfer_work_function is newly introduced.
+@pytest.mark.skipolddriver
+def test_python_reraise_file_transfer_work_fn_error_as_is(tmp_path):
+    """Tests that when reraise_error_in_file_transfer_work_function config is True,
+    exceptions are reraised immediately without continuing execution after transfer().
+    """
+    agent, test_exception, mock_client, mock_create_client_patch = (
+        _setup_test_for_reraise_file_transfer_work_fn_error(tmp_path, True)
+    )
+
+    with mock_create_client_patch as mock_create_client:
+        mock_create_client.return_value = mock_client
+
+        # Test that with the connection config
+        # reraise_error_in_file_transfer_work_function is True, the
+        # exception is reraised immediately in main thread of transfer.
+        with pytest.raises(Exception) as exc_info:
+            agent.transfer(agent._file_metadata)
+
+        # Verify it's the same exception we injected
+        assert exc_info.value is test_exception
+
+        # Verify that prepare_upload was called (showing the work function was executed)
+        mock_client.prepare_upload.assert_called_once()
+
+
+# Skip for old drivers because the connection config of
+# reraise_error_in_file_transfer_work_function is newly introduced.
+@pytest.mark.skipolddriver
+def test_python_not_reraise_file_transfer_work_fn_error_as_is(tmp_path):
+    """Tests that when reraise_error_in_file_transfer_work_function config is False (default),
+    where exceptions are stored in file metadata but execution continues.
+    """
+    agent, test_exception, mock_client, mock_create_client_patch = (
+        _setup_test_for_reraise_file_transfer_work_fn_error(tmp_path, False)
+    )
+
+    with mock_create_client_patch as mock_create_client:
+        mock_create_client.return_value = mock_client
+
+        # Verify that with the connection config
+        # reraise_error_in_file_transfer_work_function is False, the
+        # exception is not reraised (but instead stored in file metadata).
+        agent.transfer(agent._file_metadata)
+
+        # Verify that the error was stored in the file metadata
+        assert agent._file_metadata[0].error_details is test_exception
+
+        # Verify that prepare_upload was called
+        mock_client.prepare_upload.assert_called_once()