Skip to content

Commit 4c3fb1f

Browse files
vinitpayalvinitpayal
authored
Google Cloud Tasks Sensor for queue being empty (#25622)
1 parent 799b269 commit 4c3fb1f

8 files changed

Lines changed: 270 additions & 3 deletions

File tree

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
19+
"""
20+
Example Airflow DAG that sense a cloud task queue being empty.
21+
22+
This DAG relies on the following OS environment variables
23+
24+
* GCP_PROJECT_ID - Google Cloud project where the Compute Engine instance exists.
25+
* GCP_ZONE - Google Cloud zone where the cloud task queue exists.
26+
* QUEUE_NAME - Name of the cloud task queue.
27+
"""
28+
29+
import os
30+
from datetime import datetime
31+
32+
from airflow import models
33+
from airflow.providers.google.cloud.sensors.tasks import TaskQueueEmptySensor
34+
35+
GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'example-project')
36+
GCP_ZONE = os.environ.get('GCE_ZONE', 'europe-west1-b')
37+
QUEUE_NAME = os.environ.get('GCP_QUEUE_NAME', 'testqueue')
38+
39+
40+
with models.DAG(
41+
'example_gcp_cloud_tasks_sensor',
42+
start_date=datetime(2022, 8, 8),
43+
catchup=False,
44+
tags=['example'],
45+
) as dag:
46+
# [START cloud_tasks_empty_sensor]
47+
gcp_cloud_tasks_sensor = TaskQueueEmptySensor(
48+
project_id=GCP_PROJECT_ID,
49+
location=GCP_ZONE,
50+
task_id='gcp_sense_cloud_tasks_empty',
51+
queue_name=QUEUE_NAME,
52+
)
53+
# [END cloud_tasks_empty_sensor]

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

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
"""This module contains a Google Cloud Task sensor."""
19+
from typing import TYPE_CHECKING, Optional, Sequence, Union
20+
21+
from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
22+
from airflow.sensors.base import BaseSensorOperator
23+
24+
if TYPE_CHECKING:
25+
from airflow.utils.context import Context
26+
27+
28+
class TaskQueueEmptySensor(BaseSensorOperator):
29+
"""
30+
Pulls tasks count from a cloud task queue.
31+
Always waits for queue returning tasks count as 0.
32+
33+
:param project_id: the Google Cloud project ID for the subscription (templated)
34+
:param gcp_conn_id: The connection ID to use connecting to Google Cloud.
35+
:param queue_name: The queue name to for which task empty sensing is required.
36+
:param impersonation_chain: Optional service account to impersonate using short-term
37+
credentials, or chained list of accounts required to get the access_token
38+
of the last account in the list, which will be impersonated in the request.
39+
If set as a string, the account must grant the originating account
40+
the Service Account Token Creator IAM role.
41+
If set as a sequence, the identities from the list must grant
42+
Service Account Token Creator IAM role to the directly preceding identity, with first
43+
account from the list granting this role to the originating account (templated).
44+
"""
45+
46+
template_fields: Sequence[str] = (
47+
"project_id",
48+
"location",
49+
"queue_name",
50+
"gcp_conn_id",
51+
"impersonation_chain",
52+
)
53+
54+
def __init__(
55+
self,
56+
*,
57+
location: str,
58+
project_id: Optional[str] = None,
59+
queue_name: Optional[str] = None,
60+
gcp_conn_id: str = "google_cloud_default",
61+
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
62+
**kwargs,
63+
) -> None:
64+
super().__init__(**kwargs)
65+
self.location = location
66+
self.project_id = project_id
67+
self.queue_name = queue_name
68+
self.gcp_conn_id = gcp_conn_id
69+
self.impersonation_chain = impersonation_chain
70+
71+
def poke(self, context: "Context") -> bool:
72+
73+
hook = CloudTasksHook(
74+
gcp_conn_id=self.gcp_conn_id,
75+
impersonation_chain=self.impersonation_chain,
76+
)
77+
78+
# TODO uncomment page_size once https://issuetracker.google.com/issues/155978649?pli=1 gets fixed
79+
tasks = hook.list_tasks(
80+
location=self.location,
81+
queue_name=self.queue_name,
82+
# page_size=1
83+
)
84+
85+
self.log.info("tasks exhausted in cloud task queue?: %s" % (len(tasks) == 0))
86+
87+
return len(tasks) == 0

‎airflow/providers/google/provider.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -633,6 +633,9 @@ sensors:
633633
- integration-name: Google Cloud Dataform
634634
python-modules:
635635
- airflow.providers.google.cloud.sensors.dataform
636+
- integration-name: Google Cloud Tasks
637+
python-modules:
638+
- airflow.providers.google.cloud.sensors.tasks
636639

637640
hooks:
638641
- integration-name: Google Ads

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ Content
3030
Secrets backends <secrets-backends/google-cloud-secret-manager-backend>
3131
API Authentication backend <api-auth-backend/google-openid>
3232
Operators <operators/index>
33+
Sensors <sensors/index>
3334

3435
.. toctree::
3536
:maxdepth: 1

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@
1717
1818
Google Cloud Tasks
1919
==================
20-
21-
Firestore in Datastore mode is a NoSQL document database built for automatic scaling,
22-
high performance, and ease of application development.
20+
Cloud Tasks is a fully managed service that allows you to manage the execution, dispatch,
21+
and delivery of a large number of distributed tasks.
22+
Using Cloud Tasks, you can perform work asynchronously outside of a user or service-to-service request.
2323

2424
For more information about the service visit
2525
`Cloud Tasks product documentation <https://cloud.google.com/tasks/docs>`__

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
.. _google_cloud_tasks_empty_sensor:
19+
20+
Google Cloud Tasks
21+
==================
22+
Cloud Tasks is a fully managed service that allows you to manage the execution, dispatch,
23+
and delivery of a large number of distributed tasks.
24+
Using Cloud Tasks, you can perform work asynchronously outside of a user or service-to-service request.
25+
26+
For more information about the service visit
27+
`Cloud Tasks product documentation <https://cloud.google.com/tasks/docs>`__
28+
29+
Google Cloud Tasks Empty Sensor
30+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
31+
32+
To sense Queue being empty use
33+
:class:`~airflow.providers.google.cloud.sensor.tasks.TaskQueueEmptySensor`
34+
35+
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_cloud_task.py
36+
:language: python
37+
:dedent: 4
38+
:start-after: [START cloud_tasks_empty_sensor]
39+
:end-before: [END cloud_tasks_empty_sensor]

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
19+
20+
Google Sensors
21+
================
22+
23+
.. toctree::
24+
:maxdepth: 1
25+
26+
google-cloud-tasks

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
import unittest
19+
from typing import Any, Dict
20+
from unittest import mock
21+
22+
from google.cloud.tasks_v2.types import Task
23+
24+
from airflow.providers.google.cloud.sensors.tasks import TaskQueueEmptySensor
25+
26+
API_RESPONSE = {} # type: Dict[Any, Any]
27+
PROJECT_ID = "test-project"
28+
LOCATION = "asia-east2"
29+
FULL_LOCATION_PATH = "projects/test-project/locations/asia-east2"
30+
QUEUE_ID = "test-queue"
31+
FULL_QUEUE_PATH = "projects/test-project/locations/asia-east2/queues/test-queue"
32+
TASK_NAME = "test-task"
33+
FULL_TASK_PATH = "projects/test-project/locations/asia-east2/queues/test-queue/tasks/test-task"
34+
35+
36+
class TestCloudTasksEmptySensor(unittest.TestCase):
37+
@mock.patch('airflow.providers.google.cloud.sensors.tasks.CloudTasksHook')
38+
def test_queue_empty(self, mock_hook):
39+
40+
operator = TaskQueueEmptySensor(
41+
task_id=TASK_NAME, location=LOCATION, project_id=PROJECT_ID, queue_name=QUEUE_ID, poke_interval=0
42+
)
43+
44+
result = operator.poke(mock.MagicMock)
45+
46+
assert result is True
47+
48+
@mock.patch('airflow.providers.google.cloud.sensors.tasks.CloudTasksHook')
49+
def test_queue_not_empty(self, mock_hook):
50+
mock_hook.return_value.list_tasks.return_value = [Task(name=FULL_TASK_PATH)]
51+
52+
operator = TaskQueueEmptySensor(
53+
task_id=TASK_NAME, location=LOCATION, project_id=PROJECT_ID, queue_name=QUEUE_ID, poke_interval=0
54+
)
55+
56+
result = operator.poke(mock.MagicMock)
57+
58+
assert result is False

0 commit comments

Comments
 (0)