Skip to content

Commit 1f71221

Browse files
mik-lajmik-laj
authored
Support google-cloud-tasks>=2.0.0 (#13334)
1 parent 09c6549 commit 1f71221

6 files changed

Lines changed: 175 additions & 136 deletions

File tree

‎airflow/providers/google/ADDITIONAL_INFO.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ Details are covered in the UPDATING.md files for each library, but there are som
3333
| [``google-cloud-os-login``](https://pypi.org/project/google-cloud-os-login/) | ``>=1.0.0,<2.0.0`` | ``>=2.0.0,<3.0.0`` | [`UPGRADING.md`](https://github.com/googleapis/python-oslogin/blob/master/UPGRADING.md) |
3434
| [``google-cloud-pubsub``](https://pypi.org/project/google-cloud-pubsub/) | ``>=1.0.0,<2.0.0`` | ``>=2.0.0,<3.0.0`` | [`UPGRADING.md`](https://github.com/googleapis/python-pubsub/blob/master/UPGRADING.md) |
3535
| [``google-cloud-kms``](https://pypi.org/project/google-cloud-os-login/) | ``>=1.2.1,<2.0.0`` | ``>=2.0.0,<3.0.0`` | [`UPGRADING.md`](https://github.com/googleapis/python-kms/blob/master/UPGRADING.md) |
36+
| [``google-cloud-tasks``](https://pypi.org/project/google-cloud-tasks/) | ``>=1.2.1,<2.0.0`` | ``>=2.0.0,<3.0.0`` | [`UPGRADING.md`](https://github.com/googleapis/python-tasks/blob/master/UPGRADING.md) |
3637

3738

3839
### The field names use the snake_case convention

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

Lines changed: 62 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,13 @@
2121
which allows you to connect to Google Cloud Tasks service,
2222
performing actions to queues or tasks.
2323
"""
24+
2425
from typing import Dict, List, Optional, Sequence, Tuple, Union
2526

2627
from google.api_core.retry import Retry
27-
from google.cloud.tasks_v2 import CloudTasksClient, enums
28-
from google.cloud.tasks_v2.types import FieldMask, Queue, Task
28+
from google.cloud.tasks_v2 import CloudTasksClient
29+
from google.cloud.tasks_v2.types import Queue, Task
30+
from google.protobuf.field_mask_pb2 import FieldMask
2931

3032
from airflow.exceptions import AirflowException
3133
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
@@ -120,20 +122,19 @@ def create_queue(
120122
client = self.get_conn()
121123

122124
if queue_name:
123-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
125+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
124126
if isinstance(task_queue, Queue):
125127
task_queue.name = full_queue_name
126128
elif isinstance(task_queue, dict):
127129
task_queue['name'] = full_queue_name
128130
else:
129131
raise AirflowException('Unable to set queue_name.')
130-
full_location_path = CloudTasksClient.location_path(project_id, location)
132+
full_location_path = f"projects/{project_id}/locations/{location}"
131133
return client.create_queue(
132-
parent=full_location_path,
133-
queue=task_queue,
134+
request={'parent': full_location_path, 'queue': task_queue},
134135
retry=retry,
135136
timeout=timeout,
136-
metadata=metadata,
137+
metadata=metadata or (),
137138
)
138139

139140
@GoogleBaseHook.fallback_to_default_project_id
@@ -167,7 +168,7 @@ def update_queue(
167168
:param update_mask: A mast used to specify which fields of the queue are being updated.
168169
If empty, then all fields will be updated.
169170
If a dict is provided, it must be of the same form as the protobuf message.
170-
:type update_mask: dict or google.cloud.tasks_v2.types.FieldMask
171+
:type update_mask: dict or google.protobuf.field_mask_pb2.FieldMask
171172
:param retry: (Optional) A retry object used to retry requests.
172173
If None is specified, requests will not be retried.
173174
:type retry: google.api_core.retry.Retry
@@ -182,19 +183,18 @@ def update_queue(
182183
client = self.get_conn()
183184

184185
if queue_name and location:
185-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
186+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
186187
if isinstance(task_queue, Queue):
187188
task_queue.name = full_queue_name
188189
elif isinstance(task_queue, dict):
189190
task_queue['name'] = full_queue_name
190191
else:
191192
raise AirflowException('Unable to set queue_name.')
192193
return client.update_queue(
193-
queue=task_queue,
194-
update_mask=update_mask,
194+
request={'queue': task_queue, 'update_mask': update_mask},
195195
retry=retry,
196196
timeout=timeout,
197-
metadata=metadata,
197+
metadata=metadata or (),
198198
)
199199

200200
@GoogleBaseHook.fallback_to_default_project_id
@@ -230,8 +230,10 @@ def get_queue(
230230
"""
231231
client = self.get_conn()
232232

233-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
234-
return client.get_queue(name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata)
233+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
234+
return client.get_queue(
235+
request={'name': full_queue_name}, retry=retry, timeout=timeout, metadata=metadata or ()
236+
)
235237

236238
@GoogleBaseHook.fallback_to_default_project_id
237239
def list_queues(
@@ -270,14 +272,12 @@ def list_queues(
270272
"""
271273
client = self.get_conn()
272274

273-
full_location_path = CloudTasksClient.location_path(project_id, location)
275+
full_location_path = f"projects/{project_id}/locations/{location}"
274276
queues = client.list_queues(
275-
parent=full_location_path,
276-
filter_=results_filter,
277-
page_size=page_size,
277+
request={'parent': full_location_path, 'filter': results_filter, 'page_size': page_size},
278278
retry=retry,
279279
timeout=timeout,
280-
metadata=metadata,
280+
metadata=metadata or (),
281281
)
282282
return list(queues)
283283

@@ -313,8 +313,10 @@ def delete_queue(
313313
"""
314314
client = self.get_conn()
315315

316-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
317-
client.delete_queue(name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata)
316+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
317+
client.delete_queue(
318+
request={'name': full_queue_name}, retry=retry, timeout=timeout, metadata=metadata or ()
319+
)
318320

319321
@GoogleBaseHook.fallback_to_default_project_id
320322
def purge_queue(
@@ -349,8 +351,10 @@ def purge_queue(
349351
"""
350352
client = self.get_conn()
351353

352-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
353-
return client.purge_queue(name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata)
354+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
355+
return client.purge_queue(
356+
request={'name': full_queue_name}, retry=retry, timeout=timeout, metadata=metadata or ()
357+
)
354358

355359
@GoogleBaseHook.fallback_to_default_project_id
356360
def pause_queue(
@@ -385,8 +389,10 @@ def pause_queue(
385389
"""
386390
client = self.get_conn()
387391

388-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
389-
return client.pause_queue(name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata)
392+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
393+
return client.pause_queue(
394+
request={'name': full_queue_name}, retry=retry, timeout=timeout, metadata=metadata or ()
395+
)
390396

391397
@GoogleBaseHook.fallback_to_default_project_id
392398
def resume_queue(
@@ -421,8 +427,10 @@ def resume_queue(
421427
"""
422428
client = self.get_conn()
423429

424-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
425-
return client.resume_queue(name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata)
430+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
431+
return client.resume_queue(
432+
request={'name': full_queue_name}, retry=retry, timeout=timeout, metadata=metadata or ()
433+
)
426434

427435
@GoogleBaseHook.fallback_to_default_project_id
428436
def create_task(
@@ -432,7 +440,7 @@ def create_task(
432440
task: Union[Dict, Task],
433441
project_id: str,
434442
task_name: Optional[str] = None,
435-
response_view: Optional[enums.Task.View] = None,
443+
response_view: Optional = None,
436444
retry: Optional[Retry] = None,
437445
timeout: Optional[float] = None,
438446
metadata: Optional[Sequence[Tuple[str, str]]] = None,
@@ -455,7 +463,7 @@ def create_task(
455463
:type task_name: str
456464
:param response_view: (Optional) This field specifies which subset of the Task will
457465
be returned.
458-
:type response_view: google.cloud.tasks_v2.enums.Task.View
466+
:type response_view: google.cloud.tasks_v2.Task.View
459467
:param retry: (Optional) A retry object used to retry requests.
460468
If None is specified, requests will not be retried.
461469
:type retry: google.api_core.retry.Retry
@@ -470,21 +478,21 @@ def create_task(
470478
client = self.get_conn()
471479

472480
if task_name:
473-
full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name)
481+
full_task_name = (
482+
f"projects/{project_id}/locations/{location}/queues/{queue_name}/tasks/{task_name}"
483+
)
474484
if isinstance(task, Task):
475485
task.name = full_task_name
476486
elif isinstance(task, dict):
477487
task['name'] = full_task_name
478488
else:
479489
raise AirflowException('Unable to set task_name.')
480-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
490+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
481491
return client.create_task(
482-
parent=full_queue_name,
483-
task=task,
484-
response_view=response_view,
492+
request={'parent': full_queue_name, 'task': task, 'response_view': response_view},
485493
retry=retry,
486494
timeout=timeout,
487-
metadata=metadata,
495+
metadata=metadata or (),
488496
)
489497

490498
@GoogleBaseHook.fallback_to_default_project_id
@@ -494,7 +502,7 @@ def get_task(
494502
queue_name: str,
495503
task_name: str,
496504
project_id: str,
497-
response_view: Optional[enums.Task.View] = None,
505+
response_view: Optional = None,
498506
retry: Optional[Retry] = None,
499507
timeout: Optional[float] = None,
500508
metadata: Optional[Sequence[Tuple[str, str]]] = None,
@@ -513,7 +521,7 @@ def get_task(
513521
:type project_id: str
514522
:param response_view: (Optional) This field specifies which subset of the Task will
515523
be returned.
516-
:type response_view: google.cloud.tasks_v2.enums.Task.View
524+
:type response_view: google.cloud.tasks_v2.Task.View
517525
:param retry: (Optional) A retry object used to retry requests.
518526
If None is specified, requests will not be retried.
519527
:type retry: google.api_core.retry.Retry
@@ -527,13 +535,12 @@ def get_task(
527535
"""
528536
client = self.get_conn()
529537

530-
full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name)
538+
full_task_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}/tasks/{task_name}"
531539
return client.get_task(
532-
name=full_task_name,
533-
response_view=response_view,
540+
request={'name': full_task_name, 'response_view': response_view},
534541
retry=retry,
535542
timeout=timeout,
536-
metadata=metadata,
543+
metadata=metadata or (),
537544
)
538545

539546
@GoogleBaseHook.fallback_to_default_project_id
@@ -542,7 +549,7 @@ def list_tasks(
542549
location: str,
543550
queue_name: str,
544551
project_id: str,
545-
response_view: Optional[enums.Task.View] = None,
552+
response_view: Optional = None,
546553
page_size: Optional[int] = None,
547554
retry: Optional[Retry] = None,
548555
timeout: Optional[float] = None,
@@ -560,7 +567,7 @@ def list_tasks(
560567
:type project_id: str
561568
:param response_view: (Optional) This field specifies which subset of the Task will
562569
be returned.
563-
:type response_view: google.cloud.tasks_v2.enums.Task.View
570+
:type response_view: google.cloud.tasks_v2.Task.View
564571
:param page_size: (Optional) The maximum number of resources contained in the
565572
underlying API response.
566573
:type page_size: int
@@ -576,14 +583,12 @@ def list_tasks(
576583
:rtype: list[google.cloud.tasks_v2.types.Task]
577584
"""
578585
client = self.get_conn()
579-
full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name)
586+
full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
580587
tasks = client.list_tasks(
581-
parent=full_queue_name,
582-
response_view=response_view,
583-
page_size=page_size,
588+
request={'parent': full_queue_name, 'response_view': response_view, 'page_size': page_size},
584589
retry=retry,
585590
timeout=timeout,
586-
metadata=metadata,
591+
metadata=metadata or (),
587592
)
588593
return list(tasks)
589594

@@ -622,8 +627,10 @@ def delete_task(
622627
"""
623628
client = self.get_conn()
624629

625-
full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name)
626-
client.delete_task(name=full_task_name, retry=retry, timeout=timeout, metadata=metadata)
630+
full_task_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}/tasks/{task_name}"
631+
client.delete_task(
632+
request={'name': full_task_name}, retry=retry, timeout=timeout, metadata=metadata or ()
633+
)
627634

628635
@GoogleBaseHook.fallback_to_default_project_id
629636
def run_task(
@@ -632,7 +639,7 @@ def run_task(
632639
queue_name: str,
633640
task_name: str,
634641
project_id: str,
635-
response_view: Optional[enums.Task.View] = None,
642+
response_view: Optional = None,
636643
retry: Optional[Retry] = None,
637644
timeout: Optional[float] = None,
638645
metadata: Optional[Sequence[Tuple[str, str]]] = None,
@@ -651,7 +658,7 @@ def run_task(
651658
:type project_id: str
652659
:param response_view: (Optional) This field specifies which subset of the Task will
653660
be returned.
654-
:type response_view: google.cloud.tasks_v2.enums.Task.View
661+
:type response_view: google.cloud.tasks_v2.Task.View
655662
:param retry: (Optional) A retry object used to retry requests.
656663
If None is specified, requests will not be retried.
657664
:type retry: google.api_core.retry.Retry
@@ -665,11 +672,10 @@ def run_task(
665672
"""
666673
client = self.get_conn()
667674

668-
full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name)
675+
full_task_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}/tasks/{task_name}"
669676
return client.run_task(
670-
name=full_task_name,
671-
response_view=response_view,
677+
request={'name': full_task_name, 'response_view': response_view},
672678
retry=retry,
673679
timeout=timeout,
674-
metadata=metadata,
680+
metadata=metadata or (),
675681
)

0 commit comments

Comments
 (0)