Google Cloud Tasks Sensor for queue being empty (#25622)

airflow/providers/google/cloud/example_dags/example_cloud_task.py

@@ -0,0 +1,53 @@
+#
+# 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.
+
+"""
+Example Airflow DAG that sense a cloud task queue being empty.
+
+This DAG relies on the following OS environment variables
+
+* GCP_PROJECT_ID - Google Cloud project where the Compute Engine instance exists.
+* GCP_ZONE - Google Cloud zone where the cloud task queue exists.
+* QUEUE_NAME - Name of the cloud task queue.
+"""
+
+import os
+from datetime import datetime
+
+from airflow import models
+from airflow.providers.google.cloud.sensors.tasks import TaskQueueEmptySensor
+
+GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'example-project')
+GCP_ZONE = os.environ.get('GCE_ZONE', 'europe-west1-b')
+QUEUE_NAME = os.environ.get('GCP_QUEUE_NAME', 'testqueue')
+
+
+with models.DAG(
+    'example_gcp_cloud_tasks_sensor',
+    start_date=datetime(2022, 8, 8),
+    catchup=False,
+    tags=['example'],
+) as dag:
+    # [START cloud_tasks_empty_sensor]
+    gcp_cloud_tasks_sensor = TaskQueueEmptySensor(
+        project_id=GCP_PROJECT_ID,
+        location=GCP_ZONE,
+        task_id='gcp_sense_cloud_tasks_empty',
+        queue_name=QUEUE_NAME,
+    )
+    # [END cloud_tasks_empty_sensor]

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

@@ -0,0 +1,87 @@
+#
+# 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.
+"""This module contains a Google Cloud Task sensor."""
+from typing import TYPE_CHECKING, Optional, Sequence, Union
+
+from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class TaskQueueEmptySensor(BaseSensorOperator):
+    """
+    Pulls tasks count from a cloud task queue.
+    Always waits for queue returning tasks count as 0.
+
+    :param project_id: the Google Cloud project ID for the subscription (templated)
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param queue_name: The queue name to for which task empty sensing is required.
+    :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).
+    """
+
+    template_fields: Sequence[str] = (
+        "project_id",
+        "location",
+        "queue_name",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        *,
+        location: str,
+        project_id: Optional[str] = None,
+        queue_name: Optional[str] = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.location = location
+        self.project_id = project_id
+        self.queue_name = queue_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def poke(self, context: "Context") -> bool:
+
+        hook = CloudTasksHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        # TODO uncomment page_size once https://issuetracker.google.com/issues/155978649?pli=1 gets fixed
+        tasks = hook.list_tasks(
+            location=self.location,
+            queue_name=self.queue_name,
+            # page_size=1
+        )
+
+        self.log.info("tasks exhausted in cloud task queue?: %s" % (len(tasks) == 0))
+
+        return len(tasks) == 0

airflow/providers/google/provider.yaml

@@ -633,6 +633,9 @@ sensors:
   - integration-name: Google Cloud Dataform
     python-modules:
       - airflow.providers.google.cloud.sensors.dataform
+  - integration-name: Google Cloud Tasks
+    python-modules:
+      - airflow.providers.google.cloud.sensors.tasks
 
 hooks:
   - integration-name: Google Ads

docs/apache-airflow-providers-google/index.rst

@@ -30,6 +30,7 @@ Content
     Secrets backends <secrets-backends/google-cloud-secret-manager-backend>
     API Authentication backend <api-auth-backend/google-openid>
     Operators <operators/index>
+    Sensors <sensors/index>
 
 .. toctree::
     :maxdepth: 1

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

@@ -17,9 +17,9 @@
 
 Google Cloud Tasks
 ==================
-
-Firestore in Datastore mode is a NoSQL document database built for automatic scaling,
-high performance, and ease of application development.
+Cloud Tasks is a fully managed service that allows you to manage the execution, dispatch,
+and delivery of a large number of distributed tasks.
+Using Cloud Tasks, you can perform work asynchronously outside of a user or service-to-service request.
 
 For more information about the service visit
 `Cloud Tasks product documentation <https://cloud.google.com/tasks/docs>`__

docs/apache-airflow-providers-google/sensors/google-cloud-tasks.rst

@@ -0,0 +1,39 @@
+ .. 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.
+
+.. _google_cloud_tasks_empty_sensor:
+
+Google Cloud Tasks
+==================
+Cloud Tasks is a fully managed service that allows you to manage the execution, dispatch,
+and delivery of a large number of distributed tasks.
+Using Cloud Tasks, you can perform work asynchronously outside of a user or service-to-service request.
+
+For more information about the service visit
+`Cloud Tasks product documentation <https://cloud.google.com/tasks/docs>`__
+
+Google Cloud Tasks Empty Sensor
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To sense Queue being empty use
+:class:`~airflow.providers.google.cloud.sensor.tasks.TaskQueueEmptySensor`
+
+.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_cloud_task.py
+    :language: python
+    :dedent: 4
+    :start-after: [START cloud_tasks_empty_sensor]
+    :end-before: [END cloud_tasks_empty_sensor]

docs/apache-airflow-providers-google/sensors/index.rst

@@ -0,0 +1,26 @@
+ .. 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.
+
+
+
+Google Sensors
+================
+
+.. toctree::
+    :maxdepth: 1
+
+    google-cloud-tasks

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

@@ -0,0 +1,58 @@
+#
+# 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.
+import unittest
+from typing import Any, Dict
+from unittest import mock
+
+from google.cloud.tasks_v2.types import Task
+
+from airflow.providers.google.cloud.sensors.tasks import TaskQueueEmptySensor
+
+API_RESPONSE = {}  # type: Dict[Any, Any]
+PROJECT_ID = "test-project"
+LOCATION = "asia-east2"
+FULL_LOCATION_PATH = "projects/test-project/locations/asia-east2"
+QUEUE_ID = "test-queue"
+FULL_QUEUE_PATH = "projects/test-project/locations/asia-east2/queues/test-queue"
+TASK_NAME = "test-task"
+FULL_TASK_PATH = "projects/test-project/locations/asia-east2/queues/test-queue/tasks/test-task"
+
+
+class TestCloudTasksEmptySensor(unittest.TestCase):
+    @mock.patch('airflow.providers.google.cloud.sensors.tasks.CloudTasksHook')
+    def test_queue_empty(self, mock_hook):
+
+        operator = TaskQueueEmptySensor(
+            task_id=TASK_NAME, location=LOCATION, project_id=PROJECT_ID, queue_name=QUEUE_ID, poke_interval=0
+        )
+
+        result = operator.poke(mock.MagicMock)
+
+        assert result is True
+
+    @mock.patch('airflow.providers.google.cloud.sensors.tasks.CloudTasksHook')
+    def test_queue_not_empty(self, mock_hook):
+        mock_hook.return_value.list_tasks.return_value = [Task(name=FULL_TASK_PATH)]
+
+        operator = TaskQueueEmptySensor(
+            task_id=TASK_NAME, location=LOCATION, project_id=PROJECT_ID, queue_name=QUEUE_ID, poke_interval=0
+        )
+
+        result = operator.poke(mock.MagicMock)
+
+        assert result is False