Implement concurrent directory transfer in SFTPOperator

providers/sftp/src/airflow/providers/sftp/hooks/sftp.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import concurrent.futures
 import datetime
 import os
 import stat
@@ -339,6 +340,63 @@ def retrieve_directory(self, remote_full_path: str, local_full_path: str, prefet
                 new_local_path = os.path.join(local_full_path, os.path.relpath(file_path, remote_full_path))
                 self.retrieve_file(file_path, new_local_path, prefetch)
 
+    def retrieve_directory_concurrently(
+        self, remote_full_path: str, local_full_path: str, workers: int = os.cpu_count() or 2
+    ) -> None:
+        """
+        Transfer the remote directory to a local location concurrently.
+
+        If local_full_path is a string path, the directory will be put
+        at that location.
+
+        :param remote_full_path: full path to the remote directory
+        :param local_full_path: full path to the local directory
+        :param prefetch: controls whether prefetch is performed (default: True)
+        :param workers: number of workers to use for concurrent transfer (default: number of CPUs or 2 if undetermined)
+        """
+
+        def retrieve_file_chunk(
+            conn: SFTPClient, local_file_chunk: list[str], remote_file_chunk: list[str], prefetch: bool = True
+        ):
+            for local_file, remote_file in zip(local_file_chunk, remote_file_chunk):
+                conn.get(remote_file, local_file, prefetch=prefetch)
+
+        with self.get_managed_conn():
+            if Path(local_full_path).exists():
+                raise AirflowException(f"{local_full_path} already exists")
+            Path(local_full_path).mkdir(parents=True)
+            new_local_file_paths, remote_file_paths = [], []
+            files, dirs, _ = self.get_tree_map(remote_full_path)
+            for dir_path in dirs:
+                new_local_path = os.path.join(local_full_path, os.path.relpath(dir_path, remote_full_path))
+                Path(new_local_path).mkdir(parents=True, exist_ok=True)
+            for file in files:
+                remote_file_paths.append(file)
+                new_local_file_paths.append(
+                    os.path.join(local_full_path, os.path.relpath(file, remote_full_path))
+                )
+        remote_file_chunks = [remote_file_paths[i::workers] for i in range(workers)]
+        local_file_chunks = [new_local_file_paths[i::workers] for i in range(workers)]
+        self.log.info("Opening %s new SFTP connections", workers)
+        conns = [SFTPHook(ssh_conn_id=self.ssh_conn_id).get_conn() for _ in range(workers)]
+        try:
+            self.log.info("Retrieving files concurrently with %s threads", workers)
+            with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [
+                    executor.submit(
+                        retrieve_file_chunk,
+                        conns[i],
+                        local_file_chunks[i],
+                        remote_file_chunks[i],
+                    )
+                    for i in range(workers)
+                ]
+                for future in concurrent.futures.as_completed(futures):
+                    future.result()
+        finally:
+            for conn in conns:
+                conn.close()
+
     def store_directory(self, remote_full_path: str, local_full_path: str, confirm: bool = True) -> None:
         """
         Transfer a local directory to the remote location.
@@ -367,6 +425,71 @@ def store_directory(self, remote_full_path: str, local_full_path: str, confirm:
                     )
                     self.store_file(new_remote_path, file_path, confirm)
 
+    def store_directory_concurrently(
+        self,
+        remote_full_path: str,
+        local_full_path: str,
+        confirm: bool = True,
+        workers: int = os.cpu_count() or 2,
+    ) -> None:
+        """
+        Transfer a local directory to the remote location concurrently.
+
+        If local_full_path is a string path, the directory will be read
+        from that location.
+
+        :param remote_full_path: full path to the remote directory
+        :param local_full_path: full path to the local directory
+        :param confirm: whether to confirm the file size after transfer (default: True)
+        :param workers: number of workers to use for concurrent transfer (default: number of CPUs or 2 if undetermined)
+        """
+
+        def store_file_chunk(
+            conn: SFTPClient, local_file_chunk: list[str], remote_file_chunk: list[str], confirm: bool
+        ):
+            for local_file, remote_file in zip(local_file_chunk, remote_file_chunk):
+                conn.put(local_file, remote_file, confirm=confirm)
+
+        with self.get_managed_conn():
+            if self.path_exists(remote_full_path):
+                raise AirflowException(f"{remote_full_path} already exists")
+            self.create_directory(remote_full_path)
+
+            local_file_paths, new_remote_file_paths = [], []
+            for root, dirs, files in os.walk(local_full_path):
+                for dir_name in dirs:
+                    dir_path = os.path.join(root, dir_name)
+                    new_remote_path = os.path.join(
+                        remote_full_path, os.path.relpath(dir_path, local_full_path)
+                    )
+                    self.create_directory(new_remote_path)
+                for file_name in files:
+                    file_path = os.path.join(root, file_name)
+                    new_remote_path = os.path.join(
+                        remote_full_path, os.path.relpath(file_path, local_full_path)
+                    )
+                    local_file_paths.append(file_path)
+                    new_remote_file_paths.append(new_remote_path)
+
+        remote_file_chunks = [new_remote_file_paths[i::workers] for i in range(workers)]
+        local_file_chunks = [local_file_paths[i::workers] for i in range(workers)]
+        self.log.info("Opening %s new SFTP connections", workers)
+        conns = [SFTPHook(ssh_conn_id=self.ssh_conn_id).get_conn() for _ in range(workers)]
+        try:
+            self.log.info("Storing files concurrently with %s threads", workers)
+            with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [
+                    executor.submit(
+                        store_file_chunk, conns[i], local_file_chunks[i], remote_file_chunks[i], confirm
+                    )
+                    for i in range(workers)
+                ]
+                for future in concurrent.futures.as_completed(futures):
+                    future.result()
+        finally:
+            for conn in conns:
+                conn.close()
+
     def get_mod_time(self, path: str) -> str:
         """
         Get an entry's modification time.

providers/sftp/src/airflow/providers/sftp/operators/sftp.py

@@ -74,6 +74,8 @@ class SFTPOperator(BaseOperator):
                 create_intermediate_dirs=True,
                 dag=dag,
             )
+    :param concurrency: Number of threads when transferring directories. Each thread opens a new SFTP connection.
+        This parameter is used only when transferring directories, not individual files. (Default is 1)
 
     """
 
@@ -90,6 +92,7 @@ def __init__(
         operation: str = SFTPOperation.PUT,
         confirm: bool = True,
         create_intermediate_dirs: bool = False,
+        concurrency: int = 1,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -101,6 +104,7 @@ def __init__(
         self.create_intermediate_dirs = create_intermediate_dirs
         self.local_filepath = local_filepath
         self.remote_filepath = remote_filepath
+        self.concurrency = concurrency
 
     def execute(self, context: Any) -> str | list[str] | None:
         if self.local_filepath is None:
@@ -132,6 +136,9 @@ def execute(self, context: Any) -> str | list[str] | None:
                 f"expected {SFTPOperation.GET} or {SFTPOperation.PUT} or {SFTPOperation.DELETE}."
             )
 
+        if self.concurrency < 1:
+            raise ValueError(f"concurrency should be greater than 0, got {self.concurrency}")
+
         file_msg = None
         try:
             if self.ssh_conn_id:
@@ -161,7 +168,14 @@ def execute(self, context: Any) -> str | list[str] | None:
                         file_msg = f"from {_remote_filepath} to {_local_filepath}"
                         self.log.info("Starting to transfer %s", file_msg)
                         if self.sftp_hook.isdir(_remote_filepath):
-                            self.sftp_hook.retrieve_directory(_remote_filepath, _local_filepath)
+                            if self.concurrency > 1:
+                                self.sftp_hook.retrieve_directory_concurrently(
+                                    _remote_filepath,
+                                    _local_filepath,
+                                    workers=self.concurrency,
+                                )
+                            elif self.concurrency == 1:
+                                self.sftp_hook.retrieve_directory(_remote_filepath, _local_filepath)
                         else:
                             self.sftp_hook.retrieve_file(_remote_filepath, _local_filepath)
                     elif self.operation.lower() == SFTPOperation.PUT:
@@ -171,9 +185,17 @@ def execute(self, context: Any) -> str | list[str] | None:
                         file_msg = f"from {_local_filepath} to {_remote_filepath}"
                         self.log.info("Starting to transfer file %s", file_msg)
                         if os.path.isdir(_local_filepath):
-                            self.sftp_hook.store_directory(
-                                _remote_filepath, _local_filepath, confirm=self.confirm
-                            )
+                            if self.concurrency > 1:
+                                self.sftp_hook.store_directory_concurrently(
+                                    _remote_filepath,
+                                    _local_filepath,
+                                    confirm=self.confirm,
+                                    workers=self.concurrency,
+                                )
+                            elif self.concurrency == 1:
+                                self.sftp_hook.store_directory(
+                                    _remote_filepath, _local_filepath, confirm=self.confirm
+                                )
                         else:
                             self.sftp_hook.store_file(_remote_filepath, _local_filepath, confirm=self.confirm)
             elif self.operation.lower() == SFTPOperation.DELETE:

providers/sftp/tests/unit/sftp/hooks/test_sftp.py

@@ -530,6 +530,23 @@ def test_store_and_retrieve_directory(self):
         )
         assert retrieved_dir_name in os.listdir(os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS))
 
+    def test_store_and_retrieve_directory_concurrently(self):
+        stored_dir_name = "stored_dir"
+        self.hook.store_directory_concurrently(
+            remote_full_path=os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS, stored_dir_name),
+            local_full_path=os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS, SUB_DIR),
+        )
+        output = self.hook.list_directory(
+            path=os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS, stored_dir_name)
+        )
+        assert output == [TMP_FILE_FOR_TESTS]
+        retrieved_dir_name = "retrieved_dir"
+        self.hook.retrieve_directory_concurrently(
+            remote_full_path=os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS, stored_dir_name),
+            local_full_path=os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS, retrieved_dir_name),
+        )
+        assert retrieved_dir_name in os.listdir(os.path.join(self.temp_dir, TMP_DIR_FOR_TESTS))
+
     @patch("paramiko.SSHClient")
     @patch("paramiko.ProxyCommand")
     def test_sftp_hook_with_proxy_command(self, mock_proxy_command, mock_ssh_client):

providers/sftp/tests/unit/sftp/operators/test_sftp.py

@@ -429,6 +429,22 @@ def test_str_dirpaths_get(self, mock_get):
         args, _ = mock_get.call_args_list[0]
         assert args == (remote_dirpath, local_dirpath)
 
+    @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.retrieve_directory_concurrently")
+    def test_str_dirpaths_get_concurrently(self, mock_get):
+        local_dirpath = "/tmp_local"
+        remote_dirpath = "/tmp"
+        SFTPOperator(
+            task_id="test_str_to_list",
+            sftp_hook=self.sftp_hook,
+            local_filepath=local_dirpath,
+            remote_filepath=remote_dirpath,
+            operation=SFTPOperation.GET,
+            concurrency=2,
+        ).execute(None)
+        assert mock_get.call_count == 1
+        args, _ = mock_get.call_args_list[0]
+        assert args == (remote_dirpath, local_dirpath)
+
     @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.store_file")
     def test_str_filepaths_put(self, mock_get):
         local_filepath = "/tmp/test"
@@ -477,6 +493,22 @@ def test_str_dirpaths_put(self, mock_get):
         args, _ = mock_get.call_args_list[0]
         assert args == (remote_dirpath, local_dirpath)
 
+    @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.store_directory_concurrently")
+    def test_str_dirpaths_put_concurrently(self, mock_get):
+        local_dirpath = "/tmp"
+        remote_dirpath = "/tmp_remote"
+        SFTPOperator(
+            task_id="test_str_dirpaths_put",
+            sftp_hook=self.sftp_hook,
+            local_filepath=local_dirpath,
+            remote_filepath=remote_dirpath,
+            operation=SFTPOperation.PUT,
+            concurrency=2,
+        ).execute(None)
+        assert mock_get.call_count == 1
+        args, _ = mock_get.call_args_list[0]
+        assert args == (remote_dirpath, local_dirpath)
+
     @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.retrieve_file")
     def test_return_str_when_local_filepath_was_str(self, mock_get):
         local_filepath = "/tmp/ltest1"