Skip to content

Commit 8230ccc

Browse files
turbaszekturbaszek
authored
[AIRFLOW-6926] Fix Google Tasks operators return types and idempotency (#7547)
1 parent 0d1e308 commit 8230ccc

3 files changed

Lines changed: 57 additions & 33 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939
timestamp = timestamp_pb2.Timestamp()
4040
timestamp.FromDatetime(datetime.now() + timedelta(hours=12)) # pylint: disable=no-member
4141

42-
LOCATION = "asia-east2"
42+
LOCATION = "europe-west1"
4343
QUEUE_ID = "cloud-tasks-queue"
4444
TASK_NAME = "task-to-run"
4545

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

Lines changed: 43 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,11 @@
2323
"""
2424
from typing import Dict, Optional, Sequence, Tuple, Union
2525

26+
from google.api_core.exceptions import AlreadyExists
2627
from google.api_core.retry import Retry
2728
from google.cloud.tasks_v2 import enums
2829
from google.cloud.tasks_v2.types import FieldMask, Queue, Task
30+
from google.protobuf.json_format import MessageToDict
2931

3032
from airflow.models import BaseOperator
3133
from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
@@ -98,15 +100,27 @@ def __init__(
98100

99101
def execute(self, context):
100102
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
101-
return hook.create_queue(
102-
location=self.location,
103-
task_queue=self.task_queue,
104-
project_id=self.project_id,
105-
queue_name=self.queue_name,
106-
retry=self.retry,
107-
timeout=self.timeout,
108-
metadata=self.metadata,
109-
)
103+
try:
104+
queue = hook.create_queue(
105+
location=self.location,
106+
task_queue=self.task_queue,
107+
project_id=self.project_id,
108+
queue_name=self.queue_name,
109+
retry=self.retry,
110+
timeout=self.timeout,
111+
metadata=self.metadata,
112+
)
113+
except AlreadyExists:
114+
queue = hook.get_queue(
115+
location=self.location,
116+
project_id=self.project_id,
117+
queue_name=self.queue_name,
118+
retry=self.retry,
119+
timeout=self.timeout,
120+
metadata=self.metadata,
121+
)
122+
123+
return MessageToDict(queue)
110124

111125

112126
class CloudTasksQueueUpdateOperator(BaseOperator):
@@ -181,7 +195,7 @@ def __init__(
181195

182196
def execute(self, context):
183197
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
184-
return hook.update_queue(
198+
queue = hook.update_queue(
185199
task_queue=self.task_queue,
186200
project_id=self.project_id,
187201
location=self.location,
@@ -191,6 +205,7 @@ def execute(self, context):
191205
timeout=self.timeout,
192206
metadata=self.metadata,
193207
)
208+
return MessageToDict(queue)
194209

195210

196211
class CloudTasksQueueGetOperator(BaseOperator):
@@ -244,14 +259,15 @@ def __init__(
244259

245260
def execute(self, context):
246261
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
247-
return hook.get_queue(
262+
queue = hook.get_queue(
248263
location=self.location,
249264
queue_name=self.queue_name,
250265
project_id=self.project_id,
251266
retry=self.retry,
252267
timeout=self.timeout,
253268
metadata=self.metadata,
254269
)
270+
return MessageToDict(queue)
255271

256272

257273
class CloudTasksQueuesListOperator(BaseOperator):
@@ -311,7 +327,7 @@ def __init__(
311327

312328
def execute(self, context):
313329
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
314-
return hook.list_queues(
330+
queues = hook.list_queues(
315331
location=self.location,
316332
project_id=self.project_id,
317333
results_filter=self.results_filter,
@@ -320,6 +336,7 @@ def execute(self, context):
320336
timeout=self.timeout,
321337
metadata=self.metadata,
322338
)
339+
return [MessageToDict(q) for q in queues]
323340

324341

325342
class CloudTasksQueueDeleteOperator(BaseOperator):
@@ -433,14 +450,15 @@ def __init__(
433450

434451
def execute(self, context):
435452
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
436-
return hook.purge_queue(
453+
queue = hook.purge_queue(
437454
location=self.location,
438455
queue_name=self.queue_name,
439456
project_id=self.project_id,
440457
retry=self.retry,
441458
timeout=self.timeout,
442459
metadata=self.metadata,
443460
)
461+
return MessageToDict(queue)
444462

445463

446464
class CloudTasksQueuePauseOperator(BaseOperator):
@@ -494,14 +512,15 @@ def __init__(
494512

495513
def execute(self, context):
496514
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
497-
return hook.pause_queue(
515+
queues = hook.pause_queue(
498516
location=self.location,
499517
queue_name=self.queue_name,
500518
project_id=self.project_id,
501519
retry=self.retry,
502520
timeout=self.timeout,
503521
metadata=self.metadata,
504522
)
523+
return [MessageToDict(q) for q in queues]
505524

506525

507526
class CloudTasksQueueResumeOperator(BaseOperator):
@@ -555,14 +574,15 @@ def __init__(
555574

556575
def execute(self, context):
557576
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
558-
return hook.resume_queue(
577+
queue = hook.resume_queue(
559578
location=self.location,
560579
queue_name=self.queue_name,
561580
project_id=self.project_id,
562581
retry=self.retry,
563582
timeout=self.timeout,
564583
metadata=self.metadata,
565584
)
585+
return MessageToDict(queue)
566586

567587

568588
class CloudTasksTaskCreateOperator(BaseOperator):
@@ -638,7 +658,7 @@ def __init__( # pylint: disable=too-many-arguments
638658

639659
def execute(self, context):
640660
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
641-
return hook.create_task(
661+
task = hook.create_task(
642662
location=self.location,
643663
queue_name=self.queue_name,
644664
task=self.task,
@@ -649,6 +669,7 @@ def execute(self, context):
649669
timeout=self.timeout,
650670
metadata=self.metadata,
651671
)
672+
return MessageToDict(task)
652673

653674

654675
class CloudTasksTaskGetOperator(BaseOperator):
@@ -717,7 +738,7 @@ def __init__(
717738

718739
def execute(self, context):
719740
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
720-
return hook.get_task(
741+
task = hook.get_task(
721742
location=self.location,
722743
queue_name=self.queue_name,
723744
task_name=self.task_name,
@@ -727,6 +748,7 @@ def execute(self, context):
727748
timeout=self.timeout,
728749
metadata=self.metadata,
729750
)
751+
return MessageToDict(task)
730752

731753

732754
class CloudTasksTasksListOperator(BaseOperator):
@@ -790,7 +812,7 @@ def __init__(
790812

791813
def execute(self, context):
792814
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
793-
return hook.list_tasks(
815+
tasks = hook.list_tasks(
794816
location=self.location,
795817
queue_name=self.queue_name,
796818
project_id=self.project_id,
@@ -800,6 +822,7 @@ def execute(self, context):
800822
timeout=self.timeout,
801823
metadata=self.metadata,
802824
)
825+
return [MessageToDict(t) for t in tasks]
803826

804827

805828
class CloudTasksTaskDeleteOperator(BaseOperator):
@@ -939,7 +962,7 @@ def __init__(
939962

940963
def execute(self, context):
941964
hook = CloudTasksHook(gcp_conn_id=self.gcp_conn_id)
942-
return hook.run_task(
965+
task = hook.run_task(
943966
location=self.location,
944967
queue_name=self.queue_name,
945968
task_name=self.task_name,
@@ -949,3 +972,4 @@ def execute(self, context):
949972
timeout=self.timeout,
950973
metadata=self.metadata,
951974
)
975+
return MessageToDict(task)

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

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
class TestCloudTasksQueueCreate(unittest.TestCase):
4545
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
4646
def test_create_queue(self, mock_hook):
47-
mock_hook.return_value.create_queue.return_value = {}
47+
mock_hook.return_value.create_queue.return_value = mock.MagicMock()
4848
operator = CloudTasksQueueCreateOperator(
4949
location=LOCATION, task_queue=Queue(), task_id="id"
5050
)
@@ -64,7 +64,7 @@ def test_create_queue(self, mock_hook):
6464
class TestCloudTasksQueueUpdate(unittest.TestCase):
6565
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
6666
def test_update_queue(self, mock_hook):
67-
mock_hook.return_value.update_queue.return_value = {}
67+
mock_hook.return_value.update_queue.return_value = mock.MagicMock()
6868
operator = CloudTasksQueueUpdateOperator(
6969
task_queue=Queue(name=FULL_QUEUE_PATH), task_id="id"
7070
)
@@ -85,7 +85,7 @@ def test_update_queue(self, mock_hook):
8585
class TestCloudTasksQueueGet(unittest.TestCase):
8686
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
8787
def test_get_queue(self, mock_hook):
88-
mock_hook.return_value.get_queue.return_value = {}
88+
mock_hook.return_value.get_queue.return_value = mock.MagicMock()
8989
operator = CloudTasksQueueGetOperator(
9090
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
9191
)
@@ -104,7 +104,7 @@ def test_get_queue(self, mock_hook):
104104
class TestCloudTasksQueuesList(unittest.TestCase):
105105
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
106106
def test_list_queues(self, mock_hook):
107-
mock_hook.return_value.list_queues.return_value = {}
107+
mock_hook.return_value.list_queues.return_value = mock.MagicMock()
108108
operator = CloudTasksQueuesListOperator(location=LOCATION, task_id="id")
109109
operator.execute(context=None)
110110
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
@@ -122,7 +122,7 @@ def test_list_queues(self, mock_hook):
122122
class TestCloudTasksQueueDelete(unittest.TestCase):
123123
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
124124
def test_delete_queue(self, mock_hook):
125-
mock_hook.return_value.delete_queue.return_value = {}
125+
mock_hook.return_value.delete_queue.return_value = mock.MagicMock()
126126
operator = CloudTasksQueueDeleteOperator(
127127
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
128128
)
@@ -141,7 +141,7 @@ def test_delete_queue(self, mock_hook):
141141
class TestCloudTasksQueuePurge(unittest.TestCase):
142142
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
143143
def test_delete_queue(self, mock_hook):
144-
mock_hook.return_value.purge_queue.return_value = {}
144+
mock_hook.return_value.purge_queue.return_value = mock.MagicMock()
145145
operator = CloudTasksQueuePurgeOperator(
146146
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
147147
)
@@ -160,7 +160,7 @@ def test_delete_queue(self, mock_hook):
160160
class TestCloudTasksQueuePause(unittest.TestCase):
161161
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
162162
def test_pause_queue(self, mock_hook):
163-
mock_hook.return_value.pause_queue.return_value = {}
163+
mock_hook.return_value.pause_queue.return_value = mock.MagicMock()
164164
operator = CloudTasksQueuePauseOperator(
165165
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
166166
)
@@ -179,7 +179,7 @@ def test_pause_queue(self, mock_hook):
179179
class TestCloudTasksQueueResume(unittest.TestCase):
180180
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
181181
def test_resume_queue(self, mock_hook):
182-
mock_hook.return_value.resume_queue.return_value = {}
182+
mock_hook.return_value.resume_queue.return_value = mock.MagicMock()
183183
operator = CloudTasksQueueResumeOperator(
184184
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
185185
)
@@ -198,7 +198,7 @@ def test_resume_queue(self, mock_hook):
198198
class TestCloudTasksTaskCreate(unittest.TestCase):
199199
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
200200
def test_create_task(self, mock_hook):
201-
mock_hook.return_value.create_task.return_value = {}
201+
mock_hook.return_value.create_task.return_value = mock.MagicMock()
202202
operator = CloudTasksTaskCreateOperator(
203203
location=LOCATION, queue_name=QUEUE_ID, task=Task(), task_id="id"
204204
)
@@ -220,7 +220,7 @@ def test_create_task(self, mock_hook):
220220
class TestCloudTasksTaskGet(unittest.TestCase):
221221
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
222222
def test_get_task(self, mock_hook):
223-
mock_hook.return_value.get_task.return_value = {}
223+
mock_hook.return_value.get_task.return_value = mock.MagicMock()
224224
operator = CloudTasksTaskGetOperator(
225225
location=LOCATION, queue_name=QUEUE_ID, task_name=TASK_NAME, task_id="id"
226226
)
@@ -241,7 +241,7 @@ def test_get_task(self, mock_hook):
241241
class TestCloudTasksTasksList(unittest.TestCase):
242242
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
243243
def test_list_tasks(self, mock_hook):
244-
mock_hook.return_value.list_tasks.return_value = {}
244+
mock_hook.return_value.list_tasks.return_value = mock.MagicMock()
245245
operator = CloudTasksTasksListOperator(
246246
location=LOCATION, queue_name=QUEUE_ID, task_id="id"
247247
)
@@ -262,7 +262,7 @@ def test_list_tasks(self, mock_hook):
262262
class TestCloudTasksTaskDelete(unittest.TestCase):
263263
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
264264
def test_delete_task(self, mock_hook):
265-
mock_hook.return_value.delete_task.return_value = {}
265+
mock_hook.return_value.delete_task.return_value = mock.MagicMock()
266266
operator = CloudTasksTaskDeleteOperator(
267267
location=LOCATION, queue_name=QUEUE_ID, task_name=TASK_NAME, task_id="id"
268268
)
@@ -282,7 +282,7 @@ def test_delete_task(self, mock_hook):
282282
class TestCloudTasksTaskRun(unittest.TestCase):
283283
@mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook")
284284
def test_run_task(self, mock_hook):
285-
mock_hook.return_value.run_task.return_value = {}
285+
mock_hook.return_value.run_task.return_value = mock.MagicMock()
286286
operator = CloudTasksTaskRunOperator(
287287
location=LOCATION, queue_name=QUEUE_ID, task_name=TASK_NAME, task_id="id"
288288
)

0 commit comments

Comments
 (0)