Add deferrable mode for S3KeySensor

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

@@ -23,7 +23,6 @@
 import gzip as gz
 import io
 import logging
-import os
 import re
 import shutil
 import warnings
@@ -612,119 +611,6 @@ async def _list_keys_async(
 
         return keys
 
-    async def is_keys_unchanged_async(
-        self,
-        client: AioBaseClient,
-        bucket_name: str,
-        prefix: str,
-        inactivity_period: float = 60 * 60,
-        min_objects: int = 1,
-        previous_objects: set[str] | None = None,
-        inactivity_seconds: int = 0,
-        allow_delete: bool = True,
-        last_activity_time: datetime | None = None,
-    ) -> dict[str, Any]:
-        """
-        Checks whether new objects have been uploaded and the inactivity_period
-        has passed and updates the state of the sensor accordingly.
-
-        :param client: aiobotocore client
-        :param bucket_name: the name of the bucket
-        :param prefix: a key prefix
-        :param inactivity_period:  the total seconds of inactivity to designate
-            keys unchanged. Note, this mechanism is not real time and
-            this operator may not return until a poke_interval after this period
-            has passed with no additional objects sensed.
-        :param min_objects: the minimum number of objects needed for keys unchanged
-            sensor to be considered valid.
-        :param previous_objects: the set of object ids found during the last poke.
-        :param inactivity_seconds: number of inactive seconds
-        :param last_activity_time: last activity datetime
-        :param allow_delete: Should this sensor consider objects being deleted
-            between pokes valid behavior. If true a warning message will be logged
-            when this happens. If false an error will be raised.
-        :return: dictionary with status and message
-        :rtype: Dict
-        """
-        if previous_objects is None:
-            previous_objects = set()
-        list_keys = await self._list_keys_async(client=client, bucket_name=bucket_name, prefix=prefix)
-        current_objects = set(list_keys)
-        current_num_objects = len(current_objects)
-        if current_num_objects > len(previous_objects):
-            # When new objects arrived, reset the inactivity_seconds
-            # and update previous_objects for the next poke.
-            self.log.info(
-                "New objects found at %s, resetting last_activity_time.",
-                os.path.join(bucket_name, prefix),
-            )
-            self.log.debug("New objects: %s", current_objects - previous_objects)
-            last_activity_time = datetime.now()
-            inactivity_seconds = 0
-            previous_objects = current_objects
-            return {
-                "status": "pending",
-                "previous_objects": previous_objects,
-                "last_activity_time": last_activity_time,
-                "inactivity_seconds": inactivity_seconds,
-            }
-
-        if len(previous_objects) - len(current_objects):
-            # During the last poke interval objects were deleted.
-            if allow_delete:
-                deleted_objects = previous_objects - current_objects
-                previous_objects = current_objects
-                last_activity_time = datetime.now()
-                self.log.info(
-                    "Objects were deleted during the last poke interval. Updating the "
-                    "file counter and resetting last_activity_time:\n%s",
-                    deleted_objects,
-                )
-                return {
-                    "status": "pending",
-                    "previous_objects": previous_objects,
-                    "last_activity_time": last_activity_time,
-                    "inactivity_seconds": inactivity_seconds,
-                }
-
-            return {
-                "status": "error",
-                "message": f" {os.path.join(bucket_name, prefix)} between pokes.",
-            }
-
-        if last_activity_time:
-            inactivity_seconds = int((datetime.now() - last_activity_time).total_seconds())
-        else:
-            # Handles the first poke where last inactivity time is None.
-            last_activity_time = datetime.now()
-            inactivity_seconds = 0
-
-        if inactivity_seconds >= inactivity_period:
-            path = os.path.join(bucket_name, prefix)
-
-            if current_num_objects >= min_objects:
-                success_message = (
-                    "SUCCESS: Sensor found %s objects at %s. "
-                    "Waited at least %s seconds, with no new objects uploaded."
-                )
-                self.log.info(success_message, current_num_objects, path, inactivity_period)
-                return {
-                    "status": "success",
-                    "message": success_message % (current_num_objects, path, inactivity_period),
-                }
-
-            self.log.error("FAILURE: Inactivity Period passed, not enough objects found in %s", path)
-            return {
-                "status": "error",
-                "message": f"FAILURE: Inactivity Period passed, not enough objects found in {path}",
-            }
-        return {
-            "status": "pending",
-            "previous_objects": previous_objects,
-            "last_activity_time": last_activity_time,
-            "inactivity_seconds": inactivity_seconds,
-        }
-
     def _list_key_object_filter(
         self, keys: list, from_datetime: datetime | None = None, to_datetime: datetime | None = None
     ) -> list:

tests/providers/amazon/aws/hooks/test_s3.py

@@ -231,6 +231,12 @@ def test_list_keys(self, s3_bucket):
         bucket.put_object(Key="bxb", Body=b"axb")
         bucket.put_object(Key="dir/b", Body=b"b")
 
+        from_datetime = datetime(1992, 3, 8, 18, 52, 51)
+        to_datetime = datetime(1993, 3, 14, 21, 52, 42)
+
+        def dummy_object_filter(keys, from_datetime=None, to_datetime=None):
+            return []
+
         assert [] == hook.list_keys(s3_bucket, prefix="non-existent/")
         assert ["a", "ba", "bxa", "bxb", "dir/b"] == hook.list_keys(s3_bucket)
         assert ["a", "ba", "bxa", "bxb"] == hook.list_keys(s3_bucket, delimiter="/")
@@ -551,131 +557,6 @@ async def test_s3_key_hook_list_keys_async(self, mock_client):
         response = await s3_hook_async._list_keys_async(mock_client, "test_bucket", "test*")
         assert response == ["test_key", "test_key2"]
 
-    @pytest.mark.asyncio
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook.async_conn")
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook._list_keys_async")
-    async def test_s3_key_hook_is_keys_unchanged_false_async(self, mock_list_keys, mock_client):
-        """
-        Test is_key_unchanged gives False response
-        :return:
-        """
-
-        mock_list_keys.return_value = ["test"]
-
-        s3_hook_async = S3Hook(client_type="S3", resource_type="S3")
-        response = await s3_hook_async.is_keys_unchanged_async(
-            client=mock_client.return_value,
-            bucket_name="test_bucket",
-            prefix="test",
-            inactivity_period=1,
-            min_objects=1,
-            previous_objects=set(),
-            inactivity_seconds=0,
-            allow_delete=True,
-            last_activity_time=None,
-        )
-
-        assert response.get("status") == "pending"
-
-        # test for the case when current_objects < previous_objects
-        mock_list_keys.return_value = []
-
-        s3_hook_async = S3Hook(client_type="S3", resource_type="S3")
-        response = await s3_hook_async.is_keys_unchanged_async(
-            client=mock_client.return_value,
-            bucket_name="test_bucket",
-            prefix="test",
-            inactivity_period=1,
-            min_objects=1,
-            previous_objects=set("test"),
-            inactivity_seconds=0,
-            allow_delete=True,
-            last_activity_time=None,
-        )
-
-        assert response.get("status") == "pending"
-
-    @pytest.mark.asyncio
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook.async_conn")
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook._list_keys_async")
-    async def test_s3_key_hook_is_keys_unchanged_exception_async(self, mock_list_keys, mock_client):
-        """
-        Test is_key_unchanged gives AirflowException
-        :return:
-        """
-        mock_list_keys.return_value = []
-
-        s3_hook_async = S3Hook(client_type="S3", resource_type="S3")
-
-        response = await s3_hook_async.is_keys_unchanged_async(
-            client=mock_client.return_value,
-            bucket_name="test_bucket",
-            prefix="test",
-            inactivity_period=1,
-            min_objects=1,
-            previous_objects=set("test"),
-            inactivity_seconds=0,
-            allow_delete=False,
-            last_activity_time=None,
-        )
-
-        assert response == {"message": " test_bucket/test between pokes.", "status": "error"}
-
-    @pytest.mark.asyncio
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook.async_conn")
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook._list_keys_async")
-    async def test_s3_key_hook_is_keys_unchanged_pending_async(self, mock_list_keys, mock_client):
-        """
-        Test is_key_unchanged gives AirflowException
-        :return:
-        """
-        mock_list_keys.return_value = []
-
-        s3_hook_async = S3Hook(client_type="S3", resource_type="S3")
-
-        response = await s3_hook_async.is_keys_unchanged_async(
-            client=mock_client.return_value,
-            bucket_name="test_bucket",
-            prefix="test",
-            inactivity_period=1,
-            min_objects=0,
-            previous_objects=set(),
-            inactivity_seconds=0,
-            allow_delete=False,
-            last_activity_time=None,
-        )
-
-        assert response.get("status") == "pending"
-
-    @pytest.mark.asyncio
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook.async_conn")
-    @async_mock.patch("airflow.providers.amazon.aws.triggers.s3.S3Hook._list_keys_async")
-    async def test_s3_key_hook_is_keys_unchanged_inactivity_error_async(self, mock_list_keys, mock_client):
-        """
-        Test is_key_unchanged gives AirflowException
-        :return:
-        """
-        mock_list_keys.return_value = []
-
-        s3_hook_async = S3Hook(client_type="S3", resource_type="S3")
-
-        response = await s3_hook_async.is_keys_unchanged_async(
-            client=mock_client.return_value,
-            bucket_name="test_bucket",
-            prefix="test",
-            inactivity_period=0,
-            min_objects=5,
-            previous_objects=set(),
-            inactivity_seconds=5,
-            allow_delete=False,
-            last_activity_time=None,
-        )
-
-        assert response == {
-            "status": "error",
-            "message": "FAILURE: Inactivity Period passed, not enough objects found in test_bucket/test",
-        }
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         "test_first_prefix, test_second_prefix",