From 2591fba4c3c9c799adffe683550333c2b6c324c4 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 27 Jun 2026 17:04:24 +0530 Subject: [PATCH] Close FTP file handles when a transfer raises retrieve_file and store_file opened a local file, then ran the FTP transfer before calling close(). When the transfer raised, the handle was never closed and leaked for the lifetime of the process. Wrapping the transfer in try/finally guarantees the handle we opened is closed on every path, while still leaving caller-supplied buffers open. closes: #69077 --- .../src/airflow/providers/ftp/hooks/ftp.py | 27 ++++++++++++------- .../ftp/tests/unit/ftp/hooks/test_ftp.py | 18 +++++++++++++ 2 files changed, 35 insertions(+), 10 deletions(-) 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()