Add deferrable GCSObjectExistenceSensorAsync (apache#28763)

This PR donates the following developed GCSObjectExistenceSensorAsync` in [astronomer-providers](https://github.com/astronomer/astronomer-providers) repo to apache airflow.

`GCSObjectExistenceSensorAsync`

airflow/providers/google/cloud/hooks/gcs.py

@@ -32,18 +32,21 @@
 from typing import IO, Callable, Generator, Sequence, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
+from aiohttp import ClientSession
+from gcloud.aio.storage import Storage
 from google.api_core.exceptions import NotFound
 from google.api_core.retry import Retry
 
 # not sure why but mypy complains on missing `storage` but it is clearly there and is importable
 from google.cloud import storage  # type: ignore[attr-defined]
 from google.cloud.exceptions import GoogleCloudError
 from google.cloud.storage.retry import DEFAULT_RETRY
+from requests import Session
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.utils import timezone
 from airflow.version import version
 
@@ -1174,3 +1177,14 @@ def _parse_gcs_url(gsurl: str) -> tuple[str, str]:
     # Remove leading '/' but NOT trailing one
     blob = parsed_url.path.lstrip("/")
     return bucket, blob
+
+
+class GCSAsyncHook(GoogleBaseAsyncHook):
+    """GCSAsyncHook run on the trigger worker, inherits from GoogleBaseHookAsync"""
+
+    sync_hook_class = GCSHook
+
+    async def get_storage_client(self, session: ClientSession) -> Storage:
+        """Returns a Google Cloud Storage service object."""
+        with await self.service_file_as_context() as file:
+            return Storage(service_file=file, session=cast(Session, session))

airflow/providers/google/cloud/sensors/gcs.py

@@ -20,14 +20,15 @@
 
 import os
 import textwrap
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Callable, Sequence
 
 from google.api_core.retry import Retry
 from google.cloud.storage.retry import DEFAULT_RETRY
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.cloud.triggers.gcs import GCSBlobTrigger
 from airflow.sensors.base import BaseSensorOperator, poke_mode_only
 
 if TYPE_CHECKING:
@@ -94,6 +95,55 @@ def poke(self, context: Context) -> bool:
         return hook.exists(self.bucket, self.object, self.retry)
 
 
+class GCSObjectExistenceAsyncSensor(GCSObjectExistenceSensor):
+    """
+    Checks for the existence of a file in Google Cloud Storage .
+
+    :param bucket: The Google Cloud Storage bucket where the object is.
+    :param object: The name of the object to check in the Google cloud storage bucket.
+    :param google_cloud_conn_id: The connection ID to use when connecting to Google Cloud Storage.
+    :param delegate_to: The account to impersonate using domain-wide delegation of authority,
+        if any. For this to work, the service account making the request must have
+        domain-wide delegation enabled.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def execute(self, context: Context) -> None:
+        """Airflow runs this method on the worker and defers using the trigger."""
+        self.defer(
+            timeout=timedelta(seconds=self.timeout),
+            trigger=GCSBlobTrigger(
+                bucket=self.bucket,
+                object_name=self.object,
+                poke_interval=self.poke_interval,
+                google_cloud_conn_id=self.google_cloud_conn_id,
+                hook_params={
+                    "delegate_to": self.delegate_to,
+                    "impersonation_chain": self.impersonation_chain,
+                },
+            ),
+            method_name="execute_complete",
+        )
+
+    def execute_complete(self, context: Context, event: dict[str, str]) -> str:
+        """
+        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"])
+        self.log.info("File %s was found in bucket %s.", self.object, self.bucket)
+        return event["message"]
+
+
 def ts_function(context):
     """
     Default callback for the GoogleCloudStorageObjectUpdatedSensor. The default

airflow/providers/google/cloud/triggers/gcs.py

@@ -0,0 +1,99 @@
+#
+# 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
+
+from aiohttp import ClientSession
+
+from airflow.providers.google.cloud.hooks.gcs import GCSAsyncHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class GCSBlobTrigger(BaseTrigger):
+    """
+    A trigger that fires and it finds the requested file or folder present in the given bucket.
+
+    :param bucket: the bucket in the google cloud storage where the objects are residing.
+    :param object_name: the file or folder present in the bucket
+    :param google_cloud_conn_id: reference to the Google Connection
+    :param poke_interval: polling period in seconds to check for file/folder
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        object_name: str,
+        poke_interval: float,
+        google_cloud_conn_id: str,
+        hook_params: dict[str, Any],
+    ):
+        super().__init__()
+        self.bucket = bucket
+        self.object_name = object_name
+        self.poke_interval = poke_interval
+        self.google_cloud_conn_id: str = google_cloud_conn_id
+        self.hook_params = hook_params
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serializes GCSBlobTrigger arguments and classpath."""
+        return (
+            "airflow.providers.google.cloud.triggers.gcs.GCSBlobTrigger",
+            {
+                "bucket": self.bucket,
+                "object_name": self.object_name,
+                "poke_interval": self.poke_interval,
+                "google_cloud_conn_id": self.google_cloud_conn_id,
+                "hook_params": self.hook_params,
+            },
+        )
+
+    async def run(self) -> AsyncIterator["TriggerEvent"]:
+        """Simple loop until the relevant file/folder is found."""
+        try:
+            hook = self._get_async_hook()
+            while True:
+                res = await self._object_exists(
+                    hook=hook, bucket_name=self.bucket, object_name=self.object_name
+                )
+                if res == "success":
+                    yield TriggerEvent({"status": "success", "message": res})
+                await asyncio.sleep(self.poke_interval)
+        except Exception as e:
+            yield TriggerEvent({"status": "error", "message": str(e)})
+            return
+
+    def _get_async_hook(self) -> GCSAsyncHook:
+        return GCSAsyncHook(gcp_conn_id=self.google_cloud_conn_id, **self.hook_params)
+
+    async def _object_exists(self, hook: GCSAsyncHook, bucket_name: str, object_name: str) -> str:
+        """
+        Checks for the existence of a file in Google Cloud Storage.
+
+        :param bucket_name: The Google Cloud Storage bucket where the object is.
+        :param object_name: The name of the blob_name to check in the Google cloud
+            storage bucket.
+        """
+        async with ClientSession() as s:
+            client = await hook.get_storage_client(s)
+            bucket = client.get_bucket(bucket_name)
+            object_response = await bucket.blob_exists(blob_name=object_name)
+            if object_response:
+                return "success"
+            return "pending"

docs/apache-airflow-providers-google/operators/cloud/gcs.rst

@@ -188,6 +188,22 @@ Use the :class:`~airflow.providers.google.cloud.sensors.gcs.GCSObjectExistenceSe
     :start-after: [START howto_sensor_object_exists_task]
     :end-before: [END howto_sensor_object_exists_task]
 
+
+.. _howto/sensor:GCSObjectExistenceAsyncSensor:
+
+GCSObjectExistenceAsyncSensor
+-----------------------------
+
+Use the :class:`~airflow.providers.google.cloud.sensors.gcs.GCSObjectExistenceAsyncSensor`
+(deferrable version) if you would like to free up the worker slots while the sensor is running.
+
+.. exampleinclude:: /../../tests/system/providers/google/cloud/gcs/example_gcs_sensor.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_object_exists_task_async]
+    :end-before: [END howto_sensor_object_exists_task_async]
+
+
 .. _howto/sensor:GCSObjectsWithPrefixExistenceSensor:
 
 GCSObjectsWithPrefixExistenceSensor

tests/providers/google/cloud/sensors/test_gcs.py

@@ -24,15 +24,17 @@
 import pytest
 from google.cloud.storage.retry import DEFAULT_RETRY
 
-from airflow.exceptions import AirflowSensorTimeout
+from airflow.exceptions import AirflowSensorTimeout, TaskDeferred
 from airflow.models.dag import DAG, AirflowException
 from airflow.providers.google.cloud.sensors.gcs import (
+    GCSObjectExistenceAsyncSensor,
     GCSObjectExistenceSensor,
     GCSObjectsWithPrefixExistenceSensor,
     GCSObjectUpdateSensor,
     GCSUploadSessionCompleteSensor,
     ts_function,
 )
+from airflow.providers.google.cloud.triggers.gcs import GCSBlobTrigger
 
 TEST_BUCKET = "TEST_BUCKET"
 
@@ -53,6 +55,15 @@
 MOCK_DATE_ARRAY = [datetime(2019, 2, 24, 12, 0, 0) - i * timedelta(seconds=10) for i in range(25)]
 
 
+@pytest.fixture()
+def context():
+    """
+    Creates an empty context.
+    """
+    context = {"data_interval_end": datetime.utcnow()}
+    yield context
+
+
 def next_time_side_effect():
     """
     This each time this is called mock a time 10 seconds later
@@ -88,6 +99,46 @@ def test_should_pass_argument_to_hook(self, mock_hook):
         mock_hook.return_value.exists.assert_called_once_with(TEST_BUCKET, TEST_OBJECT, DEFAULT_RETRY)
 
 
+class TestGoogleCloudStorageObjectSensorAsync(TestCase):
+    def test_gcs_object_existence_sensor_async(self):
+        """
+        Asserts that a task is deferred and a GCSBlobTrigger will be fired
+        when the GCSObjectExistenceAsyncSensor is executed.
+        """
+        task = GCSObjectExistenceAsyncSensor(
+            task_id="task-id",
+            bucket=TEST_BUCKET,
+            object=TEST_OBJECT,
+            google_cloud_conn_id=TEST_GCP_CONN_ID,
+        )
+        with pytest.raises(TaskDeferred) as exc:
+            task.execute(context)
+        assert isinstance(exc.value.trigger, GCSBlobTrigger), "Trigger is not a GCSBlobTrigger"
+
+    def test_gcs_object_existence_sensor_async_execute_failure(self):
+        """Tests that an AirflowException is raised in case of error event"""
+        task = GCSObjectExistenceAsyncSensor(
+            task_id="task-id",
+            bucket=TEST_BUCKET,
+            object=TEST_OBJECT,
+            google_cloud_conn_id=TEST_GCP_CONN_ID,
+        )
+        with pytest.raises(AirflowException):
+            task.execute_complete(context=None, event={"status": "error", "message": "test failure message"})
+
+    def test_gcs_object_existence_sensor_async_execute_complete(self):
+        """Asserts that logging occurs as expected"""
+        task = GCSObjectExistenceAsyncSensor(
+            task_id="task-id",
+            bucket=TEST_BUCKET,
+            object=TEST_OBJECT,
+            google_cloud_conn_id=TEST_GCP_CONN_ID,
+        )
+        with mock.patch.object(task.log, "info") as mock_log_info:
+            task.execute_complete(context=None, event={"status": "success", "message": "Job completed"})
+        mock_log_info.assert_called_with("File %s was found in bucket %s.", TEST_OBJECT, TEST_BUCKET)
+
+
 class TestTsFunction(TestCase):
     def test_should_support_datetime(self):
         context = {