Source code for airflow.gcp.hooks.tasks

# -*- coding: utf-8 -*-
#
# 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 CloudTasksHook
which allows you to connect to GCP Cloud Tasks service,
performing actions to queues or tasks.
"""

from google.cloud.tasks_v2 import CloudTasksClient
from google.cloud.tasks_v2.types import Queue, Task

from airflow import AirflowException
from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook


[docs]class CloudTasksHook(GoogleCloudBaseHook): """ Hook for Google Cloud Tasks APIs. Cloud Tasks allows developers to manage the execution of background work in their applications. All the methods in the hook where project_id is used must be called with keyword arguments rather than positional. :param gcp_conn_id: The connection ID to use when fetching connection info. :type gcp_conn_id: str :param delegate_to: The account to impersonate, if any. For this to work, the service account making the request must have domain-wide delegation enabled. :type delegate_to: str """ def __init__(self, gcp_conn_id="google_cloud_default", delegate_to=None): super().__init__(gcp_conn_id, delegate_to) self._client = None
[docs] def get_conn(self): """ Provides a client for interacting with the Cloud Tasks API. :return: GCP Cloud Tasks API Client :rtype: google.cloud.tasks_v2.CloudTasksClient """ if not self._client: self._client = CloudTasksClient( credentials=self._get_credentials(), client_info=self.client_info ) return self._client
@GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def create_queue( self, location, task_queue, project_id=None, queue_name=None, retry=None, timeout=None, metadata=None, ): """ Creates a queue in Cloud Tasks. :param location: The location name in which the queue will be created. :type location: str :param task_queue: The task queue to create. Queue's name cannot be the same as an existing queue. If a dict is provided, it must be of the same form as the protobuf message Queue. :type task_queue: dict or class google.cloud.tasks_v2.types.Queue :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param queue_name: (Optional) The queue's name. If provided, it will be used to construct the full queue path. :type queue_name: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Queue """ client = self.get_conn() if queue_name: full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) if isinstance(task_queue, Queue): task_queue.name = full_queue_name elif isinstance(task_queue, dict): task_queue['name'] = full_queue_name else: raise AirflowException('Unable to set queue_name.') full_location_path = CloudTasksClient.location_path(project_id, location) return client.create_queue( parent=full_location_path, queue=task_queue, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def update_queue( self, task_queue, project_id=None, location=None, queue_name=None, update_mask=None, retry=None, timeout=None, metadata=None, ): """ Updates a queue in Cloud Tasks. :param task_queue: The task queue to update. This method creates the queue if it does not exist and updates the queue if it does exist. The queue's name must be specified. :type task_queue: dict or class google.cloud.tasks_v2.types.Queue :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param location: (Optional) The location name in which the queue will be updated. If provided, it will be used to construct the full queue path. :type location: str :param queue_name: (Optional) The queue's name. If provided, it will be used to construct the full queue path. :type queue_name: str :param update_mask: A mast used to specify which fields of the queue are being updated. If empty, then all fields will be updated. If a dict is provided, it must be of the same form as the protobuf message. :type update_mask: dict or class google.cloud.tasks_v2.types.FieldMask :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Queue """ client = self.get_conn() if queue_name and location: full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) if isinstance(task_queue, Queue): task_queue.name = full_queue_name elif isinstance(task_queue, dict): task_queue['name'] = full_queue_name else: raise AirflowException('Unable to set queue_name.') return client.update_queue( queue=task_queue, update_mask=update_mask, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def get_queue( self, location, queue_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Gets a queue from Cloud Tasks. :param location: The location name in which the queue was created. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Queue """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) return client.get_queue( name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def list_queues( self, location, project_id=None, results_filter=None, page_size=None, retry=None, timeout=None, metadata=None, ): """ Lists queues from Cloud Tasks. :param location: The location name in which the queues were created. :type location: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param results_filter: (Optional) Filter used to specify a subset of queues. :type results_filter: str :param page_size: (Optional) The maximum number of resources contained in the underlying API response. :type page_size: int :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: list[google.cloud.tasks_v2.types.Queue] """ client = self.get_conn() full_location_path = CloudTasksClient.location_path(project_id, location) queues = client.list_queues( parent=full_location_path, filter_=results_filter, page_size=page_size, retry=retry, timeout=timeout, metadata=metadata, ) return list(queues)
@GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def delete_queue( self, location, queue_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Deletes a queue from Cloud Tasks, even if it has tasks in it. :param location: The location name in which the queue will be deleted. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) client.delete_queue( name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def purge_queue( self, location, queue_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Purges a queue by deleting all of its tasks from Cloud Tasks. :param location: The location name in which the queue will be purged. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: list[google.cloud.tasks_v2.types.Queue] """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) return client.purge_queue( name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def pause_queue( self, location, queue_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Pauses a queue in Cloud Tasks. :param location: The location name in which the queue will be paused. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: list[google.cloud.tasks_v2.types.Queue] """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) return client.pause_queue( name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def resume_queue( self, location, queue_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Resumes a queue in Cloud Tasks. :param location: The location name in which the queue will be resumed. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: list[google.cloud.tasks_v2.types.Queue] """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) return client.resume_queue( name=full_queue_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def create_task( self, location, queue_name, task, project_id=None, task_name=None, response_view=None, retry=None, timeout=None, metadata=None, ): """ Creates a task in Cloud Tasks. :param location: The location name in which the task will be created. :type location: str :param queue_name: The queue's name. :type queue_name: str :param task: The task to add. If a dict is provided, it must be of the same form as the protobuf message Task. :type task: dict or class google.cloud.tasks_v2.types.Task :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param task_name: (Optional) The task's name. If provided, it will be used to construct the full task path. :type task_name: str :param response_view: (Optional) This field specifies which subset of the Task will be returned. :type response_view: google.cloud.tasks_v2.types.Task.View :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Task """ client = self.get_conn() if task_name: full_task_name = CloudTasksClient.task_path( project_id, location, queue_name, task_name ) if isinstance(task, Task): task.name = full_task_name elif isinstance(task, dict): task['name'] = full_task_name else: raise AirflowException('Unable to set task_name.') full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) return client.create_task( parent=full_queue_name, task=task, response_view=response_view, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def get_task( self, location, queue_name, task_name, project_id=None, response_view=None, retry=None, timeout=None, metadata=None, ): """ Gets a task from Cloud Tasks. :param location: The location name in which the task was created. :type location: str :param queue_name: The queue's name. :type queue_name: str :param task_name: The task's name. :type task_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param response_view: (Optional) This field specifies which subset of the Task will be returned. :type response_view: google.cloud.tasks_v2.types.Task.View :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Task """ client = self.get_conn() full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name) return client.get_task( name=full_task_name, response_view=response_view, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def list_tasks( self, location, queue_name, project_id=None, response_view=None, page_size=None, retry=None, timeout=None, metadata=None, ): """ Lists the tasks in Cloud Tasks. :param location: The location name in which the tasks were created. :type location: str :param queue_name: The queue's name. :type queue_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param response_view: (Optional) This field specifies which subset of the Task will be returned. :type response_view: google.cloud.tasks_v2.types.Task.View :param page_size: (Optional) The maximum number of resources contained in the underlying API response. :type page_size: int :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: list[google.cloud.tasks_v2.types.Task] """ client = self.get_conn() full_queue_name = CloudTasksClient.queue_path(project_id, location, queue_name) tasks = client.list_tasks( parent=full_queue_name, response_view=response_view, page_size=page_size, retry=retry, timeout=timeout, metadata=metadata, ) return list(tasks)
@GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def delete_task( self, location, queue_name, task_name, project_id=None, retry=None, timeout=None, metadata=None, ): """ Deletes a task from Cloud Tasks. :param location: The location name in which the task will be deleted. :type location: str :param queue_name: The queue's name. :type queue_name: str :param task_name: The task's name. :type task_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] """ client = self.get_conn() full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name) client.delete_task( name=full_task_name, retry=retry, timeout=timeout, metadata=metadata
) @GoogleCloudBaseHook.catch_http_exception @GoogleCloudBaseHook.fallback_to_default_project_id
[docs] def run_task( self, location, queue_name, task_name, project_id=None, response_view=None, retry=None, timeout=None, metadata=None, ): """ Forces to run a task in Cloud Tasks. :param location: The location name in which the task was created. :type location: str :param queue_name: The queue's name. :type queue_name: str :param task_name: The task's name. :type task_name: str :param project_id: (Optional) The ID of the GCP project that owns the Cloud Tasks. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param response_view: (Optional) This field specifies which subset of the Task will be returned. :type response_view: google.cloud.tasks_v2.types.Task.View :param retry: (Optional) A retry object used to retry requests. If None is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete. Note that if retry is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: sequence[tuple[str, str]]] :rtype: google.cloud.tasks_v2.types.Task """ client = self.get_conn() full_task_name = CloudTasksClient.task_path(project_id, location, queue_name, task_name) return client.run_task( name=full_task_name, response_view=response_view, retry=retry, timeout=timeout, metadata=metadata,
)