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