Add deferrable mode for S3KeySensor

airflow/providers/amazon/aws/sensors/s3.py

@@ -20,8 +20,8 @@
 import fnmatch
 import os
 import re
-from datetime import datetime
-from typing import TYPE_CHECKING, Callable, Sequence
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Any, Callable, List, Sequence, cast
 
 from deprecated import deprecated
 
@@ -31,6 +31,7 @@
 from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.aws.triggers.s3 import S3KeyTrigger
 from airflow.sensors.base import BaseSensorOperator, poke_mode_only
 
 
@@ -48,7 +49,7 @@ class S3KeySensor(BaseSensorOperator):
         or relative path from root level. When it's specified as a full s3://
         url, please leave bucket_name as `None`
     :param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
-        is not provided as a full s3:// url. When specified, all the keys passed to ``bucket_key``
+        is not provided as a full ``s3://`` url. When specified, all the keys passed to ``bucket_key``
         refers to this bucket
     :param wildcard_match: whether the bucket_key should be interpreted as a
         Unix wildcard pattern
@@ -61,8 +62,9 @@ class S3KeySensor(BaseSensorOperator):
             def check_fn(files: List) -> bool:
                 return any(f.get('Size', 0) > 1048576 for f in files)
     :param aws_conn_id: a reference to the s3 connection
-    :param verify: Whether or not to verify SSL certificates for S3 connection.
-        By default SSL certificates are verified.
+    :param deferrable: Run operator in the deferrable mode
+    :param verify: Whether to verify SSL certificates for S3 connection.
+        By default, SSL certificates are verified.
         You can provide the following values:
 
         - ``False``: do not validate SSL certificates. SSL will still be used
@@ -84,6 +86,7 @@ def __init__(
         check_fn: Callable[..., bool] | None = None,
         aws_conn_id: str = "aws_default",
         verify: str | bool | None = None,
+        deferrable: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -93,6 +96,7 @@ def __init__(
         self.check_fn = check_fn
         self.aws_conn_id = aws_conn_id
         self.verify = verify
+        self.deferrable = deferrable
 
     def _check_key(self, key):
         bucket_name, key = S3Hook.get_s3_bucket_key(self.bucket_name, key, "bucket_name", "bucket_key")
@@ -131,6 +135,43 @@ def poke(self, context: Context):
         else:
             return all(self._check_key(key) for key in self.bucket_key)
 
+    def execute(self, context: Context) -> None:
+        """
+        Defers to Trigger class to poll for state of the job run until
+        it reaches a failure state or success state
+        """
+        if not self.deferrable:
+            super().execute(context)
+        else:
+            if not self.poke(context=context):
+                self.defer(
+                    timeout=timedelta(seconds=self.timeout),
+                    trigger=S3KeyTrigger(
+                        bucket_name=cast(str, self.bucket_name),
+                        bucket_key=self.bucket_key,
+                        wildcard_match=self.wildcard_match,
+                        check_fn=self.check_fn,
+                        aws_conn_id=self.aws_conn_id,
+                        verify=self.verify,
+                        poke_interval=self.poke_interval,
+                    ),
+                    method_name="execute_complete",
+                )
+
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> bool | None:
+        """
+        Callback for when the trigger fires - returns immediately.
+        Relies on trigger to throw an exception, otherwise it assumes execution was
+        successful.
+        """
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+        elif event["status"] == "success" and "s3_objects" in event:
+            files = cast(List[str], event["s3_objects"])
+            if self.check_fn:
+                return self.check_fn(files)
+        return None
+
     @deprecated(reason="use `hook` property instead.")
     def get_hook(self) -> S3Hook:
         """Create and return an S3Hook"""

airflow/providers/amazon/aws/triggers/s3.py

@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+from typing import Any, AsyncIterator, Callable
+
+from airflow.providers.amazon.aws.hooks.s3 import S3AsyncHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class S3KeyTrigger(BaseTrigger):
+    """
+    S3KeyTrigger is fired as deferred class with params to run the task in trigger worker
+
+    :param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
+        is not provided as a full s3:// url.
+    :param bucket_key:  The key being waited on. Supports full s3:// style url
+        or relative path from root level. When it's specified as a full s3://
+        url, please leave bucket_name as `None`.
+    :param wildcard_match: whether the bucket_key should be interpreted as a
+        Unix wildcard pattern
+    :param aws_conn_id: reference to the s3 connection
+    :param hook_params: params for hook its optional
+    :param check_fn: Function that receives the list of the S3 objects,
+        and returns a boolean
+    """
+
+    def __init__(
+        self,
+        bucket_name: str,
+        bucket_key: str | list[str],
+        wildcard_match: bool = False,
+        check_fn: Callable[..., bool] | None = None,
+        aws_conn_id: str = "aws_default",
+        poke_interval: float = 5.0,
+        **hook_params: Any,
+    ):
+        super().__init__()
+        self.bucket_name = bucket_name
+        self.bucket_key = bucket_key
+        self.wildcard_match = wildcard_match
+        self.check_fn = check_fn
+        self.aws_conn_id = aws_conn_id
+        self.hook_params = hook_params
+        self.poke_interval = poke_interval
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize S3KeyTrigger arguments and classpath."""
+        return (
+            "airflow.providers.amazon.aws.triggers.s3.S3KeyTrigger",
+            {
+                "bucket_name": self.bucket_name,
+                "bucket_key": self.bucket_key,
+                "wildcard_match": self.wildcard_match,
+                "check_fn": self.check_fn,
+                "aws_conn_id": self.aws_conn_id,
+                "hook_params": self.hook_params,
+                "poke_interval": self.poke_interval,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Make an asynchronous connection using S3HookAsync."""
+        try:
+            hook = self._get_async_hook()
+            async with await hook.get_client_async() as client:
+                while True:
+                    if await hook.check_key(client, self.bucket_name, self.bucket_key, self.wildcard_match):
+                        if self.check_fn is None:
+                            yield TriggerEvent({"status": "success"})
+                        else:
+                            s3_objects = await hook.get_files(
+                                client, self.bucket_name, self.bucket_key, self.wildcard_match
+                            )
+                            yield TriggerEvent({"status": "success", "s3_objects": s3_objects})
+                    await asyncio.sleep(self.poke_interval)
+
+        except Exception as e:
+            yield TriggerEvent({"status": "error", "message": str(e)})
+
+    def _get_async_hook(self) -> S3AsyncHook:
+        return S3AsyncHook(aws_conn_id=self.aws_conn_id, verify=self.hook_params.get("verify"))

airflow/providers/amazon/provider.yaml

@@ -72,9 +72,10 @@ dependencies:
   - mypy-boto3-rds>=1.24.0
   - mypy-boto3-redshift-data>=1.24.0
   - mypy-boto3-appflow>=1.24.0
+  - aiohttp
+  - asgiref
   - mypy-boto3-s3>=1.24.0
 
-
 integrations:
   - integration-name: Amazon Athena
     external-doc-url: https://aws.amazon.com/athena/

docs/apache-airflow-providers-amazon/operators/s3/s3.rst

@@ -248,6 +248,26 @@ multiple files can match one key. The list of matched S3 object attributes conta
     :start-after: [START howto_sensor_s3_key_function]
     :end-before: [END howto_sensor_s3_key_function]
 
+You can also run this operator in deferrable mode by setting the parameter ``deferrable`` to True.
+This will lead to efficient utilization of Airflow workers as polling for job status happens on
+the triggerer asynchronously. Note that this will need triggerer to be available on your Airflow deployment.
+
+To check one file:
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_s3.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_s3_key_single_key_deferrable]
+    :end-before: [END howto_sensor_s3_key_single_key_deferrable]
+
+To check multiple files:
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_s3.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_s3_key_multiple_keys_deferrable]
+    :end-before: [END howto_sensor_s3_key_multiple_keys_deferrable]
+
 .. _howto/sensor:S3KeysUnchangedSensor:
 
 Wait on Amazon S3 prefix changes

generated/provider_dependencies.json

@@ -17,9 +17,11 @@
   },
   "amazon": {
     "deps": [
+      "aiohttp",
       "apache-airflow-providers-common-sql>=1.3.1",
       "apache-airflow>=2.4.0",
       "asgiref",
+      "asgiref",
       "boto3>=1.24.0",
       "jsonpath_ng>=1.5.3",
       "mypy-boto3-appflow>=1.24.0",

tests/providers/amazon/aws/compat.py

@@ -0,0 +1,37 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+__all__ = ["async_mock", "AsyncMock"]
+
+import sys
+
+if sys.version_info < (3, 8):
+    # For compatibility with Python 3.7
+    from asynctest import mock as async_mock
+
+    # ``asynctest.mock.CoroutineMock`` which provide compatibility not working well with autospec=True
+    # as result "TypeError: object MagicMock can't be used in 'await' expression" could be raised.
+    # Best solution in this case provide as spec actual awaitable object
+    # >>> from tests.providers.apache.livy.compat import AsyncMock
+    # >>> from foo.bar import SpamEgg
+    # >>> mock_something = AsyncMock(SpamEgg)
+    from asynctest.mock import CoroutineMock as AsyncMock
+else:
+    from unittest import mock as async_mock
+    from unittest.mock import AsyncMock