From c1e8ff7b9d42ae1af1ebb7292fc2042f9b51cfd4 Mon Sep 17 00:00:00 2001
From: sanjay pillai <sanjaypillai11@gmail.com>
Date: Sat, 21 Jan 2023 03:15:58 +0530
Subject: [PATCH] FTP operator has logic in __init__

Move the logic from __init__ to executor for FTP operator
---
 airflow/providers/ftp/operators/ftp.py    | 50 ++++++++++++-----------
 airflow/providers/sftp/operators/sftp.py  |  4 +-
 tests/providers/ftp/operators/test_ftp.py |  4 +-
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/airflow/providers/ftp/operators/ftp.py b/airflow/providers/ftp/operators/ftp.py
index 1d9febd23b68e..41b913a64f495 100644
--- a/airflow/providers/ftp/operators/ftp.py
+++ b/airflow/providers/ftp/operators/ftp.py
@@ -84,21 +84,31 @@ def __init__(
         self.ftp_conn_id = ftp_conn_id
         self.operation = operation
         self.create_intermediate_dirs = create_intermediate_dirs
+        self.local_filepath = local_filepath
+        self.remote_filepath = remote_filepath
 
-        if isinstance(local_filepath, str):
-            self.local_filepath = [local_filepath]
+    @cached_property
+    def hook(self) -> FTPHook:
+        """Create and return an FTPHook."""
+        return FTPHook(ftp_conn_id=self.ftp_conn_id)
+
+    def execute(self, context: Any) -> str | list[str] | None:
+        file_msg = None
+
+        if isinstance(self.local_filepath, str):
+            local_filepath_array = [self.local_filepath]
         else:
-            self.local_filepath = local_filepath
+            local_filepath_array = self.local_filepath
 
-        if isinstance(remote_filepath, str):
-            self.remote_filepath = [remote_filepath]
+        if isinstance(self.remote_filepath, str):
+            remote_filepath_array = [self.remote_filepath]
         else:
-            self.remote_filepath = remote_filepath
+            remote_filepath_array = self.remote_filepath
 
-        if len(self.local_filepath) != len(self.remote_filepath):
+        if len(local_filepath_array) != len(remote_filepath_array):
             raise ValueError(
-                f"{len(self.local_filepath)} paths in local_filepath "
-                f"!= {len(self.remote_filepath)} paths in remote_filepath"
+                f"{len(local_filepath_array)} paths in local_filepath "
+                f"!= {len(remote_filepath_array)} paths in remote_filepath"
             )
 
         if self.operation.lower() not in [FTPOperation.GET, FTPOperation.PUT]:
@@ -107,28 +117,22 @@ def __init__(
                 f"expected {FTPOperation.GET} or {FTPOperation.PUT}."
             )
 
-    @cached_property
-    def hook(self) -> FTPHook:
-        """Create and return an FTPHook."""
-        return FTPHook(ftp_conn_id=self.ftp_conn_id)
-
-    def execute(self, context: Any) -> str | list[str] | None:
-        file_msg = None
-        for local_filepath, remote_filepath in zip(self.local_filepath, self.remote_filepath):
+        for _local_filepath, _remote_filepath in zip(local_filepath_array, remote_filepath_array):
             if self.operation.lower() == FTPOperation.GET:
-                local_folder = os.path.dirname(local_filepath)
+                local_folder = os.path.dirname(_local_filepath)
                 if self.create_intermediate_dirs:
                     Path(local_folder).mkdir(parents=True, exist_ok=True)
-                file_msg = f"from {remote_filepath} to {local_filepath}"
+                file_msg = f"from {_remote_filepath} to {_local_filepath}"
                 self.log.info("Starting to transfer %s", file_msg)
-                self.hook.retrieve_file(remote_filepath, local_filepath)
+                self.hook.retrieve_file(_remote_filepath, _local_filepath)
             else:
-                remote_folder = os.path.dirname(remote_filepath)
+                remote_folder = os.path.dirname(_remote_filepath)
                 if self.create_intermediate_dirs:
                     self.hook.create_directory(remote_folder)
-                file_msg = f"from {local_filepath} to {remote_filepath}"
+                file_msg = f"from {_local_filepath} to {_remote_filepath}"
                 self.log.info("Starting to transfer file %s", file_msg)
-                self.hook.store_file(remote_filepath, local_filepath)
+                self.hook.store_file(_remote_filepath, _local_filepath)
+
         return self.local_filepath
 
 
diff --git a/airflow/providers/sftp/operators/sftp.py b/airflow/providers/sftp/operators/sftp.py
index 2bb2cee9f8271..1ccb135d5ae6d 100644
--- a/airflow/providers/sftp/operators/sftp.py
+++ b/airflow/providers/sftp/operators/sftp.py
@@ -104,10 +104,8 @@ def __init__(
         self.remote_filepath = remote_filepath
 
     def execute(self, context: Any) -> str | list[str] | None:
-        local_filepath_was_str = False
         if isinstance(self.local_filepath, str):
             local_filepath_array = [self.local_filepath]
-            local_filepath_was_str = True
         else:
             local_filepath_array = self.local_filepath
 
@@ -189,4 +187,4 @@ def execute(self, context: Any) -> str | list[str] | None:
         except Exception as e:
             raise AirflowException(f"Error while transferring {file_msg}, error: {str(e)}")
 
-        return local_filepath_array[0] if local_filepath_was_str else local_filepath_array
+        return self.local_filepath
diff --git a/tests/providers/ftp/operators/test_ftp.py b/tests/providers/ftp/operators/test_ftp.py
index 3632e0ba7dc24..f62e3f2bd4b13 100644
--- a/tests/providers/ftp/operators/test_ftp.py
+++ b/tests/providers/ftp/operators/test_ftp.py
@@ -167,7 +167,7 @@ def test_unequal_local_remote_file_paths(self):
                 ftp_conn_id=DEFAULT_CONN_ID,
                 local_filepath="/tmp/test",
                 remote_filepath=["/tmp/test1", "/tmp/test2"],
-            )
+            ).execute(None)
 
         with pytest.raises(ValueError):
             FTPFileTransmitOperator(
@@ -175,7 +175,7 @@ def test_unequal_local_remote_file_paths(self):
                 ftp_conn_id=DEFAULT_CONN_ID,
                 local_filepath=["/tmp/test1", "/tmp/test2"],
                 remote_filepath="/tmp/test1",
-            )
+            ).execute(None)
 
 
 class TestFTPSFileTransmitOperator:
