diff --git a/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py b/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py index 922f22637f570..04052e6b1426a 100644 --- a/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py +++ b/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py @@ -184,6 +184,7 @@ def write_to_file_with_progress(data): """ conn = self.get_conn() is_path = isinstance(local_full_path_or_buffer, str) + output_handle = None # without a callback, default to writing to a user-provided file or # file-like buffer @@ -195,12 +196,15 @@ def write_to_file_with_progress(data): callback = output_handle.write - self.log.info("Retrieving file from FTP: %s", remote_full_path) - conn.retrbinary(f"RETR {remote_full_path}", callback, block_size) - self.log.info("Finished retrieving file from FTP: %s", remote_full_path) - - if is_path and output_handle: - output_handle.close() + try: + self.log.info("Retrieving file from FTP: %s", remote_full_path) + conn.retrbinary(f"RETR {remote_full_path}", callback, block_size) + self.log.info("Finished retrieving file from FTP: %s", remote_full_path) + finally: + # Only close handles we opened ourselves; a caller-supplied buffer + # must stay open per this method's contract. + if is_path and output_handle: + output_handle.close() def store_file( self, remote_full_path: str, local_full_path_or_buffer: Any, block_size: int = 8192 @@ -226,10 +230,13 @@ def store_file( else: input_handle = local_full_path_or_buffer - conn.storbinary(f"STOR {remote_full_path}", input_handle, block_size) - - if is_path: - input_handle.close() + try: + conn.storbinary(f"STOR {remote_full_path}", input_handle, block_size) + finally: + # Only close handles we opened ourselves; a caller-supplied buffer + # must stay open per this method's contract. + if is_path: + input_handle.close() def delete_file(self, path: str) -> None: """ diff --git a/providers/ftp/tests/unit/ftp/hooks/test_ftp.py b/providers/ftp/tests/unit/ftp/hooks/test_ftp.py index 0a0fcee183be8..0cfad429c5fef 100644 --- a/providers/ftp/tests/unit/ftp/hooks/test_ftp.py +++ b/providers/ftp/tests/unit/ftp/hooks/test_ftp.py @@ -126,6 +126,24 @@ def test_retrieve_file_with_callback(self): ftp_hook.retrieve_file(self.path, _buffer, callback=func) self.conn_mock.retrbinary.assert_called_once_with("RETR /some/path", func, 8192) + def test_retrieve_file_closes_handle_on_error(self): + self.conn_mock.retrbinary.side_effect = OSError("transfer failed") + handle = mock.MagicMock(name="output_handle") + with mock.patch("builtins.open", return_value=handle): + with fh.FTPHook() as ftp_hook: + with pytest.raises(OSError, match="transfer failed"): + ftp_hook.retrieve_file(self.path, "/local/path") + handle.close.assert_called_once_with() + + def test_store_file_closes_handle_on_error(self): + self.conn_mock.storbinary.side_effect = OSError("transfer failed") + handle = mock.MagicMock(name="input_handle") + with mock.patch("builtins.open", return_value=handle): + with fh.FTPHook() as ftp_hook: + with pytest.raises(OSError, match="transfer failed"): + ftp_hook.store_file(self.path, "/local/path") + handle.close.assert_called_once_with() + def test_connection_success(self): with fh.FTPHook() as ftp_hook: status, msg = ftp_hook.test_connection()