From ff30dcc1e18abf267e4381bcc64a247da3c9af35 Mon Sep 17 00:00:00 2001 From: Makrushin Evgenii Date: Tue, 31 Oct 2023 06:02:54 +0700 Subject: [PATCH] FTPHook methods should not change working directory (#35105) * FTPHook: fix describe_directory, list_directory, retrieve_file and store_file methods, they no change working path anymore * diss: https://github.com/apache/airflow/pull/35105/files/1902fbb666c79b8d3aaef440471b115c108a5c2e#r1367884370 * diss: https://github.com/apache/airflow/pull/35105/files/1902fbb666c79b8d3aaef440471b115c108a5c2e#r1367884193 provide path directly to mlsd method Co-authored-by: Hussein Awala * TestFTPHook: add test_describe_directory test, fix test_list_directory test * diss: https://github.com/apache/airflow/pull/35105/files/1902fbb666c79b8d3aaef440471b115c108a5c2e#r1367884591, https://github.com/apache/airflow/pull/35105/files/1902fbb666c79b8d3aaef440471b115c108a5c2e#r1367884629 * FTPHook: change mlsd call to fit with test assertion * FTPHook: fix static check --------- Co-authored-by: e.makrushin Co-authored-by: Hussein Awala --- airflow/providers/ftp/hooks/ftp.py | 19 +++++-------------- tests/providers/ftp/hooks/test_ftp.py | 13 +++++++++---- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/airflow/providers/ftp/hooks/ftp.py b/airflow/providers/ftp/hooks/ftp.py index aea5441557de3..1935a698638a5 100644 --- a/airflow/providers/ftp/hooks/ftp.py +++ b/airflow/providers/ftp/hooks/ftp.py @@ -19,7 +19,6 @@ import datetime import ftplib -import os from typing import Any, Callable from airflow.hooks.base import BaseHook @@ -77,9 +76,7 @@ def describe_directory(self, path: str) -> dict: :param path: full path to the remote directory """ conn = self.get_conn() - conn.cwd(path) - files = dict(conn.mlsd()) - return files + return dict(conn.mlsd(path)) def list_directory(self, path: str) -> list[str]: """ @@ -88,10 +85,7 @@ def list_directory(self, path: str) -> list[str]: :param path: full path to the remote directory to list """ conn = self.get_conn() - conn.cwd(path) - - files = conn.nlst() - return files + return conn.nlst(path) def create_directory(self, path: str) -> None: """ @@ -181,10 +175,8 @@ def write_to_file_with_progress(data): callback = output_handle.write - remote_path, remote_file_name = os.path.split(remote_full_path) - conn.cwd(remote_path) self.log.info("Retrieving file from FTP: %s", remote_full_path) - conn.retrbinary(f"RETR {remote_file_name}", callback, block_size) + 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: @@ -213,9 +205,8 @@ def store_file( input_handle = open(local_full_path_or_buffer, "rb") else: input_handle = local_full_path_or_buffer - remote_path, remote_file_name = os.path.split(remote_full_path) - conn.cwd(remote_path) - conn.storbinary(f"STOR {remote_file_name}", input_handle, block_size) + + conn.storbinary(f"STOR {remote_full_path}", input_handle, block_size) if is_path: input_handle.close() diff --git a/tests/providers/ftp/hooks/test_ftp.py b/tests/providers/ftp/hooks/test_ftp.py index bdd5c0d90c29c..2f4239b24374a 100644 --- a/tests/providers/ftp/hooks/test_ftp.py +++ b/tests/providers/ftp/hooks/test_ftp.py @@ -47,12 +47,17 @@ def test_close_conn(self): self.conn_mock.quit.assert_called_once_with() + def test_describe_directory(self): + with fh.FTPHook() as ftp_hook: + ftp_hook.describe_directory(self.path) + + self.conn_mock.mlsd.assert_called_once_with(self.path) + def test_list_directory(self): with fh.FTPHook() as ftp_hook: ftp_hook.list_directory(self.path) - self.conn_mock.cwd.assert_called_once_with(self.path) - self.conn_mock.nlst.assert_called_once_with() + self.conn_mock.nlst.assert_called_once_with(self.path) def test_create_directory(self): with fh.FTPHook() as ftp_hook: @@ -112,14 +117,14 @@ def test_retrieve_file(self): _buffer = StringIO("buffer") with fh.FTPHook() as ftp_hook: ftp_hook.retrieve_file(self.path, _buffer) - self.conn_mock.retrbinary.assert_called_once_with("RETR path", _buffer.write, 8192) + self.conn_mock.retrbinary.assert_called_once_with("RETR /some/path", _buffer.write, 8192) def test_retrieve_file_with_callback(self): func = mock.Mock() _buffer = StringIO("buffer") with fh.FTPHook() as ftp_hook: ftp_hook.retrieve_file(self.path, _buffer, callback=func) - self.conn_mock.retrbinary.assert_called_once_with("RETR path", func, 8192) + self.conn_mock.retrbinary.assert_called_once_with("RETR /some/path", func, 8192) def test_connection_success(self): with fh.FTPHook() as ftp_hook: