potiuk commented on code in PR #72336:
URL: https://github.com/apache/airflow/pull/72336#discussion_r4027315437
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1080,3 +1170,283 @@ 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 sense_files_by_pattern(
+ self,
+ path: str,
+ fnmatch_pattern: str,
+ newer_than: datetime.datetime | None = None,
+ ) -> list[str]:
+ """
+ Return the names of files at ``path`` matching ``fnmatch_pattern``.
+
+ If ``newer_than`` is provided, only files modified after that
timestamp are returned; files
+ without a reported modification time are skipped in that case.
+
+ :param path: directory on the SFTP server to search for files matching
the pattern
+ :param fnmatch_pattern: pattern used to match filenames, see the
``fnmatch`` std library module
+ :param newer_than: if provided, only files modified after this UTC
timestamp are returned
+ """
+ files = await self.get_files_and_attrs_by_pattern(path=path,
fnmatch_pattern=fnmatch_pattern)
+ if not newer_than:
+ return [str(file.filename) for file in files]
+
+ matched_files = []
+ for file in files:
+ if file.attrs.mtime is None:
+ continue
+ if newer_than <= self._mod_time_to_utc(file.attrs.mtime):
+ matched_files.append(str(file.filename))
+ return matched_files
+
+ async def sense_path(self, path: str, newer_than: datetime.datetime | None
= None) -> bool:
+ """
+ Return whether ``path`` exists and, if ``newer_than`` is provided, was
modified since.
+
+ :param path: full path to the remote file
+ :param newer_than: if provided, the file must have been modified after
this UTC timestamp
+ """
+ mod_time = await self.get_mod_time(path)
+ if not newer_than:
+ return True
+ return newer_than <= self._mod_time_to_utc(mod_time)
+
+ @staticmethod
+ def _mod_time_to_utc(mod_time: int | float | str) -> datetime.datetime:
+ """Convert a modification time, either an epoch timestamp or
``%Y%m%d%H%M%S`` string, to UTC."""
+ if not isinstance(mod_time, str):
+ mod_time =
datetime.datetime.fromtimestamp(float(mod_time)).strftime("%Y%m%d%H%M%S")
+ return timezone.convert_to_utc(datetime.datetime.strptime(mod_time,
"%Y%m%d%H%M%S"))
+
+ async def isdir(self, path: str) -> bool:
Review Comment:
mypy reports `Missing return statement [return]` here and on `path_exists`
at line 1235.
It is not a spurious warning. `__aexit__` is typed as potentially
suppressing the exception, so as far as the type checker is concerned control
can fall out of the `async with` without hitting either `return` — and if a
context manager ever did suppress, the function would return `None` from a `->
bool` signature.
Assigning the result inside the block and returning once at the end (or an
explicit `return False` after it) resolves it without a `# type: ignore`.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1080,3 +1170,283 @@ 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 sense_files_by_pattern(
+ self,
+ path: str,
+ fnmatch_pattern: str,
+ newer_than: datetime.datetime | None = None,
+ ) -> list[str]:
+ """
+ Return the names of files at ``path`` matching ``fnmatch_pattern``.
+
+ If ``newer_than`` is provided, only files modified after that
timestamp are returned; files
+ without a reported modification time are skipped in that case.
+
+ :param path: directory on the SFTP server to search for files matching
the pattern
+ :param fnmatch_pattern: pattern used to match filenames, see the
``fnmatch`` std library module
+ :param newer_than: if provided, only files modified after this UTC
timestamp are returned
+ """
+ files = await self.get_files_and_attrs_by_pattern(path=path,
fnmatch_pattern=fnmatch_pattern)
+ if not newer_than:
+ return [str(file.filename) for file in files]
+
+ matched_files = []
+ for file in files:
+ if file.attrs.mtime is None:
+ continue
+ if newer_than <= self._mod_time_to_utc(file.attrs.mtime):
+ matched_files.append(str(file.filename))
+ return matched_files
+
+ async def sense_path(self, path: str, newer_than: datetime.datetime | None
= None) -> bool:
+ """
+ Return whether ``path`` exists and, if ``newer_than`` is provided, was
modified since.
+
+ :param path: full path to the remote file
+ :param newer_than: if provided, the file must have been modified after
this UTC timestamp
+ """
+ mod_time = await self.get_mod_time(path)
+ if not newer_than:
+ return True
+ return newer_than <= self._mod_time_to_utc(mod_time)
+
+ @staticmethod
+ def _mod_time_to_utc(mod_time: int | float | str) -> datetime.datetime:
+ """Convert a modification time, either an epoch timestamp or
``%Y%m%d%H%M%S`` string, to UTC."""
+ if not isinstance(mod_time, str):
+ mod_time =
datetime.datetime.fromtimestamp(float(mod_time)).strftime("%Y%m%d%H%M%S")
+ return timezone.convert_to_utc(datetime.datetime.strptime(mod_time,
"%Y%m%d%H%M%S"))
+
+ async def isdir(self, path: str) -> bool:
+ """
+ Check if the path provided is a directory.
+
+ :param path: full path to the remote directory to check
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ attrs = await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return attrs.permissions is not None and
stat.S_ISDIR(attrs.permissions)
+
+ async def path_exists(self, path: str) -> bool:
+ """
+ Whether a remote entity exists.
+
+ :param path: full path to the remote file or directory
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return True
+
+ async def create_directory(self, path: str) -> None:
+ """
+ Create a directory (and any missing parents) on the remote system
asynchronously.
+
+ Returns silently if the target directory already exists, mirroring
+ :meth:`SFTPHook.create_directory`.
+
+ :param path: full path to the remote directory to create
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.makedirs(path, exist_ok=True)
+
+ async def delete_file(self, path: str) -> None:
+ """
+ Remove a file on the server asynchronously.
+
+ :param path: full path to the remote file
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.unlink(path)
+
+ async def delete_directory(self, path: str, include_files: bool = False)
-> None:
+ """
+ Delete a directory on the remote system asynchronously.
+
+ :param path: full path to the remote directory to delete
+ :param include_files: whether to recursively delete the directory's
contents first
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+
+ if include_files:
+ files, dirs, _ = await self.get_tree_map(path)
+ dirs = dirs[::-1] # reverse the order for deleting deepest
directories first
+
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ for file_path in files:
+ await sftp.remove(file_path)
+ for dir_path in dirs:
+ await sftp.rmdir(dir_path)
+ await sftp.rmdir(path)
+
+ async def get_tree_map(
+ self, path: str, prefix: str | None = None, delimiter: str | None =
None
+ ) -> tuple[list[str], list[str], list[str]]:
+ """
+ Get tuple with recursive lists of files, directories and unknown paths
asynchronously.
+
+ It is possible to filter results by giving prefix and/or delimiter
parameters.
+
+ :param path: path from which tree will be built
+ :param prefix: if set paths will be added if start with prefix
+ :param delimiter: if set paths will be added if end with delimiter
+ :return: tuple with list of files, dirs and unknown items
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+ unknowns: list[str] = []
+
+ def append_matching_path_callback(list_: list[str]) -> Callable:
+ return lambda item: (
+ list_.append(item) if SFTPHook._is_path_match(item, prefix,
delimiter) else None
+ )
+
+ await self.walktree(
+ path=path,
+ fcallback=append_matching_path_callback(files),
+ dcallback=append_matching_path_callback(dirs),
+ ucallback=append_matching_path_callback(unknowns),
+ recurse=True,
+ )
+
+ return files, dirs, unknowns
+
+ async def retrieve_directory(
+ self, remote_full_path: str, local_full_path: str, prefetch: bool =
True
+ ) -> None:
+ """
+ Transfer the remote directory to a local location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param prefetch: controls whether read-ahead requests are pipelined
(default: True)
+ """
+ if await asyncio.to_thread(Path(local_full_path).exists):
+ raise AirflowException(f"{local_full_path} already exists")
+ dest = await asyncio.to_thread(Path(local_full_path).resolve)
+ await asyncio.to_thread(dest.mkdir, parents=True)
+ files, dirs, _ = await self.get_tree_map(remote_full_path)
+ for dir_path in dirs:
+ relative_path = await asyncio.to_thread(os.path.relpath, dir_path,
remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await asyncio.to_thread(Path(new_local_path).mkdir, parents=True,
exist_ok=True)
+ for file_path in files:
+ relative_path = await asyncio.to_thread(os.path.relpath,
file_path, remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await self.retrieve_file(file_path, new_local_path,
prefetch=prefetch)
+
+ async def store_directory(
+ self, remote_full_path: str, local_full_path: str, confirm: bool = True
+ ) -> None:
+ """
+ Transfer a local directory to the remote location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param confirm: whether to verify each uploaded file's size (default:
True)
+ """
+ if await self.path_exists(remote_full_path):
+ raise AirflowException(f"{remote_full_path} already exists")
Review Comment:
Second `check-no-new-airflow-exceptions` hit — same reasoning as the one in
`retrieve_directory` above. `SFTPOperationError` or `FileExistsError`.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1080,3 +1170,283 @@ 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 sense_files_by_pattern(
+ self,
+ path: str,
+ fnmatch_pattern: str,
+ newer_than: datetime.datetime | None = None,
+ ) -> list[str]:
+ """
+ Return the names of files at ``path`` matching ``fnmatch_pattern``.
+
+ If ``newer_than`` is provided, only files modified after that
timestamp are returned; files
+ without a reported modification time are skipped in that case.
+
+ :param path: directory on the SFTP server to search for files matching
the pattern
+ :param fnmatch_pattern: pattern used to match filenames, see the
``fnmatch`` std library module
+ :param newer_than: if provided, only files modified after this UTC
timestamp are returned
+ """
+ files = await self.get_files_and_attrs_by_pattern(path=path,
fnmatch_pattern=fnmatch_pattern)
+ if not newer_than:
+ return [str(file.filename) for file in files]
+
+ matched_files = []
+ for file in files:
+ if file.attrs.mtime is None:
+ continue
+ if newer_than <= self._mod_time_to_utc(file.attrs.mtime):
+ matched_files.append(str(file.filename))
+ return matched_files
+
+ async def sense_path(self, path: str, newer_than: datetime.datetime | None
= None) -> bool:
+ """
+ Return whether ``path`` exists and, if ``newer_than`` is provided, was
modified since.
+
+ :param path: full path to the remote file
+ :param newer_than: if provided, the file must have been modified after
this UTC timestamp
+ """
+ mod_time = await self.get_mod_time(path)
+ if not newer_than:
+ return True
+ return newer_than <= self._mod_time_to_utc(mod_time)
+
+ @staticmethod
+ def _mod_time_to_utc(mod_time: int | float | str) -> datetime.datetime:
+ """Convert a modification time, either an epoch timestamp or
``%Y%m%d%H%M%S`` string, to UTC."""
+ if not isinstance(mod_time, str):
+ mod_time =
datetime.datetime.fromtimestamp(float(mod_time)).strftime("%Y%m%d%H%M%S")
+ return timezone.convert_to_utc(datetime.datetime.strptime(mod_time,
"%Y%m%d%H%M%S"))
+
+ async def isdir(self, path: str) -> bool:
+ """
+ Check if the path provided is a directory.
+
+ :param path: full path to the remote directory to check
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ attrs = await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return attrs.permissions is not None and
stat.S_ISDIR(attrs.permissions)
+
+ async def path_exists(self, path: str) -> bool:
+ """
+ Whether a remote entity exists.
+
+ :param path: full path to the remote file or directory
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return True
+
+ async def create_directory(self, path: str) -> None:
+ """
+ Create a directory (and any missing parents) on the remote system
asynchronously.
+
+ Returns silently if the target directory already exists, mirroring
+ :meth:`SFTPHook.create_directory`.
+
+ :param path: full path to the remote directory to create
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.makedirs(path, exist_ok=True)
+
+ async def delete_file(self, path: str) -> None:
+ """
+ Remove a file on the server asynchronously.
+
+ :param path: full path to the remote file
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.unlink(path)
+
+ async def delete_directory(self, path: str, include_files: bool = False)
-> None:
+ """
+ Delete a directory on the remote system asynchronously.
+
+ :param path: full path to the remote directory to delete
+ :param include_files: whether to recursively delete the directory's
contents first
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+
+ if include_files:
+ files, dirs, _ = await self.get_tree_map(path)
+ dirs = dirs[::-1] # reverse the order for deleting deepest
directories first
+
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ for file_path in files:
+ await sftp.remove(file_path)
+ for dir_path in dirs:
+ await sftp.rmdir(dir_path)
+ await sftp.rmdir(path)
+
+ async def get_tree_map(
+ self, path: str, prefix: str | None = None, delimiter: str | None =
None
+ ) -> tuple[list[str], list[str], list[str]]:
+ """
+ Get tuple with recursive lists of files, directories and unknown paths
asynchronously.
+
+ It is possible to filter results by giving prefix and/or delimiter
parameters.
+
+ :param path: path from which tree will be built
+ :param prefix: if set paths will be added if start with prefix
+ :param delimiter: if set paths will be added if end with delimiter
+ :return: tuple with list of files, dirs and unknown items
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+ unknowns: list[str] = []
+
+ def append_matching_path_callback(list_: list[str]) -> Callable:
+ return lambda item: (
+ list_.append(item) if SFTPHook._is_path_match(item, prefix,
delimiter) else None
+ )
+
+ await self.walktree(
+ path=path,
+ fcallback=append_matching_path_callback(files),
+ dcallback=append_matching_path_callback(dirs),
+ ucallback=append_matching_path_callback(unknowns),
+ recurse=True,
+ )
+
+ return files, dirs, unknowns
+
+ async def retrieve_directory(
+ self, remote_full_path: str, local_full_path: str, prefetch: bool =
True
+ ) -> None:
+ """
+ Transfer the remote directory to a local location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param prefetch: controls whether read-ahead requests are pipelined
(default: True)
+ """
+ if await asyncio.to_thread(Path(local_full_path).exists):
+ raise AirflowException(f"{local_full_path} already exists")
+ dest = await asyncio.to_thread(Path(local_full_path).resolve)
+ await asyncio.to_thread(dest.mkdir, parents=True)
+ files, dirs, _ = await self.get_tree_map(remote_full_path)
+ for dir_path in dirs:
+ relative_path = await asyncio.to_thread(os.path.relpath, dir_path,
remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await asyncio.to_thread(Path(new_local_path).mkdir, parents=True,
exist_ok=True)
+ for file_path in files:
+ relative_path = await asyncio.to_thread(os.path.relpath,
file_path, remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await self.retrieve_file(file_path, new_local_path,
prefetch=prefetch)
+
+ async def store_directory(
+ self, remote_full_path: str, local_full_path: str, confirm: bool = True
+ ) -> None:
+ """
+ Transfer a local directory to the remote location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param confirm: whether to verify each uploaded file's size (default:
True)
+ """
+ if await self.path_exists(remote_full_path):
+ raise AirflowException(f"{remote_full_path} already exists")
+ await self.create_directory(remote_full_path)
+ for root, dirs, files in os.walk(local_full_path):
Review Comment:
`os.walk()` runs inline in an `async def`, so it blocks the Triggerer's
event loop for the whole walk — and with it every other deferred task in that
process. That is the specific cost `deferrable=True` exists to avoid, so it is
worth fixing before this lands.
What makes it stand out is that the wrapping is inverted either side of it:
`os.path.relpath` just below is pure string manipulation and gains nothing from
`asyncio.to_thread`, while `os.walk` is the one doing real filesystem I/O and
is the one left unwrapped.
Collecting the walk in one `asyncio.to_thread` call and then iterating the
result would fix it, and the `relpath` hops could go back to being plain calls.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -874,28 +948,35 @@ async def retrieve_file(
:param remote_full_path: Full path to the remote file.
:param local_full_path: Full path to the local file or a binary
file-like buffer.
:param chunk_size: Size of chunks to read at a time (default: 64KB).
+ :param prefetch: Whether to allow pipelined read-ahead requests
(default: True). When
+ ``False``, only one request is kept in flight at a time, mirroring
+ :meth:`SFTPHook.retrieve_file`'s ``prefetch`` semantics.
"""
+ if isinstance(local_full_path, (str, os.PathLike)):
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ get_kwargs: dict[str, Any] = {"block_size": chunk_size}
+ if not prefetch:
+ get_kwargs["max_requests"] = 1
+ await sftp.get(remote_full_path, local_full_path,
**get_kwargs)
Review Comment:
mypy: `Argument 2 to "get" of "SFTPClient" has incompatible type "str |
PathLike[str]"; expected "bytes | str | PurePath | None"`.
The narrowing above proves this branch only runs for `str |
os.PathLike[str]`, but asyncssh types the parameter as accepting `PurePath`
rather than the wider `os.PathLike`. Passing `os.fspath(local_full_path)` (or
`str(...)`, matching what `store_file` already does on its `sftp.put` call)
satisfies it without weakening the signature.
##########
providers/sftp/src/airflow/providers/sftp/operators/sftp.py:
##########
@@ -136,106 +139,85 @@ def execute(self, context: Any) -> str | list[str] |
None:
if self.operation.lower() not in (SFTPOperation.GET,
SFTPOperation.PUT, SFTPOperation.DELETE):
raise TypeError(
f"Unsupported operation value {self.operation}, "
- f"expected {SFTPOperation.GET} or {SFTPOperation.PUT} or
{SFTPOperation.DELETE}."
+ f"expected {SFTPOperation.GET!r}, {SFTPOperation.PUT!r}, "
+ f"or {SFTPOperation.DELETE!r}."
)
if self.concurrency < 1:
- raise ValueError(f"concurrency should be greater than 0, got
{self.concurrency}")
+ raise ValueError(f"concurrency should be >= 1, got
{self.concurrency}")
- file_msg = None
- try:
- if self.remote_host is not None:
- self.log.info(
- "remote_host is provided explicitly. "
- "It will replace the remote_host which was defined "
- "in sftp_hook or predefined in connection of ssh_conn_id."
+ # ------------------------------------------------------------------ #
+ # Synchronous path — delegate all transfer logic to the hook #
+ # ------------------------------------------------------------------ #
+ if self.remote_host is not None:
+ self.log.info(
+ "remote_host is provided explicitly. "
+ "It will replace the remote_host which was defined "
+ "in sftp_hook or predefined in connection of ssh_conn_id."
+ )
+
+ if self.ssh_conn_id:
+ if self.sftp_hook and isinstance(self.sftp_hook, SFTPHook):
+ self.log.info("ssh_conn_id is ignored when sftp_hook is
provided.")
+ else:
+ self.log.info("sftp_hook not provided or invalid. Trying
ssh_conn_id to create SFTPHook.")
+ self.sftp_hook = SFTPHook(
+ ssh_conn_id=self.ssh_conn_id,
+ remote_host=self.remote_host or "",
)
- if self.ssh_conn_id:
- if self.sftp_hook and isinstance(self.sftp_hook, SFTPHook):
- self.log.info("ssh_conn_id is ignored when sftp_hook is
provided.")
- else:
- self.log.info("sftp_hook not provided or invalid. Trying
ssh_conn_id to create SFTPHook.")
- self.sftp_hook = SFTPHook(
- ssh_conn_id=self.ssh_conn_id,
remote_host=self.remote_host or ""
- )
-
- if not self.sftp_hook:
- raise AirflowException("Cannot operate without sftp_hook or
ssh_conn_id.")
-
- if self.operation.lower() in (SFTPOperation.GET,
SFTPOperation.PUT):
- for _local_filepath, _remote_filepath in
zip(local_filepath_array, remote_filepath_array):
- if self.operation.lower() == SFTPOperation.GET:
- 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}"
- self.log.info("Starting to transfer %s", file_msg)
- if self.sftp_hook.isdir(_remote_filepath):
- if self.concurrency > 1:
- self.sftp_hook.retrieve_directory_concurrently(
- _remote_filepath,
- _local_filepath,
- workers=self.concurrency,
- prefetch=self.prefetch,
- )
- 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:
- remote_folder = os.path.dirname(_remote_filepath)
- if self.create_intermediate_dirs:
- self.sftp_hook.create_directory(remote_folder)
- 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):
- 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:
- for _remote_filepath in remote_filepath_array:
- file_msg = f"{_remote_filepath}"
- self.log.info("Starting to delete %s", file_msg)
- try:
- if self.sftp_hook.isdir(_remote_filepath):
- self.sftp_hook.delete_directory(_remote_filepath,
include_files=True)
- else:
- self.sftp_hook.delete_file(_remote_filepath)
- except OSError as exc:
- if self._is_missing_path_error(exc):
- self.log.warning(
- "Remote path %s does not exist. Skipping
delete.", _remote_filepath
- )
- continue
- raise
+ if not self.sftp_hook:
+ raise AirflowException("Cannot operate without sftp_hook or
ssh_conn_id.")
+
+ if self.deferrable:
+ from airflow.providers.sftp.triggers.sftp import
SFTPTransferTrigger
+
+ self.defer(
+ trigger=SFTPTransferTrigger(
+ sftp_conn_id=self.ssh_conn_id or
self.sftp_hook.ssh_conn_id,
Review Comment:
mypy: `Argument "sftp_conn_id" ... has incompatible type "str | None";
expected "str"`.
The fix itself is right — this is the behaviour I asked for. The type falls
out of `SFTPHook.ssh_conn_id` being `str | None`, so when neither `ssh_conn_id`
nor the hook's own conn id is set, `None` reaches a parameter typed `str`.
Worth deciding what that case should actually do rather than casting it
away: reaching here with no connection id at all is a misconfiguration the
deferred path cannot recover from, so failing loudly in `execute()` seems
better than handing `None` to the trigger.
##########
providers/sftp/tests/unit/sftp/operators/test_sftp.py:
##########
@@ -673,3 +691,170 @@ def test_extract_sftp_hook(self, get_connection,
get_conn, operation, expected):
assert lineage.inputs == expected[0]
assert lineage.outputs == expected[1]
+
+
+class TestSFTPOperatorDeferrable:
+ """Tests for SFTPOperator deferrable mode."""
+
+ def test_sftp_operator_defers_when_deferrable_true(self):
+ """Test that SFTPOperator defers when deferrable=True."""
+ operator = SFTPOperator(
+ task_id="test_sftp_defer",
+ ssh_conn_id="ssh_default",
+ local_filepath="/tmp/test.txt",
+ remote_filepath="/remote/test.txt",
+ operation=SFTPOperation.PUT,
+ deferrable=True,
+ )
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context={})
+ assert isinstance(exc.value.trigger, SFTPTransferTrigger)
+ assert exc.value.method_name == "execute_complete"
+
+ def
test_sftp_operator_defer_uses_sftp_hook_conn_id_when_ssh_conn_id_unset(self):
+ """
+ Assert that deferring honors a supplied sftp_hook's connection id.
+
+ Regression test: previously, when only ``sftp_hook`` (not
``ssh_conn_id``) was
+ provided, the trigger silently fell back to
``SFTPHookAsync.default_conn_name``
+ ("sftp_default") instead of the hook's actual connection, redirecting
the
+ deferred transfer to the wrong server.
+ """
+ operator = SFTPOperator(
+ task_id="test_sftp_defer_hook_conn_id",
+ sftp_hook=SFTPHook(ssh_conn_id="my_prod_sftp"),
Review Comment:
This is the regression test for the `sftp_hook` fix, and it is currently
failing in six jobs (LowestDeps plus all five Compat runs) with
`AirflowNotFoundException: The conn_id 'my_prod_sftp' isn't defined`.
It fails on this line rather than inside `execute()`: `SSHHook.__init__`
resolves the connection eagerly, so the hook cannot be constructed unless
`my_prod_sftp` exists. The operator code under test is fine.
Giving the test a connection fixes it — the `create_connection_without_db`
fixture, or `monkeypatch.setenv("AIRFLOW_CONN_MY_PROD_SFTP", '{"conn_type":
"sftp", "host": "example.com"}')` if you would rather keep it db-free.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1080,3 +1170,283 @@ 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 sense_files_by_pattern(
+ self,
+ path: str,
+ fnmatch_pattern: str,
+ newer_than: datetime.datetime | None = None,
+ ) -> list[str]:
+ """
+ Return the names of files at ``path`` matching ``fnmatch_pattern``.
+
+ If ``newer_than`` is provided, only files modified after that
timestamp are returned; files
+ without a reported modification time are skipped in that case.
+
+ :param path: directory on the SFTP server to search for files matching
the pattern
+ :param fnmatch_pattern: pattern used to match filenames, see the
``fnmatch`` std library module
+ :param newer_than: if provided, only files modified after this UTC
timestamp are returned
+ """
+ files = await self.get_files_and_attrs_by_pattern(path=path,
fnmatch_pattern=fnmatch_pattern)
+ if not newer_than:
+ return [str(file.filename) for file in files]
+
+ matched_files = []
+ for file in files:
+ if file.attrs.mtime is None:
+ continue
+ if newer_than <= self._mod_time_to_utc(file.attrs.mtime):
+ matched_files.append(str(file.filename))
+ return matched_files
+
+ async def sense_path(self, path: str, newer_than: datetime.datetime | None
= None) -> bool:
+ """
+ Return whether ``path`` exists and, if ``newer_than`` is provided, was
modified since.
+
+ :param path: full path to the remote file
+ :param newer_than: if provided, the file must have been modified after
this UTC timestamp
+ """
+ mod_time = await self.get_mod_time(path)
+ if not newer_than:
+ return True
+ return newer_than <= self._mod_time_to_utc(mod_time)
+
+ @staticmethod
+ def _mod_time_to_utc(mod_time: int | float | str) -> datetime.datetime:
+ """Convert a modification time, either an epoch timestamp or
``%Y%m%d%H%M%S`` string, to UTC."""
+ if not isinstance(mod_time, str):
+ mod_time =
datetime.datetime.fromtimestamp(float(mod_time)).strftime("%Y%m%d%H%M%S")
+ return timezone.convert_to_utc(datetime.datetime.strptime(mod_time,
"%Y%m%d%H%M%S"))
+
+ async def isdir(self, path: str) -> bool:
+ """
+ Check if the path provided is a directory.
+
+ :param path: full path to the remote directory to check
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ attrs = await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return attrs.permissions is not None and
stat.S_ISDIR(attrs.permissions)
+
+ async def path_exists(self, path: str) -> bool:
+ """
+ Whether a remote entity exists.
+
+ :param path: full path to the remote file or directory
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return True
+
+ async def create_directory(self, path: str) -> None:
+ """
+ Create a directory (and any missing parents) on the remote system
asynchronously.
+
+ Returns silently if the target directory already exists, mirroring
+ :meth:`SFTPHook.create_directory`.
+
+ :param path: full path to the remote directory to create
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.makedirs(path, exist_ok=True)
+
+ async def delete_file(self, path: str) -> None:
+ """
+ Remove a file on the server asynchronously.
+
+ :param path: full path to the remote file
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.unlink(path)
+
+ async def delete_directory(self, path: str, include_files: bool = False)
-> None:
+ """
+ Delete a directory on the remote system asynchronously.
+
+ :param path: full path to the remote directory to delete
+ :param include_files: whether to recursively delete the directory's
contents first
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+
+ if include_files:
+ files, dirs, _ = await self.get_tree_map(path)
+ dirs = dirs[::-1] # reverse the order for deleting deepest
directories first
+
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ for file_path in files:
+ await sftp.remove(file_path)
+ for dir_path in dirs:
+ await sftp.rmdir(dir_path)
+ await sftp.rmdir(path)
+
+ async def get_tree_map(
+ self, path: str, prefix: str | None = None, delimiter: str | None =
None
+ ) -> tuple[list[str], list[str], list[str]]:
+ """
+ Get tuple with recursive lists of files, directories and unknown paths
asynchronously.
+
+ It is possible to filter results by giving prefix and/or delimiter
parameters.
+
+ :param path: path from which tree will be built
+ :param prefix: if set paths will be added if start with prefix
+ :param delimiter: if set paths will be added if end with delimiter
+ :return: tuple with list of files, dirs and unknown items
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+ unknowns: list[str] = []
+
+ def append_matching_path_callback(list_: list[str]) -> Callable:
+ return lambda item: (
+ list_.append(item) if SFTPHook._is_path_match(item, prefix,
delimiter) else None
+ )
+
+ await self.walktree(
+ path=path,
+ fcallback=append_matching_path_callback(files),
+ dcallback=append_matching_path_callback(dirs),
+ ucallback=append_matching_path_callback(unknowns),
+ recurse=True,
+ )
+
+ return files, dirs, unknowns
+
+ async def retrieve_directory(
+ self, remote_full_path: str, local_full_path: str, prefetch: bool =
True
+ ) -> None:
+ """
+ Transfer the remote directory to a local location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param prefetch: controls whether read-ahead requests are pipelined
(default: True)
+ """
+ if await asyncio.to_thread(Path(local_full_path).exists):
+ raise AirflowException(f"{local_full_path} already exists")
Review Comment:
This is one of the two lines failing `check-no-new-airflow-exceptions`,
which is why Static checks is red.
`AGENTS.md` is explicit that new direct `raise AirflowException(...)` usages
are not accepted in `providers/`. The synchronous `retrieve_directory` you
mirrored does raise it, but that one is pre-existing — copying it into a new
method counts as a new usage, and the hook scans the whole tree rather than
just the diff.
This PR already adds `SFTPOperationError` in
`providers/sftp/src/airflow/providers/sftp/exceptions.py`, which fits; a plain
`FileExistsError` would read well here too, since that is literally the
condition.
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -1080,3 +1170,283 @@ 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 sense_files_by_pattern(
+ self,
+ path: str,
+ fnmatch_pattern: str,
+ newer_than: datetime.datetime | None = None,
+ ) -> list[str]:
+ """
+ Return the names of files at ``path`` matching ``fnmatch_pattern``.
+
+ If ``newer_than`` is provided, only files modified after that
timestamp are returned; files
+ without a reported modification time are skipped in that case.
+
+ :param path: directory on the SFTP server to search for files matching
the pattern
+ :param fnmatch_pattern: pattern used to match filenames, see the
``fnmatch`` std library module
+ :param newer_than: if provided, only files modified after this UTC
timestamp are returned
+ """
+ files = await self.get_files_and_attrs_by_pattern(path=path,
fnmatch_pattern=fnmatch_pattern)
+ if not newer_than:
+ return [str(file.filename) for file in files]
+
+ matched_files = []
+ for file in files:
+ if file.attrs.mtime is None:
+ continue
+ if newer_than <= self._mod_time_to_utc(file.attrs.mtime):
+ matched_files.append(str(file.filename))
+ return matched_files
+
+ async def sense_path(self, path: str, newer_than: datetime.datetime | None
= None) -> bool:
+ """
+ Return whether ``path`` exists and, if ``newer_than`` is provided, was
modified since.
+
+ :param path: full path to the remote file
+ :param newer_than: if provided, the file must have been modified after
this UTC timestamp
+ """
+ mod_time = await self.get_mod_time(path)
+ if not newer_than:
+ return True
+ return newer_than <= self._mod_time_to_utc(mod_time)
+
+ @staticmethod
+ def _mod_time_to_utc(mod_time: int | float | str) -> datetime.datetime:
+ """Convert a modification time, either an epoch timestamp or
``%Y%m%d%H%M%S`` string, to UTC."""
+ if not isinstance(mod_time, str):
+ mod_time =
datetime.datetime.fromtimestamp(float(mod_time)).strftime("%Y%m%d%H%M%S")
+ return timezone.convert_to_utc(datetime.datetime.strptime(mod_time,
"%Y%m%d%H%M%S"))
+
+ async def isdir(self, path: str) -> bool:
+ """
+ Check if the path provided is a directory.
+
+ :param path: full path to the remote directory to check
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ attrs = await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return attrs.permissions is not None and
stat.S_ISDIR(attrs.permissions)
+
+ async def path_exists(self, path: str) -> bool:
+ """
+ Whether a remote entity exists.
+
+ :param path: full path to the remote file or directory
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ try:
+ await sftp.stat(path)
+ except asyncssh.SFTPNoSuchFile:
+ return False
+ return True
+
+ async def create_directory(self, path: str) -> None:
+ """
+ Create a directory (and any missing parents) on the remote system
asynchronously.
+
+ Returns silently if the target directory already exists, mirroring
+ :meth:`SFTPHook.create_directory`.
+
+ :param path: full path to the remote directory to create
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.makedirs(path, exist_ok=True)
+
+ async def delete_file(self, path: str) -> None:
+ """
+ Remove a file on the server asynchronously.
+
+ :param path: full path to the remote file
+ """
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ await sftp.unlink(path)
+
+ async def delete_directory(self, path: str, include_files: bool = False)
-> None:
+ """
+ Delete a directory on the remote system asynchronously.
+
+ :param path: full path to the remote directory to delete
+ :param include_files: whether to recursively delete the directory's
contents first
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+
+ if include_files:
+ files, dirs, _ = await self.get_tree_map(path)
+ dirs = dirs[::-1] # reverse the order for deleting deepest
directories first
+
+ async with await self._get_conn() as ssh_conn:
+ async with ssh_conn.start_sftp_client() as sftp:
+ for file_path in files:
+ await sftp.remove(file_path)
+ for dir_path in dirs:
+ await sftp.rmdir(dir_path)
+ await sftp.rmdir(path)
+
+ async def get_tree_map(
+ self, path: str, prefix: str | None = None, delimiter: str | None =
None
+ ) -> tuple[list[str], list[str], list[str]]:
+ """
+ Get tuple with recursive lists of files, directories and unknown paths
asynchronously.
+
+ It is possible to filter results by giving prefix and/or delimiter
parameters.
+
+ :param path: path from which tree will be built
+ :param prefix: if set paths will be added if start with prefix
+ :param delimiter: if set paths will be added if end with delimiter
+ :return: tuple with list of files, dirs and unknown items
+ """
+ files: list[str] = []
+ dirs: list[str] = []
+ unknowns: list[str] = []
+
+ def append_matching_path_callback(list_: list[str]) -> Callable:
+ return lambda item: (
+ list_.append(item) if SFTPHook._is_path_match(item, prefix,
delimiter) else None
+ )
+
+ await self.walktree(
+ path=path,
+ fcallback=append_matching_path_callback(files),
+ dcallback=append_matching_path_callback(dirs),
+ ucallback=append_matching_path_callback(unknowns),
+ recurse=True,
+ )
+
+ return files, dirs, unknowns
+
+ async def retrieve_directory(
+ self, remote_full_path: str, local_full_path: str, prefetch: bool =
True
+ ) -> None:
+ """
+ Transfer the remote directory to a local location asynchronously.
+
+ :param remote_full_path: full path to the remote directory
+ :param local_full_path: full path to the local directory
+ :param prefetch: controls whether read-ahead requests are pipelined
(default: True)
+ """
+ if await asyncio.to_thread(Path(local_full_path).exists):
+ raise AirflowException(f"{local_full_path} already exists")
+ dest = await asyncio.to_thread(Path(local_full_path).resolve)
+ await asyncio.to_thread(dest.mkdir, parents=True)
+ files, dirs, _ = await self.get_tree_map(remote_full_path)
+ for dir_path in dirs:
+ relative_path = await asyncio.to_thread(os.path.relpath, dir_path,
remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await asyncio.to_thread(Path(new_local_path).mkdir, parents=True,
exist_ok=True)
+ for file_path in files:
+ relative_path = await asyncio.to_thread(os.path.relpath,
file_path, remote_full_path)
+ new_local_path = str(dest / relative_path)
+ SFTPHook._validate_within_directory(str(dest), new_local_path)
+ await self.retrieve_file(file_path, new_local_path,
prefetch=prefetch)
Review Comment:
Each `retrieve_file` opens its own connection and SFTP client, so a
directory transfer costs one SSH handshake per file, plus one for
`get_tree_map`, plus one for the `isdir` check in `transfer()`. The synchronous
path walks and downloads the whole tree over a single connection.
In fairness this is partly downstream of my last review: I asked for the
unused outer connection to go, and it did — but the result is that nothing is
shared at all, and adding directory support is what makes that bite. A private
helper that opens one connection and passes the `sftp` client down to the
per-file work would satisfy both points at once.
Not a blocker on its own, but better raised now than after someone points a
directory GET at a few thousand files.
--
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]
