dabla commented on code in PR #68298: URL: https://github.com/apache/airflow/pull/68298#discussion_r3518097117
########## providers/sftp/tests/unit/sftp/test_constants.py: ########## @@ -0,0 +1,24 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.providers.sftp.constants import SFTPOperation + + +def test_sftp_operation_values(): Review Comment: Same here, just move the test to the `test_sftp` module. ########## providers/sftp/src/airflow/providers/sftp/hooks/sftp.py: ########## @@ -1069,3 +1130,67 @@ async def get_mod_time(self, path: str) -> str: # type: ignore[return] return mod_time except asyncssh.SFTPNoSuchFile: raise AirflowException("No files matching") + + async def transfer( + self, + operation: str, + local_filepath: str | list[str] | None, + remote_filepath: str | list[str], + confirm: bool = True, + create_intermediate_dirs: bool = False, + concurrency: int = 1, + prefetch: bool = True, + ) -> None: + """Perform an SFTP transfer operation (GET, PUT, or DELETE) using native async I/O.""" + from airflow.providers.sftp.constants import SFTPOperation + + if isinstance(local_filepath, str): + local_filepath_array = [local_filepath] if local_filepath else [] + else: + local_filepath_array = local_filepath or [] + + if isinstance(remote_filepath, str): + remote_filepath_array = [remote_filepath] + else: + remote_filepath_array = list(remote_filepath) + + semaphore = asyncio.Semaphore(concurrency) + + async def _bounded(coro): + async with semaphore: + return await coro + + async with await self._get_conn() as ssh_conn: + async with ssh_conn.start_sftp_client() as sftp: + if operation.lower() == SFTPOperation.GET: + + async def _get(local: str, remote: str): + if create_intermediate_dirs: + os.makedirs(os.path.dirname(local), exist_ok=True) + await self.retrieve_file(remote, local) + + tasks = [ + asyncio.create_task(_bounded(_get(local, remote))) + for local, remote in zip(local_filepath_array, remote_filepath_array) + ] + await asyncio.gather(*tasks) + elif operation.lower() == SFTPOperation.PUT: + + async def _put(local: str, remote: str): + await self.store_file(remote, local) + + tasks = [ + asyncio.create_task(_bounded(_put(local, remote))) + for local, remote in zip(local_filepath_array, remote_filepath_array) + ] + await asyncio.gather(*tasks) + elif operation.lower() == SFTPOperation.DELETE: + + async def _delete(remote: str): + await sftp.unlink(remote) + Review Comment: Maybe here we could also make this method more DRY, just like it's sync counterpart, even though here repetition is smaller compared to the synced version. Yet the `_get `, `_put `and `_delete ` could be called from within the same list comprehension and just past as argument. ########## providers/sftp/src/airflow/providers/sftp/hooks/sftp.py: ########## @@ -735,6 +727,75 @@ def get_files_by_pattern(self, path, fnmatch_pattern) -> list[str]: return matched_files + def transfer( + self, + operation: str, + local_filepath: str | list[str] | None, + remote_filepath: str | list[str], + confirm: bool = True, + create_intermediate_dirs: bool = False, + concurrency: int = 1, + prefetch: bool = True, + ) -> None: + """ + Perform a synchronous SFTP transfer operation (GET, PUT, or DELETE). + + Centralises transfer logic so both the operator and the trigger + can delegate to the hook, in line with the DRY principle. + + :param operation: The SFTP operation - put, get, or delete. + :param local_filepath: Local file path(s). + :param remote_filepath: Remote file path(s). + :param confirm: Whether to confirm file size after PUT (default: True). + :param create_intermediate_dirs: Create missing intermediate directories (default: False). + :param concurrency: Number of threads for directory transfers (default: 1). + :param prefetch: Whether to prefetch during GET (default: True). + """ + from airflow.providers.sftp.constants import SFTPOperation + + if isinstance(local_filepath, str): + local_filepath_array = [local_filepath] if local_filepath else [] + else: + local_filepath_array = local_filepath or [] + + if isinstance(remote_filepath, str): + remote_filepath_array = [remote_filepath] + else: + remote_filepath_array = list(remote_filepath) + + if operation.lower() == SFTPOperation.GET: + for local, remote in zip(local_filepath_array, remote_filepath_array): + if create_intermediate_dirs: + Path(os.path.dirname(local)).mkdir(parents=True, exist_ok=True) + if self.isdir(remote): + if concurrency > 1: + self.retrieve_directory_concurrently( + remote, local, workers=concurrency, prefetch=prefetch + ) + else: + self.retrieve_directory(remote, local) + else: + self.retrieve_file(remote, local, prefetch=prefetch) + elif operation.lower() == SFTPOperation.PUT: + for local, remote in zip(local_filepath_array, remote_filepath_array): + if create_intermediate_dirs: + self.create_directory(os.path.dirname(remote)) + if os.path.isdir(local): + if concurrency > 1: + self.store_directory_concurrently( + remote, local, confirm=confirm, workers=concurrency + ) + else: + self.store_directory(remote, local, confirm=confirm) + else: + self.store_file(remote, local, confirm=confirm) + elif operation.lower() == SFTPOperation.DELETE: + for remote in remote_filepath_array: + if self.isdir(remote): + self.delete_directory(remote, include_files=True) + else: + self.delete_file(remote) Review Comment: I know the original SFTPOperator was implemented this way, but this seems like a good opportunity to improve it as well. There is quite a bit of duplication between the GET, PUT, and DELETE branches. All three follow the same pattern: iterate over the paths and perform an operation-specific action. Instead of branching around the entire loop, could we select the appropriate callable (or handler function) based on the requested operation first, and then have a single loop that invokes it? That would make the code more DRY and make it easier to add new operations in the future. ########## providers/sftp/src/airflow/providers/sftp/constants.py: ########## @@ -0,0 +1,28 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""SFTP constants.""" + +from __future__ import annotations + + +class SFTPOperation: + """SFTP operation constants.""" + + GET = "get" + PUT = "put" + DELETE = "delete" Review Comment: I would just move the `SFTPOperation` class back to the hook where it originally was, no need to create dedicated `constants` module for that. ########## providers/sftp/src/airflow/providers/sftp/hooks/sftp.py: ########## @@ -735,6 +727,75 @@ def get_files_by_pattern(self, path, fnmatch_pattern) -> list[str]: return matched_files + def transfer( + self, + operation: str, + local_filepath: str | list[str] | None, + remote_filepath: str | list[str], + confirm: bool = True, + create_intermediate_dirs: bool = False, + concurrency: int = 1, + prefetch: bool = True, + ) -> None: + """ + Perform a synchronous SFTP transfer operation (GET, PUT, or DELETE). + + Centralises transfer logic so both the operator and the trigger + can delegate to the hook, in line with the DRY principle. + + :param operation: The SFTP operation - put, get, or delete. + :param local_filepath: Local file path(s). + :param remote_filepath: Remote file path(s). + :param confirm: Whether to confirm file size after PUT (default: True). + :param create_intermediate_dirs: Create missing intermediate directories (default: False). + :param concurrency: Number of threads for directory transfers (default: 1). + :param prefetch: Whether to prefetch during GET (default: True). + """ + from airflow.providers.sftp.constants import SFTPOperation + + if isinstance(local_filepath, str): + local_filepath_array = [local_filepath] if local_filepath else [] + else: + local_filepath_array = local_filepath or [] + + if isinstance(remote_filepath, str): + remote_filepath_array = [remote_filepath] + else: + remote_filepath_array = list(remote_filepath) + + if operation.lower() == SFTPOperation.GET: + for local, remote in zip(local_filepath_array, remote_filepath_array): + if create_intermediate_dirs: + Path(os.path.dirname(local)).mkdir(parents=True, exist_ok=True) + if self.isdir(remote): + if concurrency > 1: + self.retrieve_directory_concurrently( + remote, local, workers=concurrency, prefetch=prefetch + ) + else: + self.retrieve_directory(remote, local) + else: + self.retrieve_file(remote, local, prefetch=prefetch) + elif operation.lower() == SFTPOperation.PUT: + for local, remote in zip(local_filepath_array, remote_filepath_array): + if create_intermediate_dirs: + self.create_directory(os.path.dirname(remote)) + if os.path.isdir(local): + if concurrency > 1: + self.store_directory_concurrently( + remote, local, confirm=confirm, workers=concurrency + ) + else: + self.store_directory(remote, local, confirm=confirm) + else: + self.store_file(remote, local, confirm=confirm) + elif operation.lower() == SFTPOperation.DELETE: + for remote in remote_filepath_array: + if self.isdir(remote): + self.delete_directory(remote, include_files=True) + else: + self.delete_file(remote) Review Comment: Something like this: ``` from pathlib import Path import os def transfer( self, operation: str, local_filepath: str | list[str] | None, remote_filepath: str | list[str], confirm: bool = True, create_intermediate_dirs: bool = False, concurrency: int = 1, prefetch: bool = True, ) -> None: """Perform a synchronous SFTP transfer operation.""" from airflow.providers.sftp.constants import SFTPOperation if isinstance(local_filepath, str): local_paths = [local_filepath] else: local_paths = local_filepath or [] if isinstance(remote_filepath, str): remote_paths = [remote_filepath] else: remote_paths = list(remote_filepath) operations = { SFTPOperation.GET: lambda local, remote: self._transfer_get( local, remote, create_intermediate_dirs=create_intermediate_dirs, concurrency=concurrency, prefetch=prefetch, ), SFTPOperation.PUT: lambda local, remote: self._transfer_put( local, remote, confirm=confirm, create_intermediate_dirs=create_intermediate_dirs, concurrency=concurrency, ), SFTPOperation.DELETE: lambda _, remote: self._transfer_delete(remote), } try: handler = operations[operation.lower()] except KeyError: raise ValueError(f"Unsupported SFTP operation: {operation}") if operation.lower() == SFTPOperation.DELETE: for remote in remote_paths: handler(None, remote) else: for local, remote in zip(local_paths, remote_paths): handler(local, remote) def _transfer_get( self, local: str, remote: str, *, create_intermediate_dirs: bool, concurrency: int, prefetch: bool, ) -> None: if create_intermediate_dirs: Path(local).parent.mkdir(parents=True, exist_ok=True) if self.isdir(remote): if concurrency > 1: self.retrieve_directory_concurrently( remote, local, workers=concurrency, prefetch=prefetch, ) else: self.retrieve_directory(remote, local) else: self.retrieve_file(remote, local, prefetch=prefetch) def _transfer_put( self, local: str, remote: str, *, confirm: bool, create_intermediate_dirs: bool, concurrency: int, ) -> None: if create_intermediate_dirs: self.create_directory(os.path.dirname(remote)) if os.path.isdir(local): if concurrency > 1: self.store_directory_concurrently( remote, local, confirm=confirm, workers=concurrency, ) else: self.store_directory(remote, local, confirm=confirm) else: self.store_file(remote, local, confirm=confirm) def _transfer_delete(self, remote: str) -> None: if self.isdir(remote): self.delete_directory(remote, include_files=True) else: self.delete_file(remote) ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
