diff --git a/providers/sftp/docs/index.rst b/providers/sftp/docs/index.rst index 63954b645e141..6bb968353bb5a 100644 --- a/providers/sftp/docs/index.rst +++ b/providers/sftp/docs/index.rst @@ -105,6 +105,7 @@ PIP package Version required ``apache-airflow-providers-ssh`` ``>=6.0.0`` ``apache-airflow-providers-common-compat`` ``>=1.12.0`` ``paramiko`` ``>=4.0.0,<5.0.0`` +``asgiref`` ``>=3.8.1`` ``asyncssh`` ``>=2.12.0; python_version < "3.14"`` ``asyncssh`` ``>=2.22.0; python_version >= "3.14"`` ========================================== ====================================== diff --git a/providers/sftp/newsfragments/68298.feature.rst b/providers/sftp/newsfragments/68298.feature.rst new file mode 100644 index 0000000000000..cca7e43ceb2ea --- /dev/null +++ b/providers/sftp/newsfragments/68298.feature.rst @@ -0,0 +1 @@ +Add deferrable mode to ``SFTPOperator`` via ``deferrable=True`` parameter, freeing Airflow worker slots during file transfers by delegating execution to the async Triggerer. diff --git a/providers/sftp/pyproject.toml b/providers/sftp/pyproject.toml index b59f46da49221..fc75ddbcb1be7 100644 --- a/providers/sftp/pyproject.toml +++ b/providers/sftp/pyproject.toml @@ -64,6 +64,7 @@ dependencies = [ "apache-airflow-providers-ssh>=6.0.0", "apache-airflow-providers-common-compat>=1.12.0", "paramiko>=4.0.0,<5.0.0", + "asgiref>=3.8.1", "asyncssh>=2.12.0; python_version < '3.14'", "asyncssh>=2.22.0; python_version >= '3.14'", ] diff --git a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py index cbb4254c8a304..117f66498020a 100644 --- a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py +++ b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py @@ -19,6 +19,7 @@ from __future__ import annotations +import asyncio import concurrent.futures import datetime import functools @@ -28,6 +29,7 @@ import warnings from collections.abc import Callable, Generator, Sequence from contextlib import contextmanager, suppress +from enum import Enum from fnmatch import fnmatch from io import BytesIO from pathlib import Path, PurePosixPath @@ -51,6 +53,14 @@ CHUNK_SIZE = 64 * 1024 # 64KB +class SFTPOperation(str, Enum): + """SFTP operation constants.""" + + GET = "get" + PUT = "put" + DELETE = "delete" + + def handle_connection_management(func: Callable) -> Callable: @functools.wraps(func) def handle_connection_management_wrapper(self, *args: Any, **kwargs: dict[str, Any]) -> Any: @@ -384,25 +394,6 @@ def delete_file(self, path: str) -> None: """ self.conn.remove(path) # type: ignore[arg-type, union-attr] - @staticmethod - def _validate_within_directory(base_dir: str, candidate: str) -> str: - """ - Ensure ``candidate`` resolves to a path inside ``base_dir``. - - Directory-entry names are returned by the remote SFTP server and may - contain ``..`` components; joining them into the local destination path - could otherwise write outside it. Containment is verified before any - local write or ``mkdir``. - """ - base_real = os.path.realpath(base_dir) - candidate_real = os.path.realpath(candidate) - if candidate_real != base_real and os.path.commonpath([base_real, candidate_real]) != base_real: - raise ValueError( - f"Refusing to write outside the destination directory: " - f"{candidate!r} resolves outside {base_dir!r}" - ) - return candidate - def retrieve_directory(self, remote_full_path: str, local_full_path: str, prefetch: bool = True) -> None: """ Transfer the remote directory to a local location. @@ -419,14 +410,10 @@ def retrieve_directory(self, remote_full_path: str, local_full_path: str, prefet Path(local_full_path).mkdir(parents=True) files, dirs, _ = self.get_tree_map(remote_full_path) for dir_path in dirs: - new_local_path = self._validate_within_directory( - local_full_path, os.path.join(local_full_path, os.path.relpath(dir_path, remote_full_path)) - ) + 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_path in files: - new_local_path = self._validate_within_directory( - local_full_path, os.path.join(local_full_path, os.path.relpath(file_path, remote_full_path)) - ) + 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( @@ -461,18 +448,12 @@ def retrieve_file_chunk( new_local_file_paths, remote_file_paths = [], [] files, dirs, _ = self.get_tree_map(remote_full_path) for dir_path in dirs: - new_local_path = self._validate_within_directory( - local_full_path, - os.path.join(local_full_path, os.path.relpath(dir_path, remote_full_path)), - ) + 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( - self._validate_within_directory( - local_full_path, - os.path.join(local_full_path, os.path.relpath(file, remote_full_path)), - ) + 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)] @@ -625,6 +606,27 @@ def _is_path_match(path: str, prefix: str | None = None, delimiter: str | None = return False return True + @staticmethod + def _validate_within_directory(base: str, target: str) -> str: + """ + Validate that target path is within the base directory. + + Prevents directory traversal attacks. + + :param base: The base/destination directory path + :param target: The target path to validate + :return: The target path if valid + :raises ValueError: If target path escapes the base directory + """ + base_real = os.path.realpath(os.path.expanduser(base)) + target_real = os.path.realpath(os.path.expanduser(target)) + + # Ensure target is within base directory + if not (target_real == base_real or target_real.startswith(base_real + os.sep)): + raise ValueError(f"Path {target} is outside the destination directory {base}") + + return target + def walktree( self, path: str, @@ -735,6 +737,71 @@ 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). + + Centralizes 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). + """ + 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) + class SFTPHookAsync(BaseHook): """ @@ -1078,3 +1145,64 @@ 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.""" + 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) diff --git a/providers/sftp/src/airflow/providers/sftp/operators/sftp.py b/providers/sftp/src/airflow/providers/sftp/operators/sftp.py index 0b47b5b7d5ddb..98037281e7faf 100644 --- a/providers/sftp/src/airflow/providers/sftp/operators/sftp.py +++ b/providers/sftp/src/airflow/providers/sftp/operators/sftp.py @@ -19,11 +19,8 @@ from __future__ import annotations -import errno -import os import socket from collections.abc import Sequence -from pathlib import Path from typing import Any import paramiko @@ -32,31 +29,26 @@ from airflow.providers.sftp.hooks.sftp import SFTPHook -class SFTPOperation: - """Operation that can be used with SFTP.""" - - PUT = "put" - GET = "get" - DELETE = "delete" - - class SFTPOperator(BaseOperator): """ SFTPOperator for transferring files from remote host to local or vice a versa. - This operator uses sftp_hook to open sftp transport channel that serve as basis for file transfer. + This operator uses sftp_hook to open an SFTP transport channel that serves as + the basis for file transfer. All transfer logic is delegated to + ``SFTPHook.transfer()`` so that both the synchronous and deferrable code paths + share a single, authoritative implementation. :param ssh_conn_id: :ref:`ssh connection id` from airflow Connections. - :param sftp_hook: predefined SFTPHook to use + :param sftp_hook: predefined SFTPHook to use. Either `sftp_hook` or `ssh_conn_id` needs to be provided. - :param remote_host: remote host to connect (templated) + :param remote_host: remote host to connect (templated). Nullable. If provided, it will replace the `remote_host` which was defined in `sftp_hook` or predefined in the connection of `ssh_conn_id`. :param local_filepath: local file path or list of local file paths to get or put. (templated) :param remote_filepath: remote file path or list of remote file paths to get, put, or delete. (templated) - :param operation: specify operation 'get', 'put', or 'delete', defaults to put - :param confirm: specify if the SFTP operation should be confirmed, defaults to True + :param operation: specify operation ``'get'``, ``'put'``, or ``'delete'``. Defaults to ``'put'``. + :param confirm: specify if the SFTP operation should be confirmed. Defaults to True. :param create_intermediate_dirs: create missing intermediate directories when copying from remote to local and vice-versa. Default is False. @@ -74,10 +66,15 @@ 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) - :param prefetch: controls whether prefetch is performed (default: True) + :param concurrency: number of threads when transferring directories. Each thread opens + a new SFTP connection. Only applies to directory transfers. (Default: 1) + :param prefetch: controls whether prefetch is performed on GET transfers. (Default: True) + :param deferrable: run the operator in deferrable mode. When True, the worker slot is + freed during the transfer and reclaimed only when the transfer completes. + Best suited for single large file transfers. For bulk directory transfers involving + many files, consider using ``async PythonOperator`` with ``SFTPClientPool`` instead, + which provides true async multiplexing via a single event loop. (Default: False) """ template_fields: Sequence[str] = ("local_filepath", "remote_filepath", "remote_host") @@ -95,6 +92,7 @@ def __init__( create_intermediate_dirs: bool = False, concurrency: int = 1, prefetch: bool = True, + deferrable: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -108,134 +106,99 @@ 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: + local_filepath_array: list[str] = [] if self.local_filepath is None: local_filepath_array = [] elif isinstance(self.local_filepath, str): local_filepath_array = [self.local_filepath] else: - local_filepath_array = self.local_filepath + local_filepath_array = list(self.local_filepath) - if isinstance(self.remote_filepath, str): - remote_filepath_array = [self.remote_filepath] - else: - remote_filepath_array = self.remote_filepath + remote_filepath_array: list[str] = ( + [self.remote_filepath] if isinstance(self.remote_filepath, str) else list(self.remote_filepath) + ) - if self.operation.lower() in (SFTPOperation.GET, SFTPOperation.PUT) and len( - local_filepath_array - ) != len(remote_filepath_array): + # ------------------------------------------------------------------ # + # Input validation # + # ------------------------------------------------------------------ # + if self.operation in (SFTPOperation.GET, SFTPOperation.PUT) and len(local_filepath_array) != len( + remote_filepath_array + ): raise ValueError( f"{len(local_filepath_array)} paths in local_filepath " f"!= {len(remote_filepath_array)} paths in remote_filepath" ) - if self.operation.lower() == SFTPOperation.DELETE and local_filepath_array: + if self.operation == SFTPOperation.DELETE and local_filepath_array: raise ValueError("local_filepath should not be provided for delete operation") - if self.operation.lower() not in (SFTPOperation.GET, SFTPOperation.PUT, SFTPOperation.DELETE): + if self.operation 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.") + try: + for idx, remote_fp in enumerate(remote_filepath_array): + local_fp = local_filepath_array[idx] if local_filepath_array else "" + self.sftp_hook.transfer( + local_filepath=local_fp, + remote_filepath=remote_fp, + operation=self.operation, + confirm=self.confirm, + create_intermediate_dirs=self.create_intermediate_dirs, + concurrency=self.concurrency, + prefetch=self.prefetch, + ) except Exception as e: raise AirflowException( - f"Error while processing {self.operation.upper()} operation {file_msg}, error: {e}" - ) + f"Error while processing {self.operation.upper()} operation, error: {e}" + ) from e return self.local_filepath - @staticmethod - def _is_missing_path_error(exc: Exception) -> bool: - if isinstance(exc, FileNotFoundError): - return True - if isinstance(exc, OSError) and exc.errno == errno.ENOENT: - return True - if exc.args and isinstance(exc.args[0], int) and exc.args[0] == errno.ENOENT: - return True - return False + def execute_complete(self, context: Any, event: dict[str, Any]) -> str | list[str] | None: + """ + Handle completion from ``SFTPOperatorTrigger``. + + :param context: Airflow task context + :param event: trigger result dict with ``status`` and ``message`` keys + :raises AirflowException: if the trigger reported an error + """ + if event.get("status") == "error": + raise AirflowException(event.get("message", "Unknown error during deferrable SFTP transfer")) + self.log.info("Deferrable SFTP transfer completed: %s", event.get("message")) + return self.local_filepath def get_openlineage_facets_on_start(self): """ @@ -279,10 +242,6 @@ def get_openlineage_facets_on_start(self): if hasattr(hook, "port"): remote_port = hook.port - # Since v4.1.0, SFTPOperator accepts both a string (single file) and a list of - # strings (multiple files) as local_filepath and remote_filepath, and internally - # keeps them as list in both cases. But before 4.1.0, only single string is - # allowed. So we consider both cases here for backward compatibility. if isinstance(self.local_filepath, str): local_filepath = [self.local_filepath] else: @@ -301,7 +260,7 @@ def get_openlineage_facets_on_start(self): for path in remote_filepath ] - if self.operation.lower() == SFTPOperation.GET: + if self.operation == SFTPOperation.GET: inputs = remote_datasets outputs = local_datasets else: diff --git a/providers/sftp/src/airflow/providers/sftp/triggers/sftp.py b/providers/sftp/src/airflow/providers/sftp/triggers/sftp.py index a46f29d5a4abe..1f1379947778d 100644 --- a/providers/sftp/src/airflow/providers/sftp/triggers/sftp.py +++ b/providers/sftp/src/airflow/providers/sftp/triggers/sftp.py @@ -24,9 +24,11 @@ from dateutil.parser import parse as parse_date -from airflow.providers.common.compat.sdk import AirflowException, timezone +from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.sftp.operators.sftp import SFTPOperation from airflow.providers.sftp.hooks.sftp import SFTPHookAsync from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.utils.timezone import convert_to_utc class SFTPTrigger(BaseTrigger): @@ -86,7 +88,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: if isinstance(self.newer_than, str): self.newer_than = parse_date(self.newer_than) - _newer_than = timezone.convert_to_utc(self.newer_than) if self.newer_than else None + _newer_than = convert_to_utc(self.newer_than) if self.newer_than else None while True: try: if self.file_pattern: @@ -101,9 +103,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: mod_time = datetime.fromtimestamp(float(file.attrs.mtime)).strftime( "%Y%m%d%H%M%S" ) - mod_time_utc = timezone.convert_to_utc( - datetime.strptime(mod_time, "%Y%m%d%H%M%S") - ) + mod_time_utc = convert_to_utc(datetime.strptime(mod_time, "%Y%m%d%H%M%S")) if _newer_than <= mod_time_utc: files_sensed.append(file.filename) else: @@ -119,7 +119,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: else: mod_time = await hook.get_mod_time(self.path) if _newer_than: - mod_time_utc = timezone.convert_to_utc(datetime.strptime(mod_time, "%Y%m%d%H%M%S")) + mod_time_utc = convert_to_utc(datetime.strptime(mod_time, "%Y%m%d%H%M%S")) if _newer_than <= mod_time_utc: yield TriggerEvent({"status": "success", "message": f"Sensed file: {self.path}"}) return @@ -140,3 +140,86 @@ async def run(self) -> AsyncIterator[TriggerEvent]: def _get_async_hook(self) -> SFTPHookAsync: return SFTPHookAsync(sftp_conn_id=self.sftp_conn_id) + + +class SFTPOperationTrigger(BaseTrigger): + """ + Trigger for SFTPOperator deferrable mode. + + Fires when a file transfer (PUT, GET, or DELETE) completes + on the SFTP server, freeing the worker slot during the transfer. + + :param ssh_conn_id: The SSH connection ID to use. + :param local_filepath: Local file path(s) to transfer. + :param remote_filepath: Remote file path(s) on the SFTP server. + :param operation: The SFTP operation - put, get, or delete. + :param confirm: Whether to confirm the file transfer. + :param create_intermediate_dirs: Whether to create intermediate dirs. + :param remote_host: Remote host to connect to (overrides connection). + :param concurrency: Number of threads for directory transfers. + :param prefetch: Whether to prefetch during file retrieval. + """ + + def __init__( + self, + ssh_conn_id: str | None = None, + local_filepath: str | list[str] | None = None, + remote_filepath: str | list[str] = "", + operation: str = SFTPOperation.PUT, + confirm: bool = True, + create_intermediate_dirs: bool = False, + remote_host: str | None = None, + concurrency: int = 1, + prefetch: bool = True, + ) -> None: + super().__init__() + self.ssh_conn_id = ssh_conn_id + self.local_filepath = local_filepath + self.remote_filepath = remote_filepath + self.operation = operation + self.confirm = confirm + self.create_intermediate_dirs = create_intermediate_dirs + self.remote_host = remote_host + self.concurrency = concurrency + self.prefetch = prefetch + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize the trigger for storage in the database.""" + return ( + f"{self.__class__.__module__}.{self.__class__.__name__}", + { + "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, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Run the file transfer asynchronously and yield a TriggerEvent when done.""" + try: + if self.ssh_conn_id is None: + raise ValueError("ssh_conn_id must be set for SFTPTrigger") + hook = SFTPHookAsync(sftp_conn_id=self.ssh_conn_id) + await hook.transfer( + operation=self.operation, + local_filepath=self.local_filepath, + remote_filepath=self.remote_filepath, + confirm=self.confirm, + create_intermediate_dirs=self.create_intermediate_dirs, + concurrency=self.concurrency, + prefetch=self.prefetch, + ) + yield TriggerEvent( + { + "status": "success", + "local_filepath": self.local_filepath, + } + ) + except Exception as e: + yield TriggerEvent({"status": "error", "message": str(e)}) \ No newline at end of file diff --git a/providers/sftp/tests/unit/sftp/operators/test_sftp.py b/providers/sftp/tests/unit/sftp/operators/test_sftp.py index 815d981320107..a01d2c79ee51f 100644 --- a/providers/sftp/tests/unit/sftp/operators/test_sftp.py +++ b/providers/sftp/tests/unit/sftp/operators/test_sftp.py @@ -1,4 +1,3 @@ -# # 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 @@ -28,11 +27,13 @@ import paramiko import pytest +from airflow.exceptions import TaskDeferred from airflow.models import DAG, Connection from airflow.providers.common.compat.openlineage.facet import Dataset from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.sftp.hooks.sftp import SFTPHook from airflow.providers.sftp.operators.sftp import SFTPOperation, SFTPOperator +from airflow.providers.sftp.triggers.sftp import SFTPOperationTrigger from airflow.providers.ssh.hooks.ssh import SSHHook from airflow.providers.ssh.operators.ssh import SSHOperator from airflow.utils import timezone @@ -675,3 +676,124 @@ 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, SFTPOperationTrigger) + assert exc.value.method_name == "execute_complete" + + def test_sftp_operator_execute_complete_success(self): + """Test execute_complete returns local_filepath on success.""" + operator = SFTPOperator( + task_id="test_sftp_complete", + ssh_conn_id="ssh_default", + local_filepath="/tmp/test.txt", + remote_filepath="/remote/test.txt", + operation=SFTPOperation.PUT, + deferrable=True, + ) + event = {"status": "success", "local_filepath": "/tmp/test.txt"} + result = operator.execute_complete(context={}, event=event) + assert result == "/tmp/test.txt" + + def test_sftp_operator_execute_complete_raises_on_error(self): + """Test execute_complete raises AirflowException on error.""" + operator = SFTPOperator( + task_id="test_sftp_error", + ssh_conn_id="ssh_default", + local_filepath="/tmp/test.txt", + remote_filepath="/remote/test.txt", + operation=SFTPOperation.PUT, + deferrable=True, + ) + event = {"status": "error", "message": "Connection refused"} + with pytest.raises(AirflowException, match="Connection refused"): + operator.execute_complete(context={}, event=event) + + +class TestSFTPOperationTrigger: + """Tests for SFTPOperationTrigger.""" + + def test_serialize_roundtrip(self): + """Test that serialize() produces correct output for reconstruction.""" + trigger = SFTPOperationTrigger( + ssh_conn_id="ssh_default", + local_filepath="/tmp/test.txt", + remote_filepath="/remote/test.txt", + operation="put", + confirm=True, + create_intermediate_dirs=False, + remote_host=None, + concurrency=1, + prefetch=True, + ) + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.sftp.triggers.sftp.SFTPOperationTrigger" + assert kwargs["ssh_conn_id"] == "ssh_default" + assert kwargs["local_filepath"] == "/tmp/test.txt" + assert kwargs["remote_filepath"] == "/remote/test.txt" + assert kwargs["operation"] == "put" + assert kwargs["confirm"] is True + assert kwargs["remote_host"] is None + assert kwargs["concurrency"] == 1 + assert kwargs["prefetch"] is True + + def test_run_success(self): + """Test run() yields TriggerEvent with status success.""" + import asyncio + + trigger = SFTPOperationTrigger( + ssh_conn_id="ssh_default", + local_filepath="/tmp/test.txt", + remote_filepath="/remote/test.txt", + operation="put", + ) + with mock.patch("airflow.providers.sftp.triggers.sftp.SFTPHookAsync.transfer", return_value=None): + events = [] + + async def collect(): + async for event in trigger.run(): + events.append(event) + + asyncio.run(collect()) + assert len(events) == 1 + assert events[0].payload["status"] == "success" + + def test_run_error(self): + """Test run() yields TriggerEvent with status error on exception.""" + import asyncio + + trigger = SFTPOperationTrigger( + ssh_conn_id="ssh_default", + local_filepath="/tmp/test.txt", + remote_filepath="/remote/test.txt", + operation="put", + ) + with mock.patch( + "airflow.providers.sftp.triggers.sftp.SFTPHookAsync.transfer", + side_effect=Exception("Connection failed"), + ): + events = [] + + async def collect(): + async for event in trigger.run(): + events.append(event) + + asyncio.run(collect()) + assert len(events) == 1 + assert events[0].payload["status"] == "error" + assert "Connection failed" in events[0].payload["message"] \ No newline at end of file diff --git a/providers/sftp/tests/unit/sftp/test_sftp.py b/providers/sftp/tests/unit/sftp/test_sftp.py new file mode 100644 index 0000000000000..5fe2f567ea48d --- /dev/null +++ b/providers/sftp/tests/unit/sftp/test_sftp.py @@ -0,0 +1,25 @@ +# +# 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 airflow.providers.sftp.hooks.sftp import SFTPOperation + + +def test_sftp_operation_values(): + assert SFTPOperation.GET == "get" + assert SFTPOperation.PUT == "put" + assert SFTPOperation.DELETE == "delete"