Skip to content

Commit

Permalink
FTPHook methods should not change working directory (#35105)
Browse files Browse the repository at this point in the history
* 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 <hussein@awala.fr>

* 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 <e.makrushin@skbkontur.ru>
Co-authored-by: Hussein Awala <hussein@awala.fr>
  • Loading branch information
3 people committed Oct 30, 2023
1 parent 5f2999e commit ff30dcc
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 18 deletions.
19 changes: 5 additions & 14 deletions airflow/providers/ftp/hooks/ftp.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@

import datetime
import ftplib
import os
from typing import Any, Callable

from airflow.hooks.base import BaseHook
Expand Down Expand Up @@ -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]:
"""
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down
13 changes: 9 additions & 4 deletions tests/providers/ftp/hooks/test_ftp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit ff30dcc

Please sign in to comment.