dabla commented on code in PR #68298:
URL: https://github.com/apache/airflow/pull/68298#discussion_r3459746775
##########
providers/sftp/src/airflow/providers/sftp/triggers/sftp.py:
##########
@@ -140,3 +140,84 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
def _get_async_hook(self) -> SFTPHookAsync:
return SFTPHookAsync(sftp_conn_id=self.sftp_conn_id)
+
+
+class SFTPOperatorTrigger(BaseTrigger):
Review Comment:
I would call this SFTPOperationTrigger instead of SFTPOperatorTrigger
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1069,3 +1109,40 @@ 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)
+
+ if operation.lower() == SFTPOperation.GET:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
Review Comment:
Looking at the implementation, the biggest inefficiency beside the for loop
is that every single retrieve_file() and store_file() call creates a brand new
SSH connection and SFTP session.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1069,3 +1109,40 @@ 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)
+
+ if operation.lower() == SFTPOperation.GET:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
+ if create_intermediate_dirs:
+ os.makedirs(os.path.dirname(local), exist_ok=True)
+ await self.retrieve_file(remote, local)
+ elif operation.lower() == SFTPOperation.PUT:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
Review Comment:
Same remark here
##########
providers/sftp/src/airflow/providers/sftp/operators/sftp.py:
##########
@@ -108,8 +99,25 @@ def __init__(
self.remote_filepath = remote_filepath
self.concurrency = concurrency
self.prefetch = prefetch
+ self.deferrable = deferrable
def execute(self, context: Any) -> str | list[str] | None:
+ if self.deferrable:
+ self.defer(
+ trigger=SFTPTrigger(
+ ssh_conn_id=self.ssh_conn_id,
+ local_filepath=self.local_filepath,
+ remote_filepath=self.remote_filepath,
+ operation=self.operation,
+ confirm=self.confirm,
+ create_intermediate_dirs=self.create_intermediate_dirs,
+ remote_host=self.remote_host,
+ concurrency=self.concurrency,
+ prefetch=self.prefetch,
+ ),
+ method_name="execute_complete",
Review Comment:
Nit, replace with:
` method_name=self.execute_complete.__name__,`
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1069,3 +1109,40 @@ 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)
+
+ if operation.lower() == SFTPOperation.GET:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
+ if create_intermediate_dirs:
+ os.makedirs(os.path.dirname(local), exist_ok=True)
+ await self.retrieve_file(remote, local)
+ elif operation.lower() == SFTPOperation.PUT:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
+ await self.store_file(remote, local)
+ elif operation.lower() == SFTPOperation.DELETE:
+ for remote in remote_filepath_array:
Review Comment:
And here
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1069,3 +1109,40 @@ 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)
+
+ if operation.lower() == SFTPOperation.GET:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
+ if create_intermediate_dirs:
+ os.makedirs(os.path.dirname(local), exist_ok=True)
+ await self.retrieve_file(remote, local)
+ elif operation.lower() == SFTPOperation.PUT:
+ for local, remote in zip(local_filepath_array,
remote_filepath_array):
+ await self.store_file(remote, local)
+ elif operation.lower() == SFTPOperation.DELETE:
+ for remote in remote_filepath_array:
Review Comment:
The method could be rewritten in something like this which would be far more
efficient and take advantage of the async:
```
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."""
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)
tasks = [
asyncio.create_task(_bounded(_delete(remote)))
for remote in remote_filepath_array
]
await asyncio.gather(*tasks)
```
--
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]
