Add deferrable param in SFTPSensor

airflow/providers/sftp/hooks/sftp.py

@@ -23,14 +23,20 @@
 import stat
 import warnings
 from fnmatch import fnmatch
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Sequence
+
+import asyncssh
+from asgiref.sync import sync_to_async
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.hooks.base import BaseHook
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
 if TYPE_CHECKING:
     import paramiko
 
+    from airflow.models.connection import Connection
+
 
 class SFTPHook(SSHHook):
     """Interact with SFTP.
@@ -400,3 +406,165 @@ def get_files_by_pattern(self, path, fnmatch_pattern) -> list[str]:
                 matched_files.append(file)
 
         return matched_files
+
+
+class SFTPHookAsync(BaseHook):
+    """
+    Interact with an SFTP server via asyncssh package.
+
+    :param sftp_conn_id: SFTP connection ID to be used for connecting to SFTP server
+    :param host: hostname of the SFTP server
+    :param port: port of the SFTP server
+    :param username: username used when authenticating to the SFTP server
+    :param password: password used when authenticating to the SFTP server.
+        Can be left blank if using a key file
+    :param known_hosts: path to the known_hosts file on the local file system. Defaults to ``~/.ssh/known_hosts``.
+    :param key_file: path to the client key file used for authentication to SFTP server
+    :param passphrase: passphrase used with the key_file for authentication to SFTP server
+    """
+
+    conn_name_attr = "ssh_conn_id"
+    default_conn_name = "sftp_default"
+    conn_type = "sftp"
+    hook_name = "SFTP"
+    default_known_hosts = "~/.ssh/known_hosts"
+
+    def __init__(  # nosec: B107
+        self,
+        sftp_conn_id: str = default_conn_name,
+        host: str = "",
+        port: int = 22,
+        username: str = "",
+        password: str = "",
+        known_hosts: str = default_known_hosts,
+        key_file: str = "",
+        passphrase: str = "",
+        private_key: str = "",
+    ) -> None:
+        self.sftp_conn_id = sftp_conn_id
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.known_hosts: bytes | str = os.path.expanduser(known_hosts)
+        self.key_file = key_file
+        self.passphrase = passphrase
+        self.private_key = private_key
+
+    def _parse_extras(self, conn: Connection) -> None:
+        """Parse extra fields from the connection into instance fields."""
+        extra_options = conn.extra_dejson
+        if "key_file" in extra_options and self.key_file == "":
+            self.key_file = extra_options["key_file"]
+        if "known_hosts" in extra_options and self.known_hosts != self.default_known_hosts:
+            self.known_hosts = extra_options["known_hosts"]
+        if ("passphrase" or "private_key_passphrase") in extra_options:
+            self.passphrase = extra_options["passphrase"]
+        if "private_key" in extra_options:
+            self.private_key = extra_options["private_key"]
+
+        host_key = extra_options.get("host_key")
+        no_host_key_check = extra_options.get("no_host_key_check")
+
+        if no_host_key_check is not None:
+            no_host_key_check = str(no_host_key_check).lower() == "true"
+            if host_key is not None and no_host_key_check:
+                raise ValueError("Host key check was skipped, but `host_key` value was given")
+            if no_host_key_check:
+                self.log.warning(
+                    "No Host Key Verification. This won't protect against Man-In-The-Middle attacks"
+                )
+                self.known_hosts = "none"
+
+        if host_key is not None:
+            self.known_hosts = f"{conn.host} {host_key}".encode()
+
+    async def _get_conn(self) -> asyncssh.SSHClientConnection:
+        """
+        Asynchronously connect to the SFTP server as an SSH client.
+
+        The following parameters are provided either in the extra json object in
+        the SFTP connection definition
+
+        - key_file
+        - known_hosts
+        - passphrase
+        """
+        conn = await sync_to_async(self.get_connection)(self.sftp_conn_id)
+        if conn.extra is not None:
+            self._parse_extras(conn)
+
+        conn_config = {
+            "host": conn.host,
+            "port": conn.port,
+            "username": conn.login,
+            "password": conn.password,
+        }
+        if self.key_file:
+            conn_config.update(client_keys=self.key_file)
+        if self.known_hosts:
+            if self.known_hosts.lower() == "none":
+                conn_config.update(known_hosts=None)
+            else:
+                conn_config.update(known_hosts=self.known_hosts)
+        if self.private_key:
+            _private_key = asyncssh.import_private_key(self.private_key, self.passphrase)
+            conn_config.update(client_keys=[_private_key])
+        if self.passphrase:
+            conn_config.update(passphrase=self.passphrase)
+        ssh_client_conn = await asyncssh.connect(**conn_config)
+        return ssh_client_conn
+
+    async def list_directory(self, path: str = "") -> list[str] | None:
+        """Returns a list of files on the SFTP server at the provided path."""
+        ssh_conn = await self._get_conn()
+        sftp_client = await ssh_conn.start_sftp_client()
+        try:
+            files = await sftp_client.listdir(path)
+            return sorted(files)
+        except asyncssh.SFTPNoSuchFile:
+            return None
+
+    async def read_directory(self, path: str = "") -> Sequence[asyncssh.sftp.SFTPName] | None:
+        """Returns a list of files along with their attributes on the SFTP server at the provided path."""
+        ssh_conn = await self._get_conn()
+        sftp_client = await ssh_conn.start_sftp_client()
+        try:
+            files = await sftp_client.readdir(path)
+            return files
+        except asyncssh.SFTPNoSuchFile:
+            return None
+
+    async def get_files_and_attrs_by_pattern(
+        self, path: str = "", fnmatch_pattern: str = ""
+    ) -> Sequence[asyncssh.sftp.SFTPName]:
+        """
+        Get the files along with their attributes matching the pattern (e.g. ``*.pdf``) at the provided path.
+
+        if one exists. Otherwise, raises an AirflowException to be handled upstream for deferring
+        """
+        files_list = await self.read_directory(path)
+        if files_list is None:
+            raise FileNotFoundError(f"No files at path {path!r} found...")
+        matched_files = [file for file in files_list if fnmatch(str(file.filename), fnmatch_pattern)]
+        return matched_files
+
+    async def get_mod_time(self, path: str) -> str:
+        """
+        Makes SFTP async connection.
+
+        Looks for last modified time in the specific file path and returns last modification time for
+         the file path.
+
+        :param path: full path to the remote file
+        """
+        ssh_conn = await self._get_conn()
+        sftp_client = await ssh_conn.start_sftp_client()
+        try:
+            ftp_mdtm = await sftp_client.stat(path)
+            modified_time = ftp_mdtm.mtime
+            mod_time = datetime.datetime.fromtimestamp(modified_time).strftime("%Y%m%d%H%M%S")  # type: ignore[arg-type]
+            self.log.info("Found File %s last modified: %s", str(path), str(mod_time))
+            return mod_time
+        except asyncssh.SFTPNoSuchFile:
+            raise AirflowException("No files matching")

airflow/providers/sftp/provider.yaml

@@ -61,6 +61,7 @@ dependencies:
   - apache-airflow>=2.6.0
   - apache-airflow-providers-ssh>=2.1.0
   - paramiko>=2.8.0
+  - asyncssh>=2.12.0
 
 integrations:
   - integration-name: SSH File Transfer Protocol (SFTP)
@@ -92,3 +93,8 @@ connection-types:
 task-decorators:
   - class-name: airflow.providers.sftp.decorators.sensors.sftp.sftp_sensor_task
     name: sftp_sensor
+
+triggers:
+  - integration-name: SSH File Transfer Protocol (SFTP)
+    python-modules:
+      - airflow.providers.sftp.triggers.sftp

airflow/providers/sftp/sensors/sftp.py

@@ -19,13 +19,15 @@
 from __future__ import annotations
 
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
 from paramiko.sftp import SFTP_NO_SUCH_FILE
 
-from airflow.exceptions import AirflowSkipException
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.sftp.hooks.sftp import SFTPHook
+from airflow.providers.sftp.triggers.sftp import SFTPTrigger
 from airflow.sensors.base import BaseSensorOperator, PokeReturnValue
 from airflow.utils.timezone import convert_to_utc
 
@@ -41,6 +43,7 @@ class SFTPSensor(BaseSensorOperator):
     :param file_pattern: The pattern that will be used to match the file (fnmatch format)
     :param sftp_conn_id: The connection to run the sensor against
     :param newer_than: DateTime for which the file or file path should be newer than, comparison is inclusive
+    :param deferrable: If waiting for completion, whether to defer the task until done, default is ``False``.
     """
 
     template_fields: Sequence[str] = (
@@ -58,6 +61,7 @@ def __init__(
         python_callable: Callable | None = None,
         op_args: list | None = None,
         op_kwargs: dict[str, Any] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -69,6 +73,7 @@ def __init__(
         self.python_callable: Callable | None = python_callable
         self.op_args = op_args or []
         self.op_kwargs = op_kwargs or {}
+        self.deferrable = deferrable
 
     def poke(self, context: Context) -> PokeReturnValue | bool:
         self.hook = SFTPHook(self.sftp_conn_id)
@@ -119,3 +124,49 @@ def poke(self, context: Context) -> PokeReturnValue | bool:
                 xcom_value={"files_found": files_found, "decorator_return_value": callable_return},
             )
         return True
+
+    def execute(self, context: Context) -> Any:
+        # Unlike other async sensors, we do not follow the pattern of calling the synchronous self.poke()
+        # method before deferring here. This is due to the current limitations we have in the synchronous
+        # SFTPHook methods. They are as follows:
+        #
+        # For file_pattern sensing, the hook implements list_directory() method which returns a list of
+        # filenames only without the attributes like modified time which is required for the file_pattern
+        # sensing when newer_than is supplied. This leads to intermittent failures potentially due to
+        # throttling by the SFTP server as the hook makes multiple calls to the server to get the
+        # attributes for each of the files in the directory.This limitation is resolved here by instead
+        # calling the read_directory() method which returns a list of files along with their attributes
+        # in a single call. We can add back the call to self.poke() before deferring once the above
+        # limitations are resolved in the sync sensor.
+        if self.deferrable:
+            self.defer(
+                timeout=timedelta(seconds=self.timeout),
+                trigger=SFTPTrigger(
+                    path=self.path,
+                    file_pattern=self.file_pattern,
+                    sftp_conn_id=self.sftp_conn_id,
+                    poke_interval=self.poke_interval,
+                    newer_than=self.newer_than,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            return super().execute(context=context)
+
+    def execute_complete(self, context: dict[str, Any], event: Any = None) -> None:
+        """
+        Callback for when the trigger fires - returns immediately.
+
+        Relies on trigger to throw an exception, otherwise it assumes execution was
+        successful.
+        """
+        if event is not None:
+            if "status" in event and event["status"] == "error":
+                raise AirflowException(event["message"])
+
+            if "status" in event and event["status"] == "success":
+                self.log.info("%s completed successfully.", self.task_id)
+                self.log.info(event["message"])
+                return None
+
+        raise AirflowException("No event received in trigger callback")

airflow/providers/sftp/triggers/__init__.py

@@ -0,0 +1,16 @@
+# 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.