def update_queue(
self,
task_queue: Queue,
project_id: Optional[str] = None,
location: Optional[str] = None,
queue_name: Optional[str] = None,
update_mask: Optional[FieldMask] = None,
retry: Optional[Retry] = None,
timeout: Optional[float] = None,
metadata: Optional[Sequence[Tuple[str, str]]] = None) -> Queue:
"""
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,
)
def setUp(self):
self._server = create_server("localhost", 9022)
self._server.start()
time.sleep(1)
transport = CloudTasksGrpcTransport(
channel=grpc.insecure_channel("127.0.0.1:9022"))
self._client = CloudTasksClient(
transport=transport,
client_options=ClientOptions(api_endpoint="127.0.0.1:9022"))
self._parent = self._client.location_path('[PROJECT]', '[LOCATION]')
# Create default queue
self._client.create_queue(self._parent,
{"name": "%s/queues/default" % self._parent})
def get_conn(self) -> CloudTasksClient:
"""
Provides a client for interacting with the Google Cloud Tasks API.
:return: Google Cloud Tasks API Client
:rtype: google.cloud.tasks_v2.CloudTasksClient
"""
if self._client is None:
self._client = CloudTasksClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
return self._client
def test_default_queue_name(self):
server = create_server("localhost", 9023, 10124,
"projects/[P]/locations/[L]/queues/[Q]")
server.start()
time.sleep(1)
transport = CloudTasksGrpcTransport(
channel=grpc.insecure_channel("127.0.0.1:9023"))
client = CloudTasksClient(
transport=transport,
client_options=ClientOptions(api_endpoint="127.0.0.1:9023"))
queues = list(client.list_queues(parent="projects/[P]/locations/[L]"))
self.assertEqual(len(queues), 1)
queue = queues[0]
self.assertEqual(queue.name, "projects/[P]/locations/[L]/queues/[Q]")
server.stop()
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
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.enums.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)
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.enums.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,
)
def list_queues(
self,
location: str,
project_id: str,
results_filter: Optional[str] = None,
page_size: Optional[int] = None,
retry: Optional[Retry] = None,
timeout: Optional[float] = None,
metadata: Optional[Sequence[Tuple[str,
str]]] = None) -> List[Queue]:
"""
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)
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)
def resume_queue(
self,
location: str,
queue_name: str,
project_id: str,
retry: Optional[Retry] = None,
timeout: Optional[float] = None,
metadata: Optional[Sequence[Tuple[str,
str]]] = None) -> List[Queue]:
"""
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)
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.enums.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,
)
def prod_client():
return CloudTasksClient()
def get_client(cls):
if cls.client is None:
cls.client = CloudTasksClient()
return cls.client
class TestCase(BaseTestCase):
def setUp(self):
self._server = create_server("localhost", 9022)
self._server.start()
time.sleep(1)
transport = CloudTasksGrpcTransport(
channel=grpc.insecure_channel("127.0.0.1:9022"))
self._client = CloudTasksClient(
transport=transport,
client_options=ClientOptions(api_endpoint="127.0.0.1:9022"))
self._parent = self._client.location_path('[PROJECT]', '[LOCATION]')
# Create default queue
self._client.create_queue(self._parent,
{"name": "%s/queues/default" % self._parent})
def tearDown(self):
self._server.stop()
def test_create_queue(self):
queue1_path = "%s/queues/test_queue1" % self._parent
queue2_path = "%s/queues/test_queue2" % self._parent
ret = self._client.create_queue(self._parent, {"name": queue1_path})
self.assertEqual(ret.name, queue1_path)
ret = self._client.create_queue(self._parent, {"name": queue2_path})
self.assertEqual(ret.name, queue2_path)
return (queue1_path, queue2_path)
def test_list_queues(self):
path = self._client.queue_path('[PROJECT]', '[LOCATION]', "default")
self._client.delete_queue(path)
self.test_create_queue() # Create a couple of queues
queues = self._client.list_queues(parent=self._parent)
self.assertEqual(len(list(queues)), 2)
def test_get_queue(self):
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
queue = self._client.get_queue(path)
self.assertEqual(queue.name, path)
def test_delete_queue(self):
path = self._client.queue_path('[PROJECT]', '[LOCATION]', "default")
self._client.delete_queue(path)
self.test_create_queue() # Create a couple of queues
queues = self._client.list_queues(parent=self._parent)
self.assertEqual(len(list(queues)), 2)
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
self._client.delete_queue(path)
queues = self._client.list_queues(parent=self._parent)
self.assertEqual(len(list(queues)), 1)
def test_pause_queue(self):
path = self._client.queue_path('[PROJECT]', '[LOCATION]', "default")
self._client.delete_queue(path)
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
response = self._client.pause_queue(path)
self.assertEqual(response.state, 2)
def test_purge_queue(self):
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
# Pause the queue as we don't want tasks to be processed
self._client.pause_queue(path)
payload = "Hello World!"
task = {
'app_engine_http_request': { # Specify the type of request.
'http_method': 'POST',
'relative_uri': '/example_task_handler',
'body': payload.encode()
}
}
# Create 3 tasks
self._client.create_task(path, task)
self._client.create_task(path, task)
self._client.create_task(path, task)
tasks = [x for x in self._client.list_tasks(path)]
self.assertEqual(len(tasks), 3)
# Check again, make sure list_tasks didn't do anything
tasks = [x for x in self._client.list_tasks(path)]
self.assertEqual(len(tasks), 3)
self._client.purge_queue(path)
tasks = [x for x in self._client.list_tasks(path)]
self.assertEqual(len(tasks), 0)
def test_create_task(self):
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
self._client.pause_queue(path)
payload = "Hello World!"
task = {
'app_engine_http_request': { # Specify the type of request.
'http_method': 'POST',
'relative_uri': '/example_task_handler',
'body': payload.encode()
}
}
response = self._client.create_task(path, task)
self.assertTrue(response.name.startswith(path))
def test_run_task(self):
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
self._client.pause_queue(path) # Don't run any tasks while testing
payload = "Hello World!"
task = {
'app_engine_http_request': { # Specify the type of request.
'http_method': 'POST',
'relative_uri': '/example_task_handler',
'body': payload.encode()
}
}
response = self._client.create_task(path, task)
self.assertTrue(response.name.startswith(path))
class FakeResponse:
status = 200
with sleuth.fake("server._make_task_request",
return_value=FakeResponse()):
self._client.run_task(response.name)
# Should return NOT_FOUND
self.assertRaises(
Unknown,
self._client.run_task,
"%s/tasks/1119129292929292929" % path, # Not a valid task
)
def test_default_queue_name(self):
server = create_server("localhost", 9023, 10124,
"projects/[P]/locations/[L]/queues/[Q]")
server.start()
time.sleep(1)
transport = CloudTasksGrpcTransport(
channel=grpc.insecure_channel("127.0.0.1:9023"))
client = CloudTasksClient(
transport=transport,
client_options=ClientOptions(api_endpoint="127.0.0.1:9023"))
queues = list(client.list_queues(parent="projects/[P]/locations/[L]"))
self.assertEqual(len(queues), 1)
queue = queues[0]
self.assertEqual(queue.name, "projects/[P]/locations/[L]/queues/[Q]")
server.stop()
class CustomPortTestCase(BaseTestCase):
@classmethod
def setUpClass(cls):
cls._server = MockServer(10123)
cls._server.start()
cls._server.is_running.wait()
@classmethod
def tearDownClass(cls):
cls._server.join(timeout=1)
def setUp(self):
self._server = create_server("localhost", 9022, 10123)
self._server.start()
time.sleep(1)
transport = CloudTasksGrpcTransport(
channel=grpc.insecure_channel("127.0.0.1:9022"))
self._client = CloudTasksClient(
transport=transport,
client_options=ClientOptions(api_endpoint="127.0.0.1:9022"))
self._parent = self._client.location_path('[PROJECT]', '[LOCATION]')
def tearDown(self):
self._server.stop()
def test_create_queue(self):
path1 = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue1")
ret = self._client.create_queue(self._parent, {"name": path1})
self.assertEqual(ret.name, path1)
path2 = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
ret = self._client.create_queue(self._parent, {"name": path2})
self.assertEqual(ret.name, path2)
def test_run_task(self):
self.test_create_queue() # Create a couple of queues
path = self._client.queue_path('[PROJECT]', '[LOCATION]',
"test_queue2")
self._client.pause_queue(path) # Don't run any tasks while testing
payload = "Hello World!"
task = {
'app_engine_http_request': { # Specify the type of request.
'http_method': 'POST',
'relative_uri': '/example_task_handler',
'body': payload.encode()
}
}
response = self._client.create_task(path, task)
self.assertTrue(response.name.startswith(path))
self._client.run_task(response.name)
def __init__(self):
self._client = CloudTasksClient()
_, self._project_id = auth.default()
def dev_client():
channel = grpc.insecure_channel(CLOUD_TASKS_HOST)
transport = CloudTasksGrpcTransport(channel=channel)
client = CloudTasksClient(transport=transport)
return client
