2121which allows you to connect to Google Cloud Tasks service,
2222performing actions to queues or tasks.
2323"""
24-
2524from typing import Dict , List , Optional , Sequence , Tuple , Union
2625
2726from google .api_core .retry import Retry
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
27+ from google .cloud .tasks_v2 import CloudTasksClient , enums
28+ from google .cloud .tasks_v2 .types import FieldMask , Queue , Task
3129
3230from airflow .exceptions import AirflowException
3331from airflow .providers .google .common .hooks .base_google import GoogleBaseHook
@@ -122,19 +120,20 @@ def create_queue(
122120 client = self .get_conn ()
123121
124122 if queue_name :
125- full_queue_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } "
123+ full_queue_name = CloudTasksClient . queue_path ( project_id , location , queue_name )
126124 if isinstance (task_queue , Queue ):
127125 task_queue .name = full_queue_name
128126 elif isinstance (task_queue , dict ):
129127 task_queue ['name' ] = full_queue_name
130128 else :
131129 raise AirflowException ('Unable to set queue_name.' )
132- full_location_path = f"projects/ { project_id } /locations/ { location } "
130+ full_location_path = CloudTasksClient . location_path ( project_id , location )
133131 return client .create_queue (
134- request = {'parent' : full_location_path , 'queue' : task_queue },
132+ parent = full_location_path ,
133+ queue = task_queue ,
135134 retry = retry ,
136135 timeout = timeout ,
137- metadata = metadata or () ,
136+ metadata = metadata ,
138137 )
139138
140139 @GoogleBaseHook .fallback_to_default_project_id
@@ -168,7 +167,7 @@ def update_queue(
168167 :param update_mask: A mast used to specify which fields of the queue are being updated.
169168 If empty, then all fields will be updated.
170169 If a dict is provided, it must be of the same form as the protobuf message.
171- :type update_mask: dict or google.protobuf.field_mask_pb2 .FieldMask
170+ :type update_mask: dict or google.cloud.tasks_v2.types .FieldMask
172171 :param retry: (Optional) A retry object used to retry requests.
173172 If None is specified, requests will not be retried.
174173 :type retry: google.api_core.retry.Retry
@@ -183,18 +182,19 @@ def update_queue(
183182 client = self .get_conn ()
184183
185184 if queue_name and location :
186- full_queue_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } "
185+ full_queue_name = CloudTasksClient . queue_path ( project_id , location , queue_name )
187186 if isinstance (task_queue , Queue ):
188187 task_queue .name = full_queue_name
189188 elif isinstance (task_queue , dict ):
190189 task_queue ['name' ] = full_queue_name
191190 else :
192191 raise AirflowException ('Unable to set queue_name.' )
193192 return client .update_queue (
194- request = {'queue' : task_queue , 'update_mask' : update_mask },
193+ queue = task_queue ,
194+ update_mask = update_mask ,
195195 retry = retry ,
196196 timeout = timeout ,
197- metadata = metadata or () ,
197+ metadata = metadata ,
198198 )
199199
200200 @GoogleBaseHook .fallback_to_default_project_id
@@ -230,10 +230,8 @@ def get_queue(
230230 """
231231 client = self .get_conn ()
232232
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- )
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 )
237235
238236 @GoogleBaseHook .fallback_to_default_project_id
239237 def list_queues (
@@ -272,12 +270,14 @@ def list_queues(
272270 """
273271 client = self .get_conn ()
274272
275- full_location_path = f"projects/ { project_id } /locations/ { location } "
273+ full_location_path = CloudTasksClient . location_path ( project_id , location )
276274 queues = client .list_queues (
277- request = {'parent' : full_location_path , 'filter' : results_filter , 'page_size' : page_size },
275+ parent = full_location_path ,
276+ filter_ = results_filter ,
277+ page_size = page_size ,
278278 retry = retry ,
279279 timeout = timeout ,
280- metadata = metadata or () ,
280+ metadata = metadata ,
281281 )
282282 return list (queues )
283283
@@ -313,10 +313,8 @@ def delete_queue(
313313 """
314314 client = self .get_conn ()
315315
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- )
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 )
320318
321319 @GoogleBaseHook .fallback_to_default_project_id
322320 def purge_queue (
@@ -351,10 +349,8 @@ def purge_queue(
351349 """
352350 client = self .get_conn ()
353351
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- )
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 )
358354
359355 @GoogleBaseHook .fallback_to_default_project_id
360356 def pause_queue (
@@ -389,10 +385,8 @@ def pause_queue(
389385 """
390386 client = self .get_conn ()
391387
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- )
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 )
396390
397391 @GoogleBaseHook .fallback_to_default_project_id
398392 def resume_queue (
@@ -427,10 +421,8 @@ def resume_queue(
427421 """
428422 client = self .get_conn ()
429423
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- )
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 )
434426
435427 @GoogleBaseHook .fallback_to_default_project_id
436428 def create_task (
@@ -440,7 +432,7 @@ def create_task(
440432 task : Union [Dict , Task ],
441433 project_id : str ,
442434 task_name : Optional [str ] = None ,
443- response_view : Optional = None ,
435+ response_view : Optional [ enums . Task . View ] = None ,
444436 retry : Optional [Retry ] = None ,
445437 timeout : Optional [float ] = None ,
446438 metadata : Optional [Sequence [Tuple [str , str ]]] = None ,
@@ -463,7 +455,7 @@ def create_task(
463455 :type task_name: str
464456 :param response_view: (Optional) This field specifies which subset of the Task will
465457 be returned.
466- :type response_view: google.cloud.tasks_v2.Task.View
458+ :type response_view: google.cloud.tasks_v2.enums. Task.View
467459 :param retry: (Optional) A retry object used to retry requests.
468460 If None is specified, requests will not be retried.
469461 :type retry: google.api_core.retry.Retry
@@ -478,21 +470,21 @@ def create_task(
478470 client = self .get_conn ()
479471
480472 if task_name :
481- full_task_name = (
482- f"projects/{ project_id } /locations/{ location } /queues/{ queue_name } /tasks/{ task_name } "
483- )
473+ full_task_name = CloudTasksClient .task_path (project_id , location , queue_name , task_name )
484474 if isinstance (task , Task ):
485475 task .name = full_task_name
486476 elif isinstance (task , dict ):
487477 task ['name' ] = full_task_name
488478 else :
489479 raise AirflowException ('Unable to set task_name.' )
490- full_queue_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } "
480+ full_queue_name = CloudTasksClient . queue_path ( project_id , location , queue_name )
491481 return client .create_task (
492- request = {'parent' : full_queue_name , 'task' : task , 'response_view' : response_view },
482+ parent = full_queue_name ,
483+ task = task ,
484+ response_view = response_view ,
493485 retry = retry ,
494486 timeout = timeout ,
495- metadata = metadata or () ,
487+ metadata = metadata ,
496488 )
497489
498490 @GoogleBaseHook .fallback_to_default_project_id
@@ -502,7 +494,7 @@ def get_task(
502494 queue_name : str ,
503495 task_name : str ,
504496 project_id : str ,
505- response_view : Optional = None ,
497+ response_view : Optional [ enums . Task . View ] = None ,
506498 retry : Optional [Retry ] = None ,
507499 timeout : Optional [float ] = None ,
508500 metadata : Optional [Sequence [Tuple [str , str ]]] = None ,
@@ -521,7 +513,7 @@ def get_task(
521513 :type project_id: str
522514 :param response_view: (Optional) This field specifies which subset of the Task will
523515 be returned.
524- :type response_view: google.cloud.tasks_v2.Task.View
516+ :type response_view: google.cloud.tasks_v2.enums. Task.View
525517 :param retry: (Optional) A retry object used to retry requests.
526518 If None is specified, requests will not be retried.
527519 :type retry: google.api_core.retry.Retry
@@ -535,12 +527,13 @@ def get_task(
535527 """
536528 client = self .get_conn ()
537529
538- full_task_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } /tasks/ { task_name } "
530+ full_task_name = CloudTasksClient . task_path ( project_id , location , queue_name , task_name )
539531 return client .get_task (
540- request = {'name' : full_task_name , 'response_view' : response_view },
532+ name = full_task_name ,
533+ response_view = response_view ,
541534 retry = retry ,
542535 timeout = timeout ,
543- metadata = metadata or () ,
536+ metadata = metadata ,
544537 )
545538
546539 @GoogleBaseHook .fallback_to_default_project_id
@@ -549,7 +542,7 @@ def list_tasks(
549542 location : str ,
550543 queue_name : str ,
551544 project_id : str ,
552- response_view : Optional = None ,
545+ response_view : Optional [ enums . Task . View ] = None ,
553546 page_size : Optional [int ] = None ,
554547 retry : Optional [Retry ] = None ,
555548 timeout : Optional [float ] = None ,
@@ -567,7 +560,7 @@ def list_tasks(
567560 :type project_id: str
568561 :param response_view: (Optional) This field specifies which subset of the Task will
569562 be returned.
570- :type response_view: google.cloud.tasks_v2.Task.View
563+ :type response_view: google.cloud.tasks_v2.enums. Task.View
571564 :param page_size: (Optional) The maximum number of resources contained in the
572565 underlying API response.
573566 :type page_size: int
@@ -583,12 +576,14 @@ def list_tasks(
583576 :rtype: list[google.cloud.tasks_v2.types.Task]
584577 """
585578 client = self .get_conn ()
586- full_queue_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } "
579+ full_queue_name = CloudTasksClient . queue_path ( project_id , location , queue_name )
587580 tasks = client .list_tasks (
588- request = {'parent' : full_queue_name , 'response_view' : response_view , 'page_size' : page_size },
581+ parent = full_queue_name ,
582+ response_view = response_view ,
583+ page_size = page_size ,
589584 retry = retry ,
590585 timeout = timeout ,
591- metadata = metadata or () ,
586+ metadata = metadata ,
592587 )
593588 return list (tasks )
594589
@@ -627,10 +622,8 @@ def delete_task(
627622 """
628623 client = self .get_conn ()
629624
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- )
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 )
634627
635628 @GoogleBaseHook .fallback_to_default_project_id
636629 def run_task (
@@ -639,7 +632,7 @@ def run_task(
639632 queue_name : str ,
640633 task_name : str ,
641634 project_id : str ,
642- response_view : Optional = None ,
635+ response_view : Optional [ enums . Task . View ] = None ,
643636 retry : Optional [Retry ] = None ,
644637 timeout : Optional [float ] = None ,
645638 metadata : Optional [Sequence [Tuple [str , str ]]] = None ,
@@ -658,7 +651,7 @@ def run_task(
658651 :type project_id: str
659652 :param response_view: (Optional) This field specifies which subset of the Task will
660653 be returned.
661- :type response_view: google.cloud.tasks_v2.Task.View
654+ :type response_view: google.cloud.tasks_v2.enums. Task.View
662655 :param retry: (Optional) A retry object used to retry requests.
663656 If None is specified, requests will not be retried.
664657 :type retry: google.api_core.retry.Retry
@@ -672,10 +665,11 @@ def run_task(
672665 """
673666 client = self .get_conn ()
674667
675- full_task_name = f"projects/ { project_id } /locations/ { location } /queues/ { queue_name } /tasks/ { task_name } "
668+ full_task_name = CloudTasksClient . task_path ( project_id , location , queue_name , task_name )
676669 return client .run_task (
677- request = {'name' : full_task_name , 'response_view' : response_view },
670+ name = full_task_name ,
671+ response_view = response_view ,
678672 retry = retry ,
679673 timeout = timeout ,
680- metadata = metadata or () ,
674+ metadata = metadata ,
681675 )
0 commit comments