From 1586c31ccb01623c340a9b53a5982f891abb657d Mon Sep 17 00:00:00 2001 From: Riddhi Shivhare Date: Wed, 7 Oct 2026 17:36:39 +0000 Subject: [PATCH 1/4] feat(taskqueue): route push queues to Cloud Tasks v2 (GA) When APPENGINE_USE_CLOUDTASK_PUSH_QUEUE is enabled, push queue operations in the TaskQueue SDK are served by the Cloud Tasks v2 API: - Queue.add (single and batch) uses CreateTask and BatchCreateTasks. - Queue.delete_tasks uses BatchDeleteTasks. - Queue.purge uses PurgeQueue. - Queue.fetch_statistics uses v2beta3 GetQueue, since queue stats are not part of the v2 API. - Transactional tasks are staged in Datastore and dispatched after commit. Each staged entity is locked in a transaction before dispatch, so the post-commit path and the /_ah/cloudtask/sweep cron handler do not dispatch the same entity concurrently. Each sweep reads at most 500 entities of each status. Failed entities are deleted after 7 days. Like the legacy API, at most 5 tasks can be added in one transaction. Each staged task is its own entity group, so a transaction that adds transactional tasks must be cross-group (xg=True) unless it touches no other entity group. Staged entities use the kind, property names and payload form of the Go and Java SDKs, a JSON object {"task": T} where T is the proto3 JSON form of the v2 Task, so a sweeper in a service written in any of them can dispatch them. Entities staged by the preview release are still dispatched. Batch operations read the synchronous operation result and map failed_requests in the operation metadata back to per-task TaskQueue errors. Reserved headers (Host, X-Google-*, X-AppEngine-*) are not forwarded, and a target without a version routes to that service's default version. Cloud Tasks API errors are raised as the matching taskqueue.Error subclass (for example UnknownQueueError, PermissionDeniedError, InvalidQueueModeError, TransientError). A deadline set on the rpc passed to the *_async methods, or passed to fetch_statistics, limits the Cloud Tasks RPC calls, and the rpc reports the Cloud Tasks result. Without a deadline, batch RPC calls use a 20 second default timeout. Adding, leasing and deleting pull tasks keeps using the legacy service; purge and stats also work for pull queues through Cloud Tasks. Cloud Tasks clients are created once per process and reused. Adds google-cloud-tasks>=2.25.0 as a dependency on Python 3.10+, the versions that release supports. On Python 3.9 the legacy TaskQueue path is unchanged, and enabling the Cloud Tasks backend raises an ImportError that explains the requirement. Unit tests are in tests/google/appengine/api/taskqueue/cloudtask_test.py. --- setup.py | 1 + src/google/appengine/api/__init__.py | 1 + .../appengine/api/taskqueue/cloudtask.py | 925 ++++++++++++ .../api/taskqueue/cloudtask_transactional.py | 506 +++++++ .../appengine/api/taskqueue/taskqueue.py | 53 +- .../appengine/datastore/datastore_rpc.py | 14 + src/google/appengine/runtime/middlewares.py | 20 + .../appengine/api/taskqueue/cloudtask_test.py | 1272 +++++++++++++++++ 8 files changed, 2790 insertions(+), 2 deletions(-) create mode 100644 src/google/appengine/api/taskqueue/cloudtask.py create mode 100644 src/google/appengine/api/taskqueue/cloudtask_transactional.py create mode 100644 tests/google/appengine/api/taskqueue/cloudtask_test.py diff --git a/setup.py b/setup.py index aded2ab..afd8eb0 100644 --- a/setup.py +++ b/setup.py @@ -30,6 +30,7 @@ "six>=1.15.0", "urllib3>=1.26.2,<2", "legacy-cgi>=2.6.2; python_version>='3.10'", + "google-cloud-tasks>=2.25.0; python_version>='3.10'", ], classifiers=[ "Programming Language :: Python :: 3", diff --git a/src/google/appengine/api/__init__.py b/src/google/appengine/api/__init__.py index f0e1a56..c94bfb6 100755 --- a/src/google/appengine/api/__init__.py +++ b/src/google/appengine/api/__init__.py @@ -87,6 +87,7 @@ def if_deferred_enabled(array): ]) + [ middlewares.ErrorLoggingMiddleware, middlewares.BackgroundAndShutdownMiddleware, + middlewares.CloudTaskSweepMiddleware, middlewares.SetNamespaceFromHeader, ] + if_deferred_enabled([ middlewares.AddDeferredMiddleware, diff --git a/src/google/appengine/api/taskqueue/cloudtask.py b/src/google/appengine/api/taskqueue/cloudtask.py new file mode 100644 index 0000000..d2fb62d --- /dev/null +++ b/src/google/appengine/api/taskqueue/cloudtask.py @@ -0,0 +1,925 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +"""Cloud Tasks backend integration for Taskqueue SDK.""" + +import base64 +from concurrent import futures +import contextvars +import datetime +import functools +import http +import json +import os +import re +import threading +import time +import urllib.error +import urllib.request +from google.appengine.api import app_identity +from google.appengine.api.taskqueue import taskqueue +from google.appengine.api.taskqueue import taskqueue_service_bytes_pb2 as taskqueue_service_pb2 +from google.protobuf import duration_pb2 +from google.protobuf import field_mask_pb2 +from google.protobuf.timestamp_pb2 import Timestamp + +# google-cloud-tasks>=2.25.0 (the first release with the v2 batch APIs) +# requires Python 3.10+, so it is only installed there. Import it lazily so +# the legacy TaskQueue path keeps working when it is not available. +try: + from google.api_core import exceptions as google_exceptions + from google.api_core import retry as google_retry + from google.cloud import tasks_v2 + from google.cloud import tasks_v2beta3 + from google.rpc import code_pb2 + _CLOUD_TASKS_IMPORT_ERROR = None +except ImportError as e: + google_exceptions = google_retry = tasks_v2 = tasks_v2beta3 = code_pb2 = None + _CLOUD_TASKS_IMPORT_ERROR = e + +# Environment variable constants +ENV_USE_CLOUDTASK_PUSH_QUEUE = 'APPENGINE_USE_CLOUDTASK_PUSH_QUEUE' +ENV_LOCATION_ID = 'LOCATION_ID' +ENV_GAE_LOCATION = 'GAE_LOCATION' +ENV_GAE_REGION = 'GAE_REGION' +ENV_REGION_ID = 'REGION_ID' +ENV_LOCAL_GCP_REGION = 'LOCAL_GCP_REGION' +ENV_GOOGLE_CLOUD_PROJECT = 'GOOGLE_CLOUD_PROJECT' +ENV_GAE_SERVICE = 'GAE_SERVICE' +ENV_GAE_VERSION = 'GAE_VERSION' + +_LEGACY_REGION_MAP = { + 'us-central': 'us-central1', + 'europe-west': 'europe-west1', +} + +# Sizing & concurrency constants +_MAX_CONCURRENT_API_CALLS = 100 +_BATCH_CREATE_TASKS_MAX_SIZE = 100 +_BATCH_DELETE_TASKS_MAX_SIZE = 1000 +_METADATA_SERVER_TIMEOUT_SECONDS = 2 +# Default RPC timeout for batch calls when the caller sets no deadline, and +# the upper bound on caller-supplied RPC deadlines. The client library sets +# 20s for non-batch RPCs (and no default for batch RPCs) because the Cloud +# Tasks service rejects request deadlines more than 30s in the future. +_DEFAULT_RPC_TIMEOUT_SECONDS = 20 + +# Cloud Tasks error details carry the backend TaskQueue error name, e.g. +# "apphosting::ExecutorServiceError::UNKNOWN_QUEUE". +_EXECUTOR_ERROR_RE = re.compile(r'ExecutorServiceError::([A-Z_]+)') + +_THREAD_POOL = futures.ThreadPoolExecutor(_MAX_CONCURRENT_API_CALLS) + +# Cloud Tasks clients are expensive to create (credential lookup and gRPC +# channel setup), so one client per API version is created lazily and reused. +# GAPIC clients are thread-safe. The cache is keyed by process ID so that a +# forked worker never reuses a gRPC channel created in its parent process. +_CLIENT_LOCK = threading.Lock() +_CLIENTS = {} +_CACHED_REGION = None + + +def _to_taskqueue_error(error): + """Maps a Cloud Tasks API error to the TaskQueue error legacy would raise. + + Returns error unchanged if it is already a taskqueue.Error or is not a + Cloud Tasks API error. + """ + if isinstance(error, taskqueue.Error): + return error + detail = str(error) + # The client library only retries transient errors, so running out of + # retries is a transient error too. + if (google_exceptions is not None + and isinstance(error, google_exceptions.RetryError)): + return taskqueue.TransientError(detail) + if (google_exceptions is None + or not isinstance(error, google_exceptions.GoogleAPICallError)): + return error + + match = _EXECUTOR_ERROR_RE.search(detail) + error_codes = taskqueue_service_pb2.TaskQueueServiceError.ErrorCode + if match and match.group(1) in error_codes.keys(): + code = error_codes.Value(match.group(1)) + if code in taskqueue._ERROR_MAPPING: + return taskqueue._TranslateError(code, detail) + + # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. + if isinstance(error, (google_exceptions.Aborted, + google_exceptions.ServiceUnavailable, + google_exceptions.DeadlineExceeded, + google_exceptions.TooManyRequests)): + return taskqueue.TransientError(detail) + if isinstance(error, (google_exceptions.AlreadyExists, + google_exceptions.Conflict)): + return taskqueue.TaskAlreadyExistsError(detail) + if isinstance(error, google_exceptions.NotFound): + return taskqueue.UnknownQueueError(detail) + if isinstance(error, (google_exceptions.PermissionDenied, + google_exceptions.Unauthenticated)): + return taskqueue.PermissionDeniedError(detail) + if isinstance(error, google_exceptions.InternalServerError): + return taskqueue.InternalError(detail) + return taskqueue.Error(detail) + + +def _translate_errors(fn): + """Re-raises Cloud Tasks API errors from fn as taskqueue.Error subclasses.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except Exception as e: # pylint: disable=broad-except + mapped = _to_taskqueue_error(e) + if mapped is e: + raise + raise mapped from e + + return wrapper + + +# ============================================================================== +# Public APIs +# ============================================================================== + + +def is_cloudtask_push_queue_enabled(): + """Checks if Cloud Tasks backend is enabled for Push Queues. + + Raises: + ImportError: If the backend is enabled but google-cloud-tasks is not + installed. + """ + enabled = ( + str(os.environ.get(ENV_USE_CLOUDTASK_PUSH_QUEUE, '')).lower() == 'true') + if enabled and _CLOUD_TASKS_IMPORT_ERROR is not None: + raise ImportError( + '%s is set, but google-cloud-tasks>=2.25.0 could not be imported. ' + 'Cloud Tasks push queues require Python 3.10 or later.' + % ENV_USE_CLOUDTASK_PUSH_QUEUE) from _CLOUD_TASKS_IMPORT_ERROR + return enabled + + +@_translate_errors +def create_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): + """Creates one or more tasks using Cloud Tasks API (supporting BatchCreateTasks).""" + for task in tasks: + if task.was_enqueued: + raise taskqueue.BadTaskStateError('The task has already been enqueued.') + deadline_at = _deadline_at(deadline) + if len(tasks) == 1: + return _create_single_task_in_cloud_tasks( + queue_name, tasks[0], multiple, deadline_at) + else: + return _create_batch_tasks_in_cloud_tasks( + queue_name, tasks, multiple, deadline_at) + + +@_translate_errors +def delete_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): + """Deletes tasks from a queue using Cloud Tasks Client SDK (supporting BatchDeleteTasks).""" + deadline_at = _deadline_at(deadline) + client = _get_client() + project = _get_project_id() + region = _get_region() + + parent = client.queue_path(project, region, queue_name) + + # Check pre-conditions (duplicate names or already deleted) + task_names_set = set() + for task in tasks: + if not task.name: + raise taskqueue.BadTaskStateError('A task name must be specified for a task') + if task.was_deleted: + raise taskqueue.BadTaskStateError( + 'The task %s has already been deleted' % task.name + ) + if task.name in task_names_set: + raise taskqueue.DuplicateTaskNameError( + 'The task name %s is duplicated' % task.name + ) + task_names_set.add(task.name) + + def _delete_batch_chunk(batch): + task_names = [ + client.task_path(project, region, queue_name, t.name) for t in batch + ] + op = client.batch_delete_tasks( + request={'parent': parent, 'names': task_names}, + **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) + ) + _extract_lro_response(op) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', None), + ) + if metadata + else None + ) + + chunk_exception = None + for idx, t in enumerate(batch): + error_status = _get_failed_request(failed_requests, idx) + if error_status: + code = getattr(error_status, 'code', None) + tq_code = _map_rest_code_to_tq_code(code) + if tq_code in [ + taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK, + taskqueue_service_pb2.TaskQueueServiceError.TOMBSTONED_TASK, + ]: + t._Task__deleted = False + elif chunk_exception is None: + chunk_exception = taskqueue._TranslateError( + tq_code, getattr(error_status, 'message', '') + ) + else: + t._Task__deleted = True + + return chunk_exception + + chunks = [ + tasks[i : i + _BATCH_DELETE_TASKS_MAX_SIZE] + for i in range(0, len(tasks), _BATCH_DELETE_TASKS_MAX_SIZE) + ] + exception = None + if len(chunks) == 1: + exception = _delete_batch_chunk(chunks[0]) + elif chunks: + chunk_futures = [ + _THREAD_POOL.submit( + contextvars.copy_context().run, _delete_batch_chunk, batch + ) + for batch in chunks + ] + futures.wait(chunk_futures) + for fut in chunk_futures: + chunk_exc = fut.result() + if exception is None and chunk_exc is not None: + exception = chunk_exc + + if exception is not None: + raise exception + + if multiple: + return tasks + else: + return tasks[0] + + +@_translate_errors +def purge_queue_in_cloud_tasks(queue_name): + """Purges all tasks in a queue using Cloud Tasks API.""" + client = _get_client() + project = _get_project_id() + region = _get_region() + + name = client.queue_path(project, region, queue_name) + try: + client.purge_queue(request={'name': name}) + except Exception as e: + raise e + + +@_translate_errors +def fetch_queue_stats_in_cloud_tasks(queues, multiple, deadline=None): + """Fetches queue statistics for given queues using Cloud Tasks API (v2beta3).""" + deadline_at = _deadline_at(deadline) + # QueueStats is retained on v2beta3 as it is out of scope for v2 GA + client = _get_v2beta3_client() + project = _get_project_id() + region = _get_region() + + queue_stats_list = [] + read_mask = field_mask_pb2.FieldMask(paths=['stats']) + + for queue in queues: + queue_name = queue.name if hasattr(queue, 'name') else str(queue) + name = client.queue_path(project, region, queue_name) + try: + call_options = _timeout_kwargs(deadline_at) + if call_options: + call_options['retry'] = _get_queue_retry(call_options['timeout']) + q_resp = client.get_queue( + request={'name': name, 'read_mask': read_mask}, **call_options) + ct_stats = getattr(q_resp, 'stats', None) + + tasks = getattr(ct_stats, 'tasks_count', 0) if ct_stats else 0 + oldest_eta_usec = None + if ct_stats and getattr(ct_stats, 'oldest_estimated_arrival_time', None): + oldest_eta = ct_stats.oldest_estimated_arrival_time + oldest_eta_usec = int(oldest_eta.timestamp() * 1e6) + + executed_last_minute = getattr(ct_stats, 'executed_last_minute_count', 0) if ct_stats else 0 + in_flight = getattr(ct_stats, 'concurrent_dispatches_count', 0) if ct_stats else 0 + enforced_rate = getattr(ct_stats, 'effective_execution_rate', 0.0) if ct_stats else 0.0 + + qs = taskqueue.QueueStatistics( + queue=queue, + tasks=tasks, + oldest_eta_usec=oldest_eta_usec, + executed_last_minute=executed_last_minute, + in_flight=in_flight, + enforced_rate=enforced_rate, + ) + queue_stats_list.append(qs) + except google_exceptions.NotFound as e: + raise taskqueue.UnknownQueueError(f'Queue {queue_name} not found: {e}') + except Exception as e: + raise e + + if multiple: + return queue_stats_list + else: + return queue_stats_list[0] if queue_stats_list else None + + +# ============================================================================== +# Private Helpers +# ============================================================================== + + +class _CloudTaskRPC(object): + """RPC object wrapping asynchronous execution via ThreadPoolExecutor. + + Matches the async model of apiproxy_rpc.py using a ThreadPoolExecutor. + Calls are scheduled onto background threads asynchronously, allowing + operations to run concurrently until .get_result() or .wait() is called. + """ + + def __init__(self, future_or_callable, callback=None): + if callable(future_or_callable): + self._future = _THREAD_POOL.submit(future_or_callable) + else: + self._future = future_or_callable + self.callback = callback + self._callback_called = False + + def get_result(self): + self.wait() + return self._future.result() + + def wait(self): + """Waits for completion and runs the callback once, like UserRPC.wait.""" + futures.wait([self._future]) + if self.callback is not None and not self._callback_called: + self._callback_called = True + self.callback() + + def check_success(self): + self.wait() + self._future.result() + + @property + def future(self): + return self._future + + +class _DeadlineOnlyRPC(object): + """Carries a deadline to _make_rpc in place of a UserRPC. + + Used by the synchronous methods that take a deadline argument. Creating a + UserRPC requires the taskqueue API stub, which Cloud Tasks mode does not + otherwise need. + """ + + def __init__(self, deadline): + self.deadline = deadline + self.callback = None + + +def _make_rpc(fn, rpc=None): + """Runs fn in the background and returns an RPC-like object for it. + + If the caller supplied a UserRPC, it is returned with wait, check_success + and get_result bound to the Cloud Tasks call, and its callback is run on + wait like the legacy path. UserRPC.wait_any and wait_all are not supported. + """ + ct_rpc = _CloudTaskRPC(fn, callback=getattr(rpc, 'callback', None)) + if rpc is None: + return ct_rpc + rpc.wait = ct_rpc.wait + rpc.check_success = ct_rpc.check_success + rpc.get_result = ct_rpc.get_result + return rpc + + +def _deadline_at(deadline): + """Returns the monotonic time at which a caller's RPC deadline expires. + + Returns None if the caller did not set a deadline. + """ + return None if deadline is None else time.monotonic() + deadline + + +def _timeout_kwargs(deadline_at, default=None): + """Returns the timeout argument for a Cloud Tasks call. + + The timeout is the time left before deadline_at, capped at + _DEFAULT_RPC_TIMEOUT_SECONDS because Cloud Tasks rejects request deadlines + more than 30s in the future. If no deadline is set, default is used, or if + default is None no timeout is passed and the client's 20s default applies. + """ + if deadline_at is None: + return {} if default is None else {'timeout': default} + remaining = max(deadline_at - time.monotonic(), 0.001) + return {'timeout': min(remaining, _DEFAULT_RPC_TIMEOUT_SECONDS)} + + +def _get_queue_retry(timeout): + """Returns the client's default GetQueue retry, limited to timeout seconds. + + A timeout alone bounds only each attempt; the default retry keeps retrying + DEADLINE_EXCEEDED for up to 20 seconds. + """ + return google_retry.Retry( + initial=0.1, + maximum=10.0, + multiplier=1.3, + predicate=google_retry.if_exception_type( + google_exceptions.DeadlineExceeded, + google_exceptions.ServiceUnavailable, + ), + deadline=timeout, + ) + + +def _get_cached_client(api_version, factory): + """Returns the cached client for api_version, creating it on first use.""" + key = (api_version, os.getpid()) + client = _CLIENTS.get(key) + if client is None: + with _CLIENT_LOCK: + client = _CLIENTS.get(key) + if client is None: + client = factory() + _CLIENTS[key] = client + return client + + +def _get_client(): + """Returns the shared Cloud Tasks v2 client.""" + return _get_cached_client('v2', lambda: tasks_v2.CloudTasksClient()) + + +def _get_v2beta3_client(): + """Returns the shared Cloud Tasks v2beta3 client, used only for QueueStats.""" + return _get_cached_client( + 'v2beta3', lambda: tasks_v2beta3.CloudTasksClient() + ) + + +def _reset_clients(): + """Drops cached clients and region. Intended for tests.""" + global _CACHED_REGION + with _CLIENT_LOCK: + _CLIENTS.clear() + _CACHED_REGION = None + + +def _get_project_id(): + """Extracts and formats the Google Cloud project ID.""" + project = os.environ.get(ENV_GOOGLE_CLOUD_PROJECT) + if project and (project.startswith('s~') or project.startswith('e~')): + project = project[2:] + return project + + +def _normalize_region(region): + """Maps legacy App Engine region names to Cloud Tasks locations.""" + if not region: + return region + region = region.strip() + return _LEGACY_REGION_MAP.get(region, region) + + +def _get_region(): + """Determines the App Engine region.""" + global _CACHED_REGION + region = ( + os.environ.get(ENV_LOCATION_ID) + or os.environ.get(ENV_GAE_LOCATION) + or os.environ.get(ENV_GAE_REGION) + or os.environ.get(ENV_REGION_ID) + ) + if region: + return _normalize_region(region) + + if _CACHED_REGION: + return _CACHED_REGION + + try: + req = urllib.request.Request( + 'http://metadata.google.internal/computeMetadata/v1/instance/region', + headers={'Metadata-Flavor': 'Google'}, + ) + with urllib.request.urlopen(req, timeout=_METADATA_SERVER_TIMEOUT_SECONDS) as response: + region_path = response.read().decode('utf-8') + resolved = _normalize_region(region_path.split('/')[-1]) + if resolved: + _CACHED_REGION = resolved + return resolved + except Exception: + pass + + # Fallback to us-central1 if we can't detect it; do not cache fallback. + return _normalize_region(os.environ.get(ENV_LOCAL_GCP_REGION, 'us-central1')) + + +def _to_duration(seconds): + if seconds is None: + return None + duration = duration_pb2.Duration() + # Rounding avoids float error, e.g. 2.3 seconds becoming 2s 299999999ns. + duration.FromNanoseconds(int(round(seconds * 1e9))) + return duration + + +def _build_retry_config(retry_options): + if not retry_options: + return None + + config = {} + + if retry_options.task_retry_limit is not None: + config['max_attempts'] = retry_options.task_retry_limit + 1 + if retry_options.task_age_limit is not None: + config['max_retry_duration'] = _to_duration(retry_options.task_age_limit) + if retry_options.min_backoff_seconds is not None: + config['min_backoff'] = _to_duration(retry_options.min_backoff_seconds) + if retry_options.max_backoff_seconds is not None: + config['max_backoff'] = _to_duration(retry_options.max_backoff_seconds) + if retry_options.max_doublings is not None: + config['max_doublings'] = retry_options.max_doublings + + if config: + return config + return None + + +def _build_ct_task_payload(queue_name, task, client, project, region): + """Builds the Cloud Tasks Task proto payload from GAE Task.""" + # Cloud Tasks does not allow Host, X-Google-* or X-AppEngine-* headers to + # be set; it sets the App Engine headers (queue name, task name, etc.) + # itself when the task is dispatched. + headers = {} + for key, value in (task.headers or {}).items(): + lower_key = key.lower() + if (lower_key == 'host' or lower_key.startswith('x-google-') + or lower_key.startswith('x-appengine-')): + continue + headers[key] = value + + body = b'' + if task.payload: + if isinstance(task.payload, str): + body = task.payload.encode('utf-8') + else: + body = task.payload + + http_method = tasks_v2.HttpMethod.POST + if task.method: + method_map = { + 'POST': tasks_v2.HttpMethod.POST, + 'GET': tasks_v2.HttpMethod.GET, + 'PUT': tasks_v2.HttpMethod.PUT, + 'DELETE': tasks_v2.HttpMethod.DELETE, + 'HEAD': tasks_v2.HttpMethod.HEAD, + } + http_method = method_map.get(task.method, tasks_v2.HttpMethod.POST) + + app_engine_http_request = { + 'http_method': http_method, + 'relative_uri': task.url or '/', + 'body': body, + 'headers': headers, + } + + routing = {} + current_service = os.environ.get(ENV_GAE_SERVICE) + current_version = os.environ.get(ENV_GAE_VERSION) + if isinstance(task.target, str) and task.target: + # A target derived from a "-dot-" hostname looks like "v1-dot-worker-dot". + # If a full hostname was passed as target, strip the project/domain suffix. + target_str = task.target.split(':', 1)[0] + app_host = project + if project and ':' in project: + domain, app_id = project.split(':', 1) + app_host = f'{app_id}.{domain}' + + matched_other_project = False + if target_str.endswith('.appspot.com'): + host_body = target_str[: -len('.appspot.com')] + if host_body.endswith('.r'): + r_dot = host_body[:-2].rfind('.') + if r_dot != -1: + host_body = host_body[:r_dot] + normalized_host = host_body.replace('-dot-', '.') + if app_host: + if normalized_host == app_host: + target_str = '' + elif normalized_host.endswith('.' + app_host): + target_str = normalized_host[: -(len(app_host) + 1)] + else: + matched_other_project = True + else: + parts = normalized_host.rsplit('.', 1) + target_str = parts[0] if len(parts) > 1 else '' + elif app_host: + for sep in ('-dot-' + app_host + '.', '.' + app_host + '.'): + if sep in target_str: + target_str = target_str[: target_str.rfind(sep)] + break + else: + if target_str.startswith(app_host + '.'): + target_str = '' + + if matched_other_project: + target_service = current_service + target_version = current_version + target_instance = None + else: + if target_str.endswith('-dot'): + target_str = target_str[:-4] + if not target_str: + target_service = 'default' + target_version = None + target_instance = None + else: + target_components = target_str.replace('-dot-', '.').rsplit('.', 3) + target_service = target_components[-1] + target_version = ( + len(target_components) > 1 and target_components[-2] or None + ) + target_instance = ( + len(target_components) > 2 and target_components[-3] or None + ) + + if target_service: + routing['service'] = target_service + if target_version: + routing['version'] = target_version + elif current_version and target_service == current_service: + # Like legacy push queues, a task for the current service runs on the + # version that enqueued it. Other services use their default version. + routing['version'] = current_version + if target_instance: + routing['instance'] = target_instance + else: + if current_service: + routing['service'] = current_service + if current_version: + routing['version'] = current_version + + if routing: + app_engine_http_request['app_engine_routing'] = routing + + ct_task = {'app_engine_http_request': app_engine_http_request} + + if task.name: + ct_task['name'] = client.task_path(project, region, queue_name, task.name) + + if task.eta: + epoch = datetime.datetime.utcfromtimestamp(0) + eta = task.eta + if eta.tzinfo is not None: + eta = eta.astimezone(datetime.timezone.utc).replace(tzinfo=None) + delta = eta - epoch + seconds = int(delta.total_seconds()) + nanos = int(delta.microseconds * 1000) + timestamp = Timestamp(seconds=seconds, nanos=nanos) + ct_task['schedule_time'] = timestamp + + if task.retry_options: + retry_config = _build_retry_config(task.retry_options) + if retry_config: + ct_task['retry_config'] = retry_config + + return ct_task + + +def _create_single_task_in_cloud_tasks(queue_name, task, multiple, + deadline_at=None): + """Helper to create a single task using CloudTasksClient CreateTask API.""" + client = _get_client() + project = _get_project_id() + region = _get_region() + + parent = client.queue_path(project, region, queue_name) + ct_task = _build_ct_task_payload(queue_name, task, client, project, region) + + try: + response_task = client.create_task( + request={'parent': parent, 'task': ct_task}, + **_timeout_kwargs(deadline_at)) + task_id = response_task.name.split('/')[-1] + task._Task__name = task_id + task._Task__queue_name = queue_name + task._Task__enqueued = True + if multiple: + return [task] + else: + return task + except google_exceptions.Aborted: + # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. + raise + except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: + raise taskqueue.TaskAlreadyExistsError(str(e)) + except google_exceptions.NotFound as e: + raise taskqueue.UnknownQueueError(str(e)) + except google_exceptions.BadRequest as e: + if 'Queue does not exist' in str(e): + raise taskqueue.UnknownQueueError(str(e)) + raise e + except Exception as e: + raise e + + +def _extract_lro_response(op): + """Extracts the unpacked response message from a synchronous LRO.""" + if hasattr(op, 'result') and callable(op.result): + if (not isinstance(op, futures.Future) + and getattr(op, 'response', None) is not None): + return op.response + return op.result() + return getattr(op, 'response', None) + + +def _set_preferred_exception(current, candidate): + """Prefers non-duplicate errors over TaskAlreadyExists/TombstonedTask errors.""" + if candidate is None: + return current + if (current is None + or isinstance(current, (taskqueue.TaskAlreadyExistsError, + taskqueue.TombstonedTaskError))): + return candidate + return current + + +def _create_batch_tasks_in_cloud_tasks(queue_name, tasks, multiple, + deadline_at=None): + """Helper to create tasks in batches using CloudTasksClient BatchCreateTasks API.""" + client = _get_client() + project = _get_project_id() + region = _get_region() + + parent = client.queue_path(project, region, queue_name) + + # Check pre-conditions + task_names = set() + for task in tasks: + if task.name: + if task.name in task_names: + raise taskqueue.DuplicateTaskNameError( + 'The task name %s is duplicated' % task.name + ) + task_names.add(task.name) + + def _create_batch_chunk(batch): + requests_payload = [ + { + 'parent': parent, + 'task': _build_ct_task_payload( + queue_name, t, client, project, region + ), + } + for t in batch + ] + try: + op = client.batch_create_tasks( + request={'parent': parent, 'requests': requests_payload}, + **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) + ) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', None), + ) + if metadata + else None + ) + response = _extract_lro_response(op) + response_tasks = getattr(response, 'tasks', []) if response else [] + + res_iter = iter(response_tasks) + chunk_created = [] + chunk_exception = None + missing_task_idx = None + + for idx, t in enumerate(batch): + error_status = _get_failed_request(failed_requests, idx) + if error_status: + code = getattr(error_status, 'code', None) + tq_code = _map_rest_code_to_tq_code(code) + err = taskqueue._TranslateError( + tq_code, getattr(error_status, 'message', '') + ) + chunk_exception = _set_preferred_exception(chunk_exception, err) + else: + try: + res_task = next(res_iter) + task_id = ( + res_task.name.split('/')[-1] + if hasattr(res_task, 'name') + else res_task['name'].split('/')[-1] + ) + t._Task__name = task_id + t._Task__queue_name = queue_name + t._Task__enqueued = True + chunk_created.append(t) + except StopIteration: + if missing_task_idx is None: + missing_task_idx = idx + + if chunk_exception is None and missing_task_idx is not None: + chunk_exception = taskqueue.InternalError( + 'BatchCreateTasks response is missing created task at index %d' + % missing_task_idx + ) + + return chunk_created, chunk_exception + except google_exceptions.Aborted: + # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. + raise + except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: + raise taskqueue.TaskAlreadyExistsError(str(e)) + except google_exceptions.NotFound as e: + raise taskqueue.UnknownQueueError(str(e)) + except google_exceptions.BadRequest as e: + if 'Queue does not exist' in str(e): + raise taskqueue.UnknownQueueError(str(e)) + raise e + + chunks = [ + tasks[i : i + _BATCH_CREATE_TASKS_MAX_SIZE] + for i in range(0, len(tasks), _BATCH_CREATE_TASKS_MAX_SIZE) + ] + created_tasks = [] + exception = None + if len(chunks) == 1: + chunk_created, exception = _create_batch_chunk(chunks[0]) + created_tasks.extend(chunk_created) + elif chunks: + chunk_futures = [ + _THREAD_POOL.submit( + contextvars.copy_context().run, _create_batch_chunk, batch + ) + for batch in chunks + ] + futures.wait(chunk_futures) + for fut in chunk_futures: + chunk_created, chunk_exc = fut.result() + created_tasks.extend(chunk_created) + exception = _set_preferred_exception(exception, chunk_exc) + + if exception is not None: + raise exception + + if multiple: + return created_tasks + else: + return created_tasks[0] + + +def _get_failed_request(failed_requests, idx): + """Safely retrieves a failed request status by index from proto map or dict.""" + if not failed_requests: + return None + if idx in failed_requests: + return failed_requests[idx] + try: + if str(idx) in failed_requests: + return failed_requests[str(idx)] + except TypeError: + pass + return None + + +def _map_rest_code_to_tq_code(code): + """Maps gRPC / HTTP error status codes to legacy TaskQueue error enum codes.""" + if code in [code_pb2.NOT_FOUND, http.HTTPStatus.NOT_FOUND]: + return taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK + if code in [code_pb2.INVALID_ARGUMENT, http.HTTPStatus.BAD_REQUEST]: + # INVALID_ARGUMENT covers many validation failures, not just task names. + return taskqueue_service_pb2.TaskQueueServiceError.INVALID_REQUEST + if code in [code_pb2.ALREADY_EXISTS, http.HTTPStatus.CONFLICT]: + return taskqueue_service_pb2.TaskQueueServiceError.TASK_ALREADY_EXISTS + if code in [code_pb2.PERMISSION_DENIED, http.HTTPStatus.FORBIDDEN]: + return taskqueue_service_pb2.TaskQueueServiceError.PERMISSION_DENIED + return taskqueue_service_pb2.TaskQueueServiceError.INTERNAL_ERROR diff --git a/src/google/appengine/api/taskqueue/cloudtask_transactional.py b/src/google/appengine/api/taskqueue/cloudtask_transactional.py new file mode 100644 index 0000000..602c831 --- /dev/null +++ b/src/google/appengine/api/taskqueue/cloudtask_transactional.py @@ -0,0 +1,506 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +"""Cloud Tasks transactional task support for Taskqueue SDK.""" + +import base64 +from concurrent import futures +import contextlib +import contextvars +import datetime +import json +import logging +import os +import uuid + +from google.appengine.api import datastore +from google.appengine.api import datastore_errors +from google.appengine.api import datastore_types +from google.appengine.api.taskqueue import cloudtask +from google.appengine.api.taskqueue import taskqueue +from google.appengine.api.taskqueue import taskqueue_service_bytes_pb2 as taskqueue_service_pb2 +from google.appengine.datastore import datastore_pb +from google.protobuf import duration_pb2 +from google.protobuf import json_format +from google.protobuf.timestamp_pb2 import Timestamp + +# Only available with google-cloud-tasks (Python 3.10+). See cloudtask.py. +google_exceptions = cloudtask.google_exceptions +tasks_v2 = cloudtask.tasks_v2 + +try: + from google.appengine.ext import ndb +except ImportError: + ndb = None + +# Constants +# The kind and property names of staged tasks are shared with the Go and Java +# App Engine SDKs, so a sweeper in a service written in any of them can +# dispatch the task. +_PENDING_TASK_KIND = '_AE_PendingCloudTask' +_SDK_LANG = 'PYTHON' +_TX_TASK_STATUS_PENDING = 'PENDING' +_TX_TASK_STATUS_PROCESSING = 'PROCESSING' +_TX_TASK_STATUS_DONE = 'DONE' +_TX_TASK_STATUS_FAILED = 'FAILED' + +_SWEEPER_MAX_RETRIES = 5 +# Maximum entities of each status read by one sweep. Later cron runs pick up +# the rest of a backlog. +_SWEEPER_BATCH_SIZE = 500 +# Like the legacy TaskQueue API, at most 5 tasks can be added in one +# transaction. +_MAX_TASKS_PER_TRANSACTION = 5 +_SWEEPER_LOCK_TIMEOUT_SECONDS = 60 +_SWEEPER_FAST_PATH_GRACE_SECONDS = 60 +# Entities that exhausted their retries are kept this long for inspection, +# then deleted by the sweeper. +_FAILED_TASK_RETENTION = datetime.timedelta(days=7) + + +# ============================================================================== +# Public APIs +# ============================================================================== + + +def add_transactional_tasks(queue_name, tasks, multiple): + """Stages transactional tasks in Datastore within the active transaction. + + Each task is staged as a root entity, i.e. its own entity group, so the + transaction must be cross-group (db: CreateTransactionOptions(xg=True); + ndb: @ndb.transactional(xg=True)) unless it touches no other entity group. + Staging N tasks consumes N of Datastore's 25 entity groups per cross-group + transaction (up to _MAX_TASKS_PER_TRANSACTION = 5 entity groups). + """ + # Check pre-conditions (duplicate names or already queued) + for task in tasks: + if task.name: + raise taskqueue.InvalidTaskNameError( + 'A task bound to a transaction cannot be named.' + ) + if task.was_enqueued: + raise taskqueue.BadTaskStateError('The task has already been enqueued.') + + if not (ndb and ndb.in_transaction()) and not datastore.IsInTransaction(): + raise taskqueue.BadTransactionStateError( + 'Transactional tasks must be added inside a transaction.' + ) + + pending_keys = _pending_keys_for_current_transaction() + if len(pending_keys) + len(tasks) > _MAX_TASKS_PER_TRANSACTION: + # Same error the legacy API raises from the datastore commit. + raise taskqueue._TranslateError( + taskqueue_service_pb2.TaskQueueServiceError.DATASTORE_ERROR + + datastore_pb.Error.BAD_REQUEST, + 'Too many messages, maximum allowed %d' % _MAX_TASKS_PER_TRANSACTION) + + for task in tasks: + task_uuid = uuid.uuid4().hex + generated_name = f"tx-{task_uuid}" + task._Task__name = generated_name + task._Task__queue_name = queue_name + + ct_task_payload = build_task_payload_for_transactional_task( + queue_name, task + ) + + entity = datastore.Entity(_PENDING_TASK_KIND) + entity['queue_name'] = queue_name + entity['cloud_task_name'] = generated_name + # Text is unindexed; indexed strings are limited to 1500 bytes and task + # payloads can be up to 100KB. + entity['cloud_task_payload'] = datastore_types.Text( + _encode_task_payload(ct_task_payload)) + entity['status'] = _TX_TASK_STATUS_PENDING + entity['created'] = datetime.datetime.utcnow() + entity['retry_count'] = 0 + entity['sdk_lang'] = _SDK_LANG + + with _use_default_datastore_adapter(): + datastore.Put(entity) + pending_keys.append(entity.key()) + task._Task__enqueued = True + + if multiple: + return tasks + else: + return tasks[0] + + +def build_task_payload_for_transactional_task(queue_name, task): + """Builds the Cloud Tasks task payload for a transactional task.""" + client = cloudtask._get_client() + project = cloudtask._get_project_id() + region = cloudtask._get_region() + + return cloudtask._build_ct_task_payload(queue_name, task, client, project, region) + + +def dispatch_task_payload(queue_name, task): + """Creates a staged task, a `tasks_v2.Task`, in Cloud Tasks.""" + client = cloudtask._get_client() + project = cloudtask._get_project_id() + region = cloudtask._get_region() + + parent = client.queue_path(project, region, queue_name) + client.create_task(request={'parent': parent, 'task': task}) + + +def _encode_task_payload(task_payload): + """Returns the staged form of a Cloud Tasks task payload dict. + + The form is shared with the Go and Java SDKs: the JSON object {"task": T}, + where T is the proto3 JSON form of the Cloud Tasks v2 Task. + """ + task = tasks_v2.Task(task_payload) + return json.dumps({'task': json_format.MessageToDict(tasks_v2.Task.pb(task))}) + + +def _decode_task_payload(payload_str): + """Parses a payload written by _encode_task_payload or another SDK.""" + obj = json.loads(payload_str) + if isinstance(obj.get('task'), dict): + obj = obj['task'] + task_pb = tasks_v2.Task.pb()() + json_format.ParseDict(obj, task_pb, ignore_unknown_fields=True) + return tasks_v2.Task.wrap(task_pb) + + +def _decode_preview_payload(payload_str): + """Parses the payload of a task staged by the preview release. + + The preview stored a dict of Task fields in the 'payload' property, with + the body base64 encoded and times as {'seconds': s, 'nanos': n}. + """ + task_payload = json.loads(payload_str) + if 'app_engine_http_request' in task_payload: + ae_req = dict(task_payload['app_engine_http_request']) + if 'body' in ae_req and isinstance(ae_req['body'], str): + ae_req['body'] = base64.b64decode(ae_req['body'].encode('utf-8')) + task_payload['app_engine_http_request'] = ae_req + + if 'schedule_time' in task_payload and isinstance(task_payload['schedule_time'], dict): + st = task_payload['schedule_time'] + task_payload['schedule_time'] = Timestamp(seconds=st.get('seconds', 0), nanos=st.get('nanos', 0)) + + if 'retry_config' in task_payload and isinstance(task_payload['retry_config'], dict): + rc = dict(task_payload['retry_config']) + for dur_field in ('max_retry_duration', 'min_backoff', 'max_backoff'): + if dur_field in rc and isinstance(rc[dur_field], dict): + d = rc[dur_field] + rc[dur_field] = duration_pb2.Duration( + seconds=d.get('seconds', 0), nanos=d.get('nanos', 0) + ) + task_payload['retry_config'] = rc + + return tasks_v2.Task(task_payload) + + +def sweep(): + """Queries Datastore for pending Cloud Tasks and dispatches them.""" + # Each status is read separately with a limit, so that memory use is + # bounded and failed entities kept for inspection cannot crowd out + # pending ones. + entities = [] + for status in (_TX_TASK_STATUS_PROCESSING, _TX_TASK_STATUS_PENDING, + _TX_TASK_STATUS_FAILED): + try: + with _use_default_datastore_adapter(): + query = datastore.Query(_PENDING_TASK_KIND, {'status =': status}) + # Fetch inside the adapter context; Run() returns a lazy iterator. + entities.extend(query.Run(limit=_SWEEPER_BATCH_SIZE)) + except Exception as e: + logging.error("Failed to query %s in sweeper: %s", _PENDING_TASK_KIND, e) + return + + now = datetime.datetime.utcnow() + keys_to_dispatch = [] + expired_failed_keys = [] + for entity in entities: + if not entity: + continue + status = entity.get('status', _TX_TASK_STATUS_PENDING) + if status == _TX_TASK_STATUS_DONE: + continue + if status == _TX_TASK_STATUS_FAILED: + failed_at = entity.get('failed_at') or entity.get('created') + if (isinstance(failed_at, datetime.datetime) + and now - failed_at > _FAILED_TASK_RETENTION): + expired_failed_keys.append(entity.key()) + continue # exceeded max sweeper retries + if not _is_dispatchable(entity, now): + continue + + created = entity.get('created') + if status == _TX_TASK_STATUS_PENDING and created and isinstance(created, datetime.datetime): + if (now - created).total_seconds() < _SWEEPER_FAST_PATH_GRACE_SECONDS: + continue # give fast-path grace period to dispatch post-commit + + keys_to_dispatch.append(entity.key()) + + if expired_failed_keys: + try: + with _use_default_datastore_adapter(): + datastore.Delete(expired_failed_keys) + logging.info("Cloud Tasks sweeper deleted %d expired failed tasks.", + len(expired_failed_keys)) + except Exception as e: + logging.error("Failed to delete expired failed tasks: %s", e) + + if keys_to_dispatch: + logging.info("Cloud Tasks sweeper found %d tasks to process.", len(keys_to_dispatch)) + _dispatch_pending_keys_now(keys_to_dispatch, handled_by_sweeper=True) + + +def sweep_wsgi_app(environ, start_response): + """WSGI app handler for /_ah/cloudtask/sweep.""" + is_cron = str(environ.get('HTTP_X_APPENGINE_CRON', '')).lower() == 'true' or str(environ.get('X-AppEngine-Cron', '')).lower() == 'true' + if not is_cron and not str(environ.get('SERVER_SOFTWARE', '')).lower().startswith('dev'): + status = '403 Forbidden' + response_headers = [('Content-Type', 'text/plain')] + start_response(status, response_headers) + return [b'Access denied: endpoint only accessible via App Engine Cron.\n'] + + try: + sweep() + status = '200 OK' + response_headers = [('Content-Type', 'text/plain')] + start_response(status, response_headers) + return [b'Sweeper completed successfully.\n'] + except Exception as e: + logging.error("Cloud Tasks sweeper failed: %s", e) + status = '500 Internal Server Error' + response_headers = [('Content-Type', 'text/plain')] + start_response(status, response_headers) + return [f'Sweeper failed: {e}\n'.encode('utf-8')] + + +# ============================================================================== +# Private Helpers +# ============================================================================== + + +@contextlib.contextmanager +def _use_default_datastore_adapter(non_transactional=False): + popped_conn = None + if non_transactional and datastore.IsInTransaction(): + popped_conn = datastore._PopConnection() + + try: + conn = datastore._GetConnection() + orig_adapter = getattr(conn, '_BaseConnection__adapter', None) + if orig_adapter is not None: + conn._BaseConnection__adapter = datastore._adapter + try: + yield + finally: + conn._BaseConnection__adapter = orig_adapter + else: + yield + finally: + if popped_conn is not None: + datastore._PushConnection(popped_conn) + + +def _is_dispatchable(entity, now): + """Returns True if a pending task entity may be (re)dispatched now.""" + status = entity.get('status', _TX_TASK_STATUS_PENDING) + if status == _TX_TASK_STATUS_PENDING: + return True + if status == _TX_TASK_STATUS_PROCESSING: + # A dispatcher holds the lock until it expires. A missing expiry means + # the lock was just taken. + lock_expires = entity.get('lock_expires') + return isinstance(lock_expires, datetime.datetime) and now >= lock_expires + return False + + +def _acquire_dispatch_lock(key, handled_by_sweeper): + """Atomically marks a pending entity as PROCESSING. + + Returns: + The locked entity, or None if it no longer exists or another dispatcher + holds it. + """ + + def txn(): + entity = datastore.Get(key) + now = datetime.datetime.utcnow() + if not _is_dispatchable(entity, now): + return None + entity['status'] = _TX_TASK_STATUS_PROCESSING + entity['lock_expires'] = now + datetime.timedelta( + seconds=_SWEEPER_LOCK_TIMEOUT_SECONDS) + entity['handled_by_sweeper'] = handled_by_sweeper + datastore.Put(entity) + return entity + + try: + with _use_default_datastore_adapter(non_transactional=True): + return datastore.RunInTransaction(txn) + except datastore_errors.EntityNotFoundError: + return None # already dispatched and deleted by another dispatcher + + +def _record_dispatch_failure(key, error): + """Atomically records a failed dispatch attempt on the entity. + + Does nothing if the entity was deleted meanwhile, so a dispatcher that + failed never recreates an entity another dispatcher already completed. + """ + + def txn(): + try: + entity = datastore.Get(key) + except datastore_errors.EntityNotFoundError: + return + now = datetime.datetime.utcnow() + retry_count = entity.get('retry_count', 0) + 1 + entity['retry_count'] = retry_count + entity['last_error'] = datastore_types.Text(str(error)[:500]) + entity['lock_expires'] = None + if retry_count >= _SWEEPER_MAX_RETRIES: + entity['status'] = _TX_TASK_STATUS_FAILED + entity['failed_at'] = now + else: + entity['status'] = _TX_TASK_STATUS_PENDING + datastore.Put(entity) + + with _use_default_datastore_adapter(non_transactional=True): + datastore.RunInTransaction(txn) + + +def _finalize_dispatched_entity(entity, task_name, dispatch_error): + """Deletes a dispatched staged entity or records its dispatch failure.""" + if dispatch_error is None: + try: + with _use_default_datastore_adapter(non_transactional=True): + datastore.Delete(entity.key()) + logging.info("Successfully dispatched transactional task %s", task_name) + except Exception as del_err: + logging.error( + "Failed to delete dispatched transactional task %s: %s", + task_name, + del_err, + ) + return + + if isinstance( + dispatch_error, + (google_exceptions.AlreadyExists, google_exceptions.Conflict), + ): + try: + with _use_default_datastore_adapter(non_transactional=True): + datastore.Delete(entity.key()) + logging.info( + "Transactional task %s already exists in Cloud Tasks; cleaned up entity", + task_name, + ) + except Exception as del_err: + logging.error( + "Failed to delete duplicate transactional task %s: %s", + task_name, + del_err, + ) + return + + logging.error( + "Failed to dispatch transactional task %s: %s", task_name, dispatch_error + ) + try: + _record_dispatch_failure(entity.key(), dispatch_error) + except Exception as put_err: + logging.error( + "Failed to record error state for task %s: %s", task_name, put_err + ) + + +def _dispatch_pending_keys_now(pending_keys, handled_by_sweeper=False): + locked_items = [] + for key in pending_keys: + try: + entity = _acquire_dispatch_lock(key, handled_by_sweeper) + except Exception as e: + logging.warning("Failed to acquire lock for task %s: %s", key, e) + continue + if entity is None: + continue + + queue_name = entity.get('queue_name') + task_name = entity.get('cloud_task_name') or entity.get('task_name') + try: + if entity.get('cloud_task_payload') is not None: + task = _decode_task_payload(entity['cloud_task_payload']) + else: + task = _decode_preview_payload(entity.get('payload')) + except Exception as e: + _finalize_dispatched_entity(entity, task_name, e) + continue + + locked_items.append((entity, queue_name, task_name, task)) + + if len(locked_items) == 1: + entity, queue_name, task_name, task = locked_items[0] + dispatch_err = None + try: + dispatch_task_payload(queue_name, task) + except Exception as e: + dispatch_err = e + _finalize_dispatched_entity(entity, task_name, dispatch_err) + elif len(locked_items) > 1: + dispatch_futures = [ + cloudtask._THREAD_POOL.submit( + contextvars.copy_context().run, + dispatch_task_payload, + queue_name, + task, + ) + for _, queue_name, _, task in locked_items + ] + futures.wait(dispatch_futures) + for (entity, _, task_name, _), fut in zip(locked_items, dispatch_futures): + _finalize_dispatched_entity(entity, task_name, fut.exception()) + + +def _pending_keys_for_current_transaction(): + """Returns the list of keys staged in the current transaction. + + On the first call in a transaction, the list is created and dispatch of the + keys it will hold is registered to run after the transaction commits. Each + transaction attempt has its own ndb context or datastore connection, so a + retried transaction starts with an empty list. + """ + if ndb and ndb.in_transaction(): + holder = ndb.get_context() + elif datastore.IsInTransaction(): + holder = datastore._GetConnection() + else: + raise taskqueue.BadTransactionStateError( + 'Transactional tasks must be added inside a transaction.' + ) + + pending_keys = getattr(holder, '_ae_pending_cloud_task_keys', None) + if pending_keys is not None: + return pending_keys + + pending_keys = [] + holder._ae_pending_cloud_task_keys = pending_keys + dispatch = lambda: _dispatch_pending_keys_now(pending_keys) + if ndb and ndb.in_transaction(): + holder.call_on_commit(dispatch) + else: + if not hasattr(holder, '_on_commit_callbacks'): + holder._on_commit_callbacks = [] + holder._on_commit_callbacks.append(dispatch) + return pending_keys diff --git a/src/google/appengine/api/taskqueue/taskqueue.py b/src/google/appengine/api/taskqueue/taskqueue.py index 20c5c53..a4f005b 100755 --- a/src/google/appengine/api/taskqueue/taskqueue.py +++ b/src/google/appengine/api/taskqueue/taskqueue.py @@ -44,6 +44,8 @@ from google.appengine.api import modules from google.appengine.api import namespace_manager from google.appengine.api import urlfetch +from google.appengine.api.taskqueue import cloudtask +from google.appengine.api.taskqueue import cloudtask_transactional from google.appengine.api.taskqueue import taskqueue_service_bytes_pb2 as taskqueue_service_pb2 from google.appengine.runtime import apiproxy_errors from google.appengine.runtime import context @@ -1557,13 +1559,21 @@ def fetch(cls, queue_or_queues, deadline=10): return [] - rpc = create_rpc(deadline) + if cloudtask.is_cloudtask_push_queue_enabled(): + rpc = cloudtask._DeadlineOnlyRPC(deadline) + else: + rpc = create_rpc(deadline) cls.fetch_async(queue_or_queues, rpc) return rpc.get_result() @classmethod def _FetchMultipleQueues(cls, queues, multiple, rpc=None): """Internal implementation of fetch stats where queues must be a list.""" + if cloudtask.is_cloudtask_push_queue_enabled(): + return cloudtask._make_rpc( + lambda: cloudtask.fetch_queue_stats_in_cloud_tasks( + queues, multiple, deadline=getattr(rpc, 'deadline', None)), + rpc) def ResultHook(rpc): """Processes the TaskQueueFetchQueueStatsResponse.""" @@ -1654,6 +1664,10 @@ def purge(self): Raises: Error-subclass on application errors. """ + if cloudtask.is_cloudtask_push_queue_enabled(): + cloudtask.purge_queue_in_cloud_tasks(self.__name) + return + request = taskqueue_service_pb2.TaskQueuePurgeQueueRequest() response = taskqueue_service_pb2.TaskQueuePurgeQueueResponse() @@ -1789,6 +1803,13 @@ def delete_tasks(self, task): def __DeleteTasks(self, tasks, multiple, rpc=None): """Internal implementation of delete_tasks_async(), tasks must be a list.""" + if (cloudtask.is_cloudtask_push_queue_enabled() + and not any(getattr(t, 'method', None) == 'PULL' for t in tasks)): + return cloudtask._make_rpc( + lambda: cloudtask.delete_tasks_in_cloud_tasks( + self.__name, tasks, multiple, + deadline=getattr(rpc, 'deadline', None)), + rpc) def ResultHook(rpc): """Processes the TaskQueueDeleteResponse.""" @@ -2083,6 +2104,11 @@ def add_async(self, task, transactional=False, rpc=None): (and therefore will never run). If `False`, the added tasks are available to run immediately; any enclosing transaction's success or failure is ignored. + When `APPENGINE_USE_CLOUDTASK_PUSH_QUEUE` is enabled, each + transactional task is staged as its own Datastore entity group + (consuming one of the 25 entity groups allowed per cross-group + transaction, up to 5 tasks per transaction), so the enclosing + transaction must be cross-group (`xg=True`). rpc: An optional UserRPC object. Returns: @@ -2133,6 +2159,21 @@ def add_async(self, task, transactional=False, rpc=None): raise InvalidTaskError( 'You cannot add both push and pull tasks in a single call.') + # Intercept for Cloud Tasks backend + if (cloudtask.is_cloudtask_push_queue_enabled() + and has_push_task + and len(tasks) >= 1): + if transactional: + cloudtask_transactional.add_transactional_tasks(self.__name, tasks, multiple) + return cloudtask._make_rpc( + lambda: tasks if multiple else tasks[0], rpc) + else: + return cloudtask._make_rpc( + lambda: cloudtask.create_tasks_in_cloud_tasks( + self.__name, tasks, multiple, + deadline=getattr(rpc, 'deadline', None)), + rpc) + if has_push_task: fill_function = self.__FillAddPushTasksRequest else: @@ -2175,6 +2216,11 @@ def add(self, task, transactional=False): (and therefore will never run). If `False`, the added tasks are available to run immediately; any enclosing transaction's success or failure is ignored. + When `APPENGINE_USE_CLOUDTASK_PUSH_QUEUE` is enabled, each + transactional task is staged as its own Datastore entity group + (consuming one of the 25 entity groups allowed per cross-group + transaction, up to 5 tasks per transaction), so the enclosing + transaction must be cross-group (`xg=True`). Returns: The task or list of tasks that was supplied to this method. Successfully @@ -2471,7 +2517,10 @@ def fetch_statistics(self, deadline=10): Error-subclass on application errors. """ _ValidateDeadline(deadline) - rpc = create_rpc(deadline) + if cloudtask.is_cloudtask_push_queue_enabled(): + rpc = cloudtask._DeadlineOnlyRPC(deadline) + else: + rpc = create_rpc(deadline) self.fetch_statistics_async(rpc) return rpc.get_result() diff --git a/src/google/appengine/datastore/datastore_rpc.py b/src/google/appengine/datastore/datastore_rpc.py index e518dd0..32e94a5 100755 --- a/src/google/appengine/datastore/datastore_rpc.py +++ b/src/google/appengine/datastore/datastore_rpc.py @@ -2626,6 +2626,13 @@ def async_commit(self, config): transaction = self.transaction if transaction is None: self._state = TransactionalConnection.CLOSED + callbacks = getattr(self, '_on_commit_callbacks', None) + if callbacks: + for cb in callbacks: + try: + cb() + except Exception as e: + logging.error('Error in on_commit callback: %s', e) return None if self._api_version == _CLOUD_DATASTORE_V1: @@ -2668,6 +2675,13 @@ def __commit_hook(self, rpc): else: raise _ToDatastoreError(err) else: + callbacks = getattr(self, '_on_commit_callbacks', None) + if callbacks: + for cb in callbacks: + try: + cb() + except Exception as e: + logging.error('Error in on_commit callback: %s', e) return True diff --git a/src/google/appengine/runtime/middlewares.py b/src/google/appengine/runtime/middlewares.py index e0efa2d..1e6c911 100755 --- a/src/google/appengine/runtime/middlewares.py +++ b/src/google/appengine/runtime/middlewares.py @@ -378,6 +378,26 @@ def BackgroundAndShutdownMiddleware(app, wsgi_env, start_response): return app(wsgi_env, start_response) +@middleware +def CloudTaskSweepMiddleware(app, wsgi_env, start_response): + """Intercept calls to the endpoint for Cloud Tasks Transactional Sweeper. + + Handle requests made to /_ah/cloudtask/sweep + + Args: + app: the WSGI app to wrap + wsgi_env: see PEP 3333 + start_response: see PEP 3333 + Returns: + The wrapped WSGI app response in UTF-8 format + """ + path = wsgi_env['PATH_INFO'] + if path == '/_ah/cloudtask/sweep': + from google.appengine.api.taskqueue import cloudtask_transactional + return cloudtask_transactional.sweep_wsgi_app(wsgi_env, start_response) + return app(wsgi_env, start_response) + + @middleware def AddDeferredMiddleware(app, wsgi_env, start_response): """Intercept calls to the default endpoint for Deferred. diff --git a/tests/google/appengine/api/taskqueue/cloudtask_test.py b/tests/google/appengine/api/taskqueue/cloudtask_test.py new file mode 100644 index 0000000..91263ac --- /dev/null +++ b/tests/google/appengine/api/taskqueue/cloudtask_test.py @@ -0,0 +1,1272 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import base64 +import datetime +import http +import json +import os +import sys +import time +import unittest +from unittest import mock + +import pytest + +if sys.version_info < (3, 10): + # google-cloud-tasks>=2.25.0 is only installed on Python 3.10+. + pytest.skip('Cloud Tasks push queues require Python 3.10+', + allow_module_level=True) + +from google.api_core import exceptions as google_exceptions +from google.api_core import operation +from google.appengine.api import datastore +from google.appengine.api import datastore_errors +from google.appengine.api import datastore_types +from google.appengine.api.taskqueue import cloudtask +from google.appengine.api.taskqueue import cloudtask_transactional +from google.appengine.api.taskqueue import taskqueue +from google.appengine.api.taskqueue import taskqueue_service_bytes_pb2 as taskqueue_service_pb2 +from google.appengine.ext import ndb +from google.appengine.ext import testbed +from google.cloud import tasks_v2 +from google.longrunning import operations_pb2 +from google.protobuf import duration_pb2 +from google.protobuf import empty_pb2 +from google.protobuf.timestamp_pb2 import Timestamp +from google.rpc import code_pb2 +from google.rpc import status_pb2 + + +def _make_batch_create_operation(task_names=None, failed_requests=None): + resp_pb = tasks_v2.BatchCreateTasksResponse.pb()( + tasks=[tasks_v2.Task.pb()(name=n) for n in (task_names or [])] + ) + meta_pb = tasks_v2.BatchCreateTasksMetadata.pb()( + failed_requests=failed_requests or {} + ) + op_pb = operations_pb2.Operation(name='operations/batch-create-1', done=True) + op_pb.response.Pack(resp_pb) + op_pb.metadata.Pack(meta_pb) + return operation.from_gapic( + op_pb, + mock.Mock(), + tasks_v2.BatchCreateTasksResponse, + metadata_type=tasks_v2.BatchCreateTasksMetadata, + ) + + +def _make_batch_delete_operation(failed_requests=None): + meta_pb = tasks_v2.BatchDeleteTasksMetadata.pb()( + failed_requests=failed_requests or {} + ) + op_pb = operations_pb2.Operation(name='operations/batch-delete-1', done=True) + op_pb.response.Pack(empty_pb2.Empty()) + op_pb.metadata.Pack(meta_pb) + return operation.from_gapic( + op_pb, + mock.Mock(), + empty_pb2.Empty, + metadata_type=tasks_v2.BatchDeleteTasksMetadata, + ) + + +class CloudtaskTest(unittest.TestCase): + + def setUp(self): + super(CloudtaskTest, self).setUp() + cloudtask._reset_clients() + self.addCleanup(cloudtask._reset_clients) + self.mock_client = mock.Mock() + self.mock_client.queue_path.side_effect = ( + lambda p, r, q: f'projects/{p}/locations/{r}/queues/{q}' + ) + self.mock_client.task_path.side_effect = ( + lambda p, r, q, t: f'projects/{p}/locations/{r}/queues/{q}/tasks/{t}' + ) + + @mock.patch.dict(os.environ, {'GAE_SERVICE': 'default-service'}, clear=False) + @mock.patch('google.appengine.api.app_identity.get_default_version_hostname', return_value='app.appspot.com') + def test_build_ct_task_payload_with_default_app_version_target(self, _): + task = taskqueue.Task(url='/test', target=taskqueue.DEFAULT_APP_VERSION) + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + self.assertIn('app_engine_http_request', payload) + self.assertEqual( + payload['app_engine_http_request'].get('app_engine_routing', {}).get('service'), + 'default-service' + ) + + @mock.patch.dict(os.environ, {}, clear=True) + @mock.patch('google.appengine.api.app_identity.get_default_version_hostname', return_value='app.appspot.com') + def test_build_ct_task_payload_with_string_target(self, _): + task = taskqueue.Task(url='/test', target='worker-dot') + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + self.assertEqual( + payload['app_engine_http_request'].get('app_engine_routing', {}).get('service'), + 'worker' + ) + + @mock.patch.dict(os.environ, {}, clear=True) + @mock.patch('google.appengine.api.app_identity.get_default_version_hostname', return_value='p.appspot.com') + def test_build_ct_task_payload_with_full_hostname_target(self, _): + task = taskqueue.Task(url='/test', target='v2-dot-worker-dot-p.uc.r.appspot.com') + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + routing = payload['app_engine_http_request'].get('app_engine_routing', {}) + self.assertEqual(routing.get('service'), 'worker') + self.assertEqual(routing.get('version'), 'v2') + + @mock.patch.dict(os.environ, {}, clear=True) + @mock.patch('google.appengine.api.app_identity.get_default_version_hostname', return_value='app.appspot.com') + def test_build_ct_task_payload_with_version_service_target(self, _): + task = taskqueue.Task(url='/test', target='v2.worker') + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + routing = payload['app_engine_http_request'].get('app_engine_routing', {}) + self.assertEqual(routing.get('service'), 'worker') + self.assertEqual(routing.get('version'), 'v2') + + @mock.patch.dict( + os.environ, {'GAE_SERVICE': 'default', 'GAE_VERSION': 'v1'}, clear=True) + @mock.patch('google.appengine.api.app_identity.get_default_version_hostname', return_value='app.appspot.com') + def test_build_ct_task_payload_routing_and_reserved_headers(self, _): + def routing_for(target, **kwargs): + task = taskqueue.Task(url='/test', target=target, **kwargs) + req = cloudtask._build_ct_task_payload( + 'q', task, self.mock_client, 'p', 'us-central1' + )['app_engine_http_request'] + return req, req.get('app_engine_routing', {}) + + # Another service without a version uses that service's default version. + req, routing = routing_for( + 'worker', headers={'X-Custom': '1', 'X-AppEngine-Foo': 'x'}) + self.assertEqual(routing, {'service': 'worker'}) + # Host and X-AppEngine-* cannot be set on Cloud Tasks tasks. + self.assertEqual(req['headers'], {'X-Custom': '1'}) + + # The current service keeps the version that enqueued the task. + _, routing = routing_for('default-dot') + self.assertEqual(routing, {'service': 'default', 'version': 'v1'}) + + # Version-specific "-dot-" hostnames are split into version and service. + _, routing = routing_for('v2-dot-worker-dot') + self.assertEqual(routing, {'service': 'worker', 'version': 'v2'}) + + # Domain-scoped project IDs in full .appspot.com targets are stripped. + task_domain = taskqueue.Task( + url='/test', + target='v2-dot-worker-dot-my-app.example.com.uc.r.appspot.com', + ) + routing_domain = cloudtask._build_ct_task_payload( + 'q', task_domain, self.mock_client, 'example.com:my-app', 'us-central1' + )['app_engine_http_request'].get('app_engine_routing', {}) + self.assertEqual(routing_domain, {'service': 'worker', 'version': 'v2'}) + + # A full .appspot.com target belonging to another project falls back to + # current GAE_SERVICE / GAE_VERSION. + _, routing_other = routing_for('v2-dot-worker-dot-other-proj.appspot.com') + self.assertEqual(routing_other, {'service': 'default', 'version': 'v1'}) + + def test_build_ct_task_payload_with_countdown(self): + now = time.time() + countdown_seconds = 60 + task = taskqueue.Task(url='/test', countdown=countdown_seconds) + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + self.assertIn('schedule_time', payload) + st = payload['schedule_time'] + self.assertAlmostEqual(st.seconds, int(now + countdown_seconds), delta=2) + + def test_build_ct_task_payload_with_eta(self): + eta = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(seconds=120) + task = taskqueue.Task(url='/test', eta=eta) + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1' + ) + self.assertIn('schedule_time', payload) + st = payload['schedule_time'] + self.assertAlmostEqual(st.seconds, int(eta.timestamp()), delta=2) + + def test_build_ct_task_payload_with_retry_options(self): + retry_opts = taskqueue.TaskRetryOptions( + task_retry_limit=4, + task_age_limit=600, + min_backoff_seconds=1.5, + max_backoff_seconds=30, + max_doublings=3, + ) + task = taskqueue.Task(url='/test', retry_options=retry_opts) + payload = cloudtask._build_ct_task_payload( + queue_name='default', + task=task, + client=self.mock_client, + project='p', + region='us-central1', + ) + self.assertIn('retry_config', payload) + rc = payload['retry_config'] + self.assertEqual(rc['max_attempts'], 5) + self.assertEqual(rc['max_retry_duration'], duration_pb2.Duration(seconds=600, nanos=0)) + self.assertEqual(rc['min_backoff'], duration_pb2.Duration(seconds=1, nanos=500000000)) + self.assertEqual(rc['max_backoff'], duration_pb2.Duration(seconds=30, nanos=0)) + self.assertEqual(rc['max_doublings'], 3) + + def test_get_project_id(self): + with mock.patch.dict(os.environ, {cloudtask.ENV_GOOGLE_CLOUD_PROJECT: 's~my-app'}): + self.assertEqual(cloudtask._get_project_id(), 'my-app') + with mock.patch.dict(os.environ, {cloudtask.ENV_GOOGLE_CLOUD_PROJECT: 'e~my-app'}): + self.assertEqual(cloudtask._get_project_id(), 'my-app') + with mock.patch.dict(os.environ, {cloudtask.ENV_GOOGLE_CLOUD_PROJECT: 'my-app'}): + self.assertEqual(cloudtask._get_project_id(), 'my-app') + + def test_get_region_normalization_and_metadata_caching(self): + with mock.patch.dict(os.environ, {'LOCATION_ID': 'us-central'}, clear=True): + self.assertEqual(cloudtask._get_region(), 'us-central1') + with mock.patch.dict(os.environ, {'GAE_REGION': 'europe-west'}, clear=True): + self.assertEqual(cloudtask._get_region(), 'europe-west1') + with mock.patch.dict(os.environ, {'REGION_ID': 'asia-northeast1'}, clear=True): + self.assertEqual(cloudtask._get_region(), 'asia-northeast1') + + cloudtask._reset_clients() + mock_resp = mock.MagicMock() + mock_resp.__enter__.return_value.read.return_value = ( + b'projects/123/regions/us-central' + ) + with mock.patch.dict(os.environ, {}, clear=True): + with mock.patch('urllib.request.urlopen', return_value=mock_resp) as m_open: + self.assertEqual(cloudtask._get_region(), 'us-central1') + self.assertEqual(cloudtask._get_region(), 'us-central1') + self.assertEqual(m_open.call_count, 1) + + def test_is_cloudtask_push_queue_enabled(self): + with mock.patch.dict(os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true'}): + self.assertTrue(cloudtask.is_cloudtask_push_queue_enabled()) + with mock.patch.dict(os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'True'}): + self.assertTrue(cloudtask.is_cloudtask_push_queue_enabled()) + with mock.patch.dict(os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'false'}): + self.assertFalse(cloudtask.is_cloudtask_push_queue_enabled()) + with mock.patch.dict(os.environ, {}, clear=True): + self.assertFalse(cloudtask.is_cloudtask_push_queue_enabled()) + + def test_enabled_without_cloud_tasks_library_raises(self): + with mock.patch.object( + cloudtask, '_CLOUD_TASKS_IMPORT_ERROR', ImportError('missing')): + with mock.patch.dict( + os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true'}): + with self.assertRaisesRegex(ImportError, 'Python 3.10'): + cloudtask.is_cloudtask_push_queue_enabled() + with mock.patch.dict(os.environ, {}, clear=True): + self.assertFalse(cloudtask.is_cloudtask_push_queue_enabled()) + + def test_map_rest_code_to_tq_code(self): + self.assertEqual( + cloudtask._map_rest_code_to_tq_code(code_pb2.NOT_FOUND), + taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK + ) + self.assertEqual( + cloudtask._map_rest_code_to_tq_code(http.HTTPStatus.NOT_FOUND), + taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK + ) + self.assertEqual( + cloudtask._map_rest_code_to_tq_code(code_pb2.ALREADY_EXISTS), + taskqueue_service_pb2.TaskQueueServiceError.TASK_ALREADY_EXISTS + ) + self.assertEqual( + cloudtask._map_rest_code_to_tq_code(code_pb2.INVALID_ARGUMENT), + taskqueue_service_pb2.TaskQueueServiceError.INVALID_REQUEST + ) + self.assertEqual( + cloudtask._map_rest_code_to_tq_code(code_pb2.PERMISSION_DENIED), + taskqueue_service_pb2.TaskQueueServiceError.PERMISSION_DENIED + ) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + def test_create_single_task_success_and_errors(self, mock_client_cls): + mock_client_cls.return_value = self.mock_client + self.mock_client.create_task.return_value = tasks_v2.Task( + name='projects/my-proj/locations/us-central1/queues/default/tasks/gen-task-1' + ) + + t = taskqueue.Task(url='/worker', payload='hello') + res = cloudtask.create_tasks_in_cloud_tasks('default', [t], multiple=False) + self.assertIs(res, t) + self.assertEqual(t.name, 'gen-task-1') + self.assertEqual(t.queue_name, 'default') + self.assertTrue(t.was_enqueued) + + # AlreadyExists -> TaskAlreadyExistsError + self.mock_client.create_task.side_effect = google_exceptions.AlreadyExists('dup') + with self.assertRaises(taskqueue.TaskAlreadyExistsError): + cloudtask.create_tasks_in_cloud_tasks('default', [taskqueue.Task(url='/worker')], multiple=False) + + # NotFound -> UnknownQueueError + self.mock_client.create_task.side_effect = google_exceptions.NotFound('missing queue') + with self.assertRaises(taskqueue.UnknownQueueError): + cloudtask.create_tasks_in_cloud_tasks('missing', [taskqueue.Task(url='/worker')], multiple=False) + + # BadRequest with Queue does not exist -> UnknownQueueError + self.mock_client.create_task.side_effect = google_exceptions.BadRequest('Queue does not exist') + with self.assertRaises(taskqueue.UnknownQueueError): + cloudtask.create_tasks_in_cloud_tasks('missing', [taskqueue.Task(url='/worker')], multiple=False) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + def test_create_batch_tasks_with_real_lro_operation(self, mock_client_cls): + mock_client_cls.return_value = self.mock_client + self.mock_client.batch_create_tasks.return_value = _make_batch_create_operation( + task_names=[ + 'projects/my-proj/locations/us-central1/queues/default/tasks/b1', + 'projects/my-proj/locations/us-central1/queues/default/tasks/b2', + ] + ) + + t1 = taskqueue.Task(url='/worker', payload='1') + t2 = taskqueue.Task(url='/worker', payload='2') + res = cloudtask.create_tasks_in_cloud_tasks('default', [t1, t2], multiple=True) + + self.assertEqual(len(res), 2) + self.assertEqual(t1.name, 'b1') + self.assertEqual(t2.name, 'b2') + self.assertTrue(t1.was_enqueued) + self.assertTrue(t2.was_enqueued) + self.mock_client.batch_create_tasks.assert_called_once() + + # Re-adding an already enqueued task raises BadTaskStateError. + with self.assertRaises(taskqueue.BadTaskStateError): + cloudtask.create_tasks_in_cloud_tasks('default', [t1], multiple=False) + + # Multi-chunk batch execution (>100 tasks) preserves order and names. + def fake_batch_create(request, **_): + reqs = request['requests'] + return _make_batch_create_operation( + task_names=[ + f"projects/my-proj/locations/us-central1/queues/default/tasks/{r['task']['app_engine_http_request']['relative_uri'].lstrip('/')}" + for r in reqs + ] + ) + + self.mock_client.batch_create_tasks.reset_mock() + self.mock_client.batch_create_tasks.side_effect = fake_batch_create + many_tasks = [taskqueue.Task(url=f'/t-{i}') for i in range(250)] + many_res = cloudtask.create_tasks_in_cloud_tasks( + 'default', many_tasks, multiple=True + ) + self.assertEqual(self.mock_client.batch_create_tasks.call_count, 3) + self.assertEqual([t.name for t in many_res], [f't-{i}' for i in range(250)]) + + # Missing task in BatchCreateTasks response raises InternalError instead of + # silently leaving task.name as None. + self.mock_client.batch_create_tasks.side_effect = None + self.mock_client.batch_create_tasks.return_value = ( + _make_batch_create_operation(task_names=[]) + ) + with self.assertRaises(taskqueue.InternalError): + cloudtask.create_tasks_in_cloud_tasks( + 'default', + [taskqueue.Task(url='/1'), taskqueue.Task(url='/2')], + multiple=True, + ) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + def test_create_batch_tasks_partial_failure_and_duplicate_names(self, mock_client_cls): + mock_client_cls.return_value = self.mock_client + + # Duplicate names pre-condition check + dup1 = taskqueue.Task(name='same-name', url='/worker') + dup2 = taskqueue.Task(name='same-name', url='/worker') + with self.assertRaises(taskqueue.DuplicateTaskNameError): + cloudtask.create_tasks_in_cloud_tasks('default', [dup1, dup2], multiple=True) + + # Partial failure via LRO metadata.failed_requests (index 1 failed with ALREADY_EXISTS) + self.mock_client.batch_create_tasks.return_value = _make_batch_create_operation( + task_names=[ + 'projects/my-proj/locations/us-central1/queues/default/tasks/ok-task', + ], + failed_requests={ + 1: status_pb2.Status(code=code_pb2.ALREADY_EXISTS, message='Task already exists') + }, + ) + t_ok = taskqueue.Task(name='ok-task', url='/worker') + t_dup = taskqueue.Task(name='dup-task', url='/worker') + with self.assertRaises(taskqueue.TaskAlreadyExistsError): + cloudtask.create_tasks_in_cloud_tasks('default', [t_ok, t_dup], multiple=True) + + self.assertTrue(t_ok.was_enqueued) + self.assertEqual(t_ok.name, 'ok-task') + self.assertFalse(t_dup.was_enqueued) + + # INVALID_ARGUMENT is a generic error that keeps the server's message. + self.mock_client.batch_create_tasks.return_value = _make_batch_create_operation( + failed_requests={ + 0: status_pb2.Status( + code=code_pb2.INVALID_ARGUMENT, message='schedule time too far') + }, + ) + with self.assertRaisesRegex(taskqueue.Error, 'schedule time too far') as cm: + cloudtask.create_tasks_in_cloud_tasks( + 'default', [taskqueue.Task(url='/a'), taskqueue.Task(url='/b')], + multiple=True) + self.assertNotIsInstance(cm.exception, taskqueue.InvalidTaskNameError) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + def test_delete_tasks_in_cloud_tasks(self, mock_client_cls): + mock_client_cls.return_value = self.mock_client + self.mock_client.batch_delete_tasks.return_value = _make_batch_delete_operation() + + t1 = taskqueue.Task(name='del-1') + t2 = taskqueue.Task(name='del-2') + res = cloudtask.delete_tasks_in_cloud_tasks('default', [t1, t2], multiple=True) + self.assertEqual(len(res), 2) + self.assertTrue(t1.was_deleted) + self.assertTrue(t2.was_deleted) + + # Already deleted task raises BadTaskStateError + with self.assertRaises(taskqueue.BadTaskStateError): + cloudtask.delete_tasks_in_cloud_tasks('default', [t1], multiple=False) + + # Unnamed task raises BadTaskStateError + with self.assertRaises(taskqueue.BadTaskStateError): + cloudtask.delete_tasks_in_cloud_tasks('default', [taskqueue.Task()], multiple=False) + + # Duplicate names raises DuplicateTaskNameError + with self.assertRaises(taskqueue.DuplicateTaskNameError): + cloudtask.delete_tasks_in_cloud_tasks( + 'default', [taskqueue.Task(name='d'), taskqueue.Task(name='d')], multiple=True + ) + + # NOT_FOUND in failed_requests marks was_deleted=False without raising + self.mock_client.batch_delete_tasks.return_value = _make_batch_delete_operation( + failed_requests={ + 0: status_pb2.Status(code=code_pb2.NOT_FOUND, message='Not found') + } + ) + t_missing = taskqueue.Task(name='missing-task') + t_present = taskqueue.Task(name='present-task') + cloudtask.delete_tasks_in_cloud_tasks('default', [t_missing, t_present], multiple=True) + self.assertFalse(t_missing.was_deleted) + self.assertTrue(t_present.was_deleted) + + # Non-ignored error in failed_requests raises translated exception + self.mock_client.batch_delete_tasks.return_value = _make_batch_delete_operation( + failed_requests={ + 0: status_pb2.Status(code=code_pb2.PERMISSION_DENIED, message='Denied') + } + ) + with self.assertRaises(taskqueue.PermissionDeniedError): + cloudtask.delete_tasks_in_cloud_tasks( + 'default', [taskqueue.Task(name='denied-task')], multiple=False + ) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + def test_purge_queue_in_cloud_tasks(self, mock_client_cls): + mock_client_cls.return_value = self.mock_client + cloudtask.purge_queue_in_cloud_tasks('my-queue') + self.mock_client.purge_queue.assert_called_once_with( + request={'name': 'projects/my-proj/locations/us-central1/queues/my-queue'} + ) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_fetch_queue_stats_in_cloud_tasks(self, mock_v2beta3_client_cls): + mock_v2beta3_client = mock.Mock() + mock_v2beta3_client.queue_path.side_effect = ( + lambda p, r, q: f'projects/{p}/locations/{r}/queues/{q}' + ) + mock_v2beta3_client_cls.return_value = mock_v2beta3_client + + eta_dt = datetime.datetime(2026, 9, 30, 12, 0, 0, tzinfo=datetime.timezone.utc) + mock_stats = mock.Mock( + tasks_count=42, + oldest_estimated_arrival_time=eta_dt, + executed_last_minute_count=7, + concurrent_dispatches_count=3, + effective_execution_rate=5.0, + ) + mock_v2beta3_client.get_queue.return_value = mock.Mock(stats=mock_stats) + + q = taskqueue.Queue('default') + stats = cloudtask.fetch_queue_stats_in_cloud_tasks([q], multiple=False) + self.assertIsInstance(stats, taskqueue.QueueStatistics) + self.assertEqual(stats.tasks, 42) + self.assertEqual(stats.oldest_eta_usec, int(eta_dt.timestamp() * 1e6)) + self.assertEqual(stats.executed_last_minute, 7) + self.assertEqual(stats.in_flight, 3) + self.assertEqual(stats.enforced_rate, 5.0) + + mock_v2beta3_client.get_queue.side_effect = google_exceptions.NotFound('no queue') + with self.assertRaises(taskqueue.UnknownQueueError): + cloudtask.fetch_queue_stats_in_cloud_tasks([q], multiple=False) + + @mock.patch.dict( + os.environ, + { + 'APPLICATION_ID': 's~my-proj', + 'GOOGLE_CLOUD_PROJECT': 'my-proj', + 'LOCATION_ID': 'us-central1', + }, + ) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.appengine.api.datastore.IsInTransaction', return_value=True) + @mock.patch('google.appengine.api.datastore._GetConnection') + @mock.patch('google.appengine.api.datastore.Put') + def test_transactional_tasks_with_retry_options_and_schedule_time( + self, mock_put, mock_get_conn, mock_in_tx, mock_client_cls + ): + mock_client_cls.return_value = self.mock_client + mock_conn = mock.Mock(_BaseConnection__adapter=None, + _ae_pending_cloud_task_keys=None) + mock_conn._on_commit_callbacks = [] + mock_get_conn.return_value = mock_conn + + saved_entities = [] + + def fake_put(entity): + entity._Entity__key = datastore.Key.from_path( + cloudtask_transactional._PENDING_TASK_KIND, 1, _app='my-proj' + ) + saved_entities.append(entity) + + mock_put.side_effect = fake_put + + retry_opts = taskqueue.TaskRetryOptions( + task_retry_limit=3, + task_age_limit=120, + min_backoff_seconds=2.5, + max_backoff_seconds=10, + max_doublings=2, + ) + task = taskqueue.Task( + url='/worker/tx', + payload=b'tx-body', + countdown=30, + retry_options=retry_opts, + ) + + res = cloudtask_transactional.add_transactional_tasks( + 'default', [task], multiple=False + ) + self.assertIs(res, task) + self.assertTrue(task.was_enqueued) + self.assertTrue(task.name.startswith('tx-')) + self.assertEqual(len(saved_entities), 1) + self.assertEqual(len(mock_conn._on_commit_callbacks), 1) + + # The entity uses the kind, property names and payload form shared with + # the Go and Java SDKs. + entity = saved_entities[0] + self.assertEqual(entity['queue_name'], 'default') + self.assertEqual(entity['cloud_task_name'], task.name) + self.assertEqual(entity['sdk_lang'], 'PYTHON') + self.assertIsInstance(entity['cloud_task_payload'], datastore_types.Text) + staged = json.loads(entity['cloud_task_payload'])['task'] + self.assertEqual( + staged['name'], + 'projects/my-proj/locations/us-central1/queues/default/tasks/' + + task.name) + ae_req = staged['appEngineHttpRequest'] + self.assertEqual(ae_req['httpMethod'], 'POST') + self.assertEqual(ae_req['relativeUri'], '/worker/tx') + self.assertEqual(base64.b64decode(ae_req['body']), b'tx-body') + self.assertIn('scheduleTime', staged) + self.assertEqual(staged['retryConfig'], { + 'maxAttempts': 4, + 'maxRetryDuration': '120s', + 'minBackoff': '2.500s', + 'maxBackoff': '10s', + 'maxDoublings': 2, + }) + + # The staged payload is dispatched as the same task. + dispatched = cloudtask_transactional._decode_task_payload( + entity['cloud_task_payload']) + cloudtask_transactional.dispatch_task_payload('default', dispatched) + self.mock_client.create_task.assert_called_once() + call_req = self.mock_client.create_task.call_args[1]['request'] + self.assertEqual( + call_req['parent'], 'projects/my-proj/locations/us-central1/queues/default' + ) + sent = call_req['task'] + self.assertEqual(sent.app_engine_http_request.body, b'tx-body') + self.assertEqual(sent.app_engine_http_request.relative_uri, '/worker/tx') + self.assertEqual(sent.app_engine_http_request.http_method, + tasks_v2.HttpMethod.POST) + self.assertEqual(sent.retry_config.max_attempts, 4) + self.assertEqual(sent.retry_config.max_retry_duration.total_seconds(), 120.0) + self.assertEqual(sent.retry_config.min_backoff.total_seconds(), 2.5) + self.assertEqual(sent.retry_config.max_backoff.total_seconds(), 10.0) + self.assertEqual(sent.retry_config.max_doublings, 2) + expected_eta = task.eta.timestamp() + self.assertAlmostEqual(sent.schedule_time.timestamp(), expected_eta, 3) + + def test_to_duration_rounds_to_nanoseconds(self): + self.assertEqual( + cloudtask._to_duration(2.3), + duration_pb2.Duration(seconds=2, nanos=300000000)) + self.assertEqual( + cloudtask._to_duration(120), duration_pb2.Duration(seconds=120)) + self.assertIsNone(cloudtask._to_duration(None)) + + @mock.patch.dict(os.environ, {'APPLICATION_ID': 's~my-proj'}) + def test_transactional_tasks_preconditions(self): + # Named task in transaction raises InvalidTaskNameError + with self.assertRaises(taskqueue.InvalidTaskNameError): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(name='named-tx', url='/worker')], multiple=False + ) + + # Outside transaction raises BadTransactionStateError + with mock.patch('google.appengine.api.datastore.IsInTransaction', return_value=False): + with mock.patch('google.appengine.api.datastore.Put'): + with mock.patch('google.appengine.api.datastore._GetConnection'): + with mock.patch('google.cloud.tasks_v2.CloudTasksClient', return_value=self.mock_client): + with self.assertRaises(taskqueue.BadTransactionStateError): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/worker')], multiple=False + ) + + @mock.patch.dict( + os.environ, + { + 'APPLICATION_ID': 's~my-proj', + 'GOOGLE_CLOUD_PROJECT': 'my-proj', + 'LOCATION_ID': 'us-central1', + 'APPENGINE_USE_CLOUDTASK_PUSH_QUEUE': 'true', + }, + ) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_queue_integration_with_cloud_tasks_enabled( + self, mock_v2beta3_cls, mock_v2_cls + ): + mock_v2_cls.return_value = self.mock_client + mock_v2beta3_client = mock.Mock() + mock_v2beta3_client.queue_path.side_effect = ( + lambda p, r, q: f'projects/{p}/locations/{r}/queues/{q}' + ) + mock_v2beta3_cls.return_value = mock_v2beta3_client + + q = taskqueue.Queue('default') + + # Single add via Queue.add + self.mock_client.create_task.return_value = tasks_v2.Task( + name='projects/my-proj/locations/us-central1/queues/default/tasks/q-t1' + ) + t1 = q.add(taskqueue.Task(url='/worker', payload=b'hello')) + self.assertEqual(t1.name, 'q-t1') + self.assertTrue(t1.was_enqueued) + + # Batch add via Queue.add + self.mock_client.batch_create_tasks.return_value = ( + _make_batch_create_operation( + task_names=[ + 'projects/my-proj/locations/us-central1/queues/default/tasks/q-b1', + 'projects/my-proj/locations/us-central1/queues/default/tasks/q-b2', + ] + ) + ) + batch_res = q.add([ + taskqueue.Task(url='/worker', payload=b'1'), + taskqueue.Task(url='/worker', payload=b'2'), + ]) + self.assertEqual([t.name for t in batch_res], ['q-b1', 'q-b2']) + + # Delete tasks via Queue.delete_tasks + self.mock_client.batch_delete_tasks.return_value = ( + _make_batch_delete_operation() + ) + del_res = q.delete_tasks(batch_res) + self.assertEqual([t.was_deleted for t in del_res], [True, True]) + + # Purge via Queue.purge + q.purge() + self.mock_client.purge_queue.assert_called_once() + + # Fetch statistics via Queue.fetch_statistics + mock_stats = mock.Mock( + tasks_count=5, + oldest_estimated_arrival_time=None, + executed_last_minute_count=1, + concurrent_dispatches_count=0, + effective_execution_rate=10.0, + ) + mock_v2beta3_client.get_queue.return_value = mock.Mock(stats=mock_stats) + stats = q.fetch_statistics() + self.assertEqual(stats.tasks, 5) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_api_errors_are_raised_as_taskqueue_errors( + self, mock_v2beta3_cls, mock_v2_cls + ): + mock_v2_cls.return_value = self.mock_client + mock_v2beta3_cls.return_value = self.mock_client + queue_mode_detail = ( + 'The task cannot be created because the queue is not a push queue. ' + '[detail: "[ORIGINAL ERROR] ExecutorError::1006: ' + 'apphosting::ExecutorServiceError::INVALID_QUEUE_MODE"]') + cases = [ + (google_exceptions.PermissionDenied('denied'), + taskqueue.PermissionDeniedError), + (google_exceptions.Unauthenticated('no creds'), + taskqueue.PermissionDeniedError), + (google_exceptions.ServiceUnavailable('unavailable'), + taskqueue.TransientError), + (google_exceptions.DeadlineExceeded('deadline'), + taskqueue.TransientError), + (google_exceptions.Aborted('aborted'), taskqueue.TransientError), + (google_exceptions.RetryError( + 'Timeout of 20.0s exceeded', + google_exceptions.DeadlineExceeded('deadline')), + taskqueue.TransientError), + (google_exceptions.InternalServerError('internal'), + taskqueue.InternalError), + (google_exceptions.NotFound('Queue does not exist.'), + taskqueue.UnknownQueueError), + (google_exceptions.InvalidArgument(queue_mode_detail), + taskqueue.InvalidQueueModeError), + (google_exceptions.FailedPrecondition( + 'Queue does not exist. apphosting::ExecutorServiceError::' + 'UNKNOWN_QUEUE'), taskqueue.UnknownQueueError), + ] + calls = { + 'create_task': lambda: cloudtask.create_tasks_in_cloud_tasks( + 'q', [taskqueue.Task(url='/w')], multiple=False), + 'batch_create_tasks': lambda: cloudtask.create_tasks_in_cloud_tasks( + 'q', [taskqueue.Task(url='/w'), taskqueue.Task(url='/w')], + multiple=True), + 'batch_delete_tasks': lambda: cloudtask.delete_tasks_in_cloud_tasks( + 'q', [taskqueue.Task(name='t1')], multiple=False), + 'purge_queue': lambda: cloudtask.purge_queue_in_cloud_tasks('q'), + 'get_queue': lambda: cloudtask.fetch_queue_stats_in_cloud_tasks( + ['q'], multiple=False), + } + for api_error, expected in cases: + for method, call in calls.items(): + with self.subTest(error=type(api_error).__name__, method=method): + getattr(self.mock_client, method).side_effect = api_error + with self.assertRaises(expected) as cm: + call() + self.assertIsInstance(cm.exception, taskqueue.Error) + getattr(self.mock_client, method).side_effect = None + + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_caller_supplied_rpc_reports_cloud_tasks_result( + self, mock_v2beta3_cls, mock_v2_cls + ): + mock_v2_cls.return_value = self.mock_client + mock_v2beta3_cls.return_value = self.mock_client + tb = testbed.Testbed() + tb.activate() + self.addCleanup(tb.deactivate) + tb.init_taskqueue_stub() + # Patched after testbed.activate so that deactivate restores the + # environment without the flag. + env = mock.patch.dict(os.environ, { + 'GOOGLE_CLOUD_PROJECT': 'my-proj', + 'LOCATION_ID': 'us-central1', + cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true', + }) + env.start() + self.addCleanup(env.stop) + self.mock_client.create_task.return_value = tasks_v2.Task( + name='projects/my-proj/locations/us-central1/queues/default/tasks/t1') + self.mock_client.get_queue.return_value = mock.Mock(stats=None) + self.mock_client.batch_delete_tasks.return_value = ( + _make_batch_delete_operation()) + q = taskqueue.Queue('default') + + callbacks = [] + rpc = taskqueue.create_rpc(callback=lambda: callbacks.append('add')) + self.assertIs(q.add_async(taskqueue.Task(url='/w'), rpc=rpc), rpc) + self.assertEqual(rpc.get_result().name, 't1') + rpc.wait() + self.assertEqual(callbacks, ['add']) # Called once, like UserRPC. + + rpc = taskqueue.create_rpc() + self.assertIs(q.fetch_statistics_async(rpc), rpc) + self.assertEqual(rpc.get_result().tasks, 0) + + rpc = taskqueue.create_rpc() + self.assertIs(q.delete_tasks_async(taskqueue.Task(name='t1'), rpc=rpc), + rpc) + self.assertTrue(rpc.get_result().was_deleted) + + self.mock_client.create_task.side_effect = ( + google_exceptions.PermissionDenied('denied')) + rpc = taskqueue.create_rpc() + q.add_async(taskqueue.Task(url='/w'), rpc=rpc) + rpc.wait() # Like UserRPC.wait, does not raise the call's error. + with self.assertRaises(taskqueue.PermissionDeniedError): + rpc.check_success() + + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_caller_deadline_is_passed_to_cloud_tasks( + self, mock_v2beta3_cls, mock_v2_cls + ): + mock_v2_cls.return_value = self.mock_client + mock_v2beta3_cls.return_value = self.mock_client + tb = testbed.Testbed() + tb.activate() + self.addCleanup(tb.deactivate) + tb.init_taskqueue_stub() + env = mock.patch.dict(os.environ, { + 'GOOGLE_CLOUD_PROJECT': 'my-proj', + 'LOCATION_ID': 'us-central1', + cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true', + }) + env.start() + self.addCleanup(env.stop) + self.mock_client.create_task.return_value = tasks_v2.Task( + name='projects/my-proj/locations/us-central1/queues/default/tasks/t1') + self.mock_client.get_queue.return_value = mock.Mock(stats=None) + q = taskqueue.Queue('default') + + def assert_timeout(mock_method, deadline): + timeout = mock_method.call_args[1]['timeout'] + self.assertGreater(timeout, 0) + self.assertLessEqual(timeout, deadline) + + q.add_async(taskqueue.Task(url='/w'), + rpc=taskqueue.create_rpc(deadline=2.5)).get_result() + assert_timeout(self.mock_client.create_task, 2.5) + + create_op = _make_batch_create_operation( + task_names=[ + 'projects/p/locations/l/queues/default/tasks/b1', + 'projects/p/locations/l/queues/default/tasks/b2', + ] + ) + self.mock_client.batch_create_tasks.return_value = create_op + q.add_async([taskqueue.Task(url='/w'), taskqueue.Task(url='/w')], + rpc=taskqueue.create_rpc(deadline=1.5)).get_result() + assert_timeout(self.mock_client.batch_create_tasks, 1.5) + + self.mock_client.batch_delete_tasks.return_value = ( + _make_batch_delete_operation()) + q.delete_tasks_async(taskqueue.Task(name='t1'), + rpc=taskqueue.create_rpc(deadline=2)).get_result() + assert_timeout(self.mock_client.batch_delete_tasks, 2) + + # The synchronous methods pass their deadline argument. GetQueue retries + # are limited to the deadline too. + q.fetch_statistics(deadline=3) + assert_timeout(self.mock_client.get_queue, 3) + self.assertLessEqual( + self.mock_client.get_queue.call_args[1]['retry'].timeout, 3) + taskqueue.QueueStatistics.fetch('default', deadline=4) + assert_timeout(self.mock_client.get_queue, 4) + + # A caller deadline above 20s is capped so Cloud Tasks does not reject the + # request for having a deadline more than 30s in the future. + q.fetch_statistics(deadline=60) + self.assertEqual( + self.mock_client.get_queue.call_args[1]['timeout'], + cloudtask._DEFAULT_RPC_TIMEOUT_SECONDS) + + # Without a deadline, the client's default timeout applies for non-batch + # calls, while batch calls pass _DEFAULT_RPC_TIMEOUT_SECONDS. + q.add(taskqueue.Task(url='/w')) + self.assertNotIn('timeout', self.mock_client.create_task.call_args[1]) + q.add([taskqueue.Task(url='/w'), taskqueue.Task(url='/w')]) + self.assertEqual( + self.mock_client.batch_create_tasks.call_args[1]['timeout'], + cloudtask._DEFAULT_RPC_TIMEOUT_SECONDS) + + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) + @mock.patch('google.cloud.tasks_v2.CloudTasksClient') + @mock.patch('google.cloud.tasks_v2beta3.CloudTasksClient') + def test_clients_are_created_once_and_reused( + self, mock_v2beta3_cls, mock_v2_cls + ): + mock_v2_cls.return_value = self.mock_client + mock_v2beta3_client = mock.Mock() + mock_v2beta3_client.queue_path.side_effect = ( + lambda p, r, q: f'projects/{p}/locations/{r}/queues/{q}' + ) + mock_v2beta3_client.get_queue.return_value = mock.Mock(stats=None) + mock_v2beta3_cls.return_value = mock_v2beta3_client + self.mock_client.create_task.return_value = tasks_v2.Task( + name='projects/my-proj/locations/us-central1/queues/default/tasks/t' + ) + + cloudtask.create_tasks_in_cloud_tasks( + 'default', [taskqueue.Task(url='/worker')], multiple=False + ) + cloudtask.create_tasks_in_cloud_tasks( + 'default', [taskqueue.Task(url='/worker')], multiple=False + ) + cloudtask.purge_queue_in_cloud_tasks('default') + cloudtask.fetch_queue_stats_in_cloud_tasks(['default'], multiple=False) + cloudtask.fetch_queue_stats_in_cloud_tasks(['default'], multiple=False) + + self.assertEqual(mock_v2_cls.call_count, 1) + self.assertEqual(mock_v2beta3_cls.call_count, 1) + + +class CloudtaskTransactionalDatastoreTest(unittest.TestCase): + """Transactional task staging and dispatch against the datastore stub.""" + + def setUp(self): + super().setUp() + env = mock.patch.dict(os.environ, { + 'APPLICATION_ID': 's~my-proj', + 'GOOGLE_CLOUD_PROJECT': 'my-proj', + 'LOCATION_ID': 'us-central1', + }) + env.start() + self.addCleanup(env.stop) + self.testbed = testbed.Testbed() + self.testbed.activate() + self.addCleanup(self.testbed.deactivate) + self.testbed.init_datastore_v3_stub() + self.testbed.init_memcache_stub() + cloudtask._reset_clients() + self.addCleanup(cloudtask._reset_clients) + self.mock_client = mock.Mock() + self.mock_client.queue_path.side_effect = ( + lambda p, r, q: f'projects/{p}/locations/{r}/queues/{q}' + ) + self.mock_client.task_path.side_effect = ( + lambda p, r, q, t: f'projects/{p}/locations/{r}/queues/{q}/tasks/{t}' + ) + client_patch = mock.patch( + 'google.cloud.tasks_v2.CloudTasksClient', return_value=self.mock_client) + client_patch.start() + self.addCleanup(client_patch.stop) + + def _pending_entities(self): + return list( + datastore.Query(cloudtask_transactional._PENDING_TASK_KIND).Run()) + + def _put_pending(self, **props): + entity = datastore.Entity(cloudtask_transactional._PENDING_TASK_KIND) + entity['queue_name'] = 'default' + entity['cloud_task_name'] = 'tx-1' + entity['cloud_task_payload'] = datastore_types.Text(json.dumps( + {'task': {'appEngineHttpRequest': {'relativeUri': '/w'}}})) + entity['status'] = cloudtask_transactional._TX_TASK_STATUS_PENDING + entity['created'] = datetime.datetime.utcnow() + entity['retry_count'] = 0 + entity.update(props) + datastore.Put(entity) + return entity.key() + + def test_sweep_dispatches_tasks_staged_by_other_sdks(self): + # Staged entities are shared with the Go and Java SDKs, so the sweeper + # dispatches tasks they staged, and tasks staged by the preview release. + queue = 'projects/my-proj/locations/us-central1/queues/default' + old = datetime.datetime.utcnow() - datetime.timedelta(minutes=5) + # As written by the Go SDK. + go_payload = ( + '{"task":{"name":"' + queue + '/tasks/tx-go","appEngineHttpRequest":' + '{"httpMethod":"POST","appEngineRouting":{"service":"worker",' + '"version":"v2"},"relativeUri":"/go","headers":{"X-Custom":"g"},' + '"body":"Z28tYm9keQ=="},"scheduleTime":"2026-10-07T20:00:00.250Z",' + '"retryConfig":{"maxAttempts":2,"minBackoff":"1.500s"}}}') + self._put_pending( + cloud_task_name='tx-go', created=old, sdk_lang='GO', + cloud_task_payload=go_payload, lock_expires=None) + # As written by the Java SDK, which stores payloads of up to 1500 bytes + # as indexed strings. fieldAddedLater stands for a Task field that is + # newer than the installed google-cloud-tasks. + java_payload = ( + '{"task":{"name":"' + queue + '/tasks/tx-java","appEngineHttpRequest":' + '{"appEngineRouting":{"service":"worker"},"httpMethod":"DELETE",' + '"relativeUri":"/java","headers":{}},"scheduleTime":' + '"2026-10-07T20:00:00Z","fieldAddedLater":1}}') + self._put_pending( + cloud_task_name='tx-java', created=old, sdk_lang='JAVA', + cloud_task_payload=java_payload, lock_expires=None, last_error='', + handled_by_sweeper=False) + # As written by the preview release of this SDK. + preview = datastore.Entity(cloudtask_transactional._PENDING_TASK_KIND) + preview['task_name'] = 'tx-preview' + preview['queue_name'] = 'default' + preview['payload'] = datastore_types.Text(json.dumps({ + 'name': queue + '/tasks/tx-preview', + 'app_engine_http_request': { + 'relative_uri': '/preview', + 'body': base64.b64encode(b'preview-body').decode('utf-8'), + }, + 'schedule_time': {'seconds': 1791403200, 'nanos': 0}, + })) + preview['status'] = cloudtask_transactional._TX_TASK_STATUS_PENDING + preview['created'] = old + preview['retry_count'] = 0 + datastore.Put(preview) + + cloudtask_transactional.sweep() + + sent = {} + for call in self.mock_client.create_task.call_args_list: + task = call[1]['request']['task'] + sent[task.app_engine_http_request.relative_uri] = task + self.assertEqual(sorted(sent), ['/go', '/java', '/preview']) + go_task = sent['/go'] + self.assertEqual(go_task.name, queue + '/tasks/tx-go') + self.assertEqual(go_task.app_engine_http_request.body, b'go-body') + self.assertEqual(go_task.app_engine_http_request.headers['X-Custom'], 'g') + self.assertEqual( + go_task.app_engine_http_request.app_engine_routing.version, 'v2') + self.assertEqual(go_task.retry_config.min_backoff.total_seconds(), 1.5) + self.assertEqual(go_task.schedule_time.timestamp(), 1791403200.25) + self.assertEqual(sent['/java'].app_engine_http_request.http_method, + tasks_v2.HttpMethod.DELETE) + self.assertEqual(sent['/preview'].app_engine_http_request.body, + b'preview-body') + self.assertEqual(sent['/preview'].schedule_time.timestamp(), 1791403200) + self.assertEqual(self._pending_entities(), []) + + def test_large_payload_is_staged_and_dispatched_after_commit(self): + body = b'x' * 5000 # Larger than the 1500-byte indexed string limit. + + def txn(): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w', payload=body)], multiple=False) + + datastore.RunInTransaction(txn) + + self.mock_client.create_task.assert_called_once() + sent = self.mock_client.create_task.call_args[1]['request']['task'] + self.assertEqual(sent.app_engine_http_request.body, body) + self.assertEqual(self._pending_entities(), []) + + def test_rolled_back_db_transaction_stages_nothing(self): + + def txn(): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False) + raise datastore_errors.Rollback() + + datastore.RunInTransaction(txn) + + self.mock_client.create_task.assert_not_called() + self.assertEqual(self._pending_entities(), []) + + def test_ndb_transaction_dispatches_after_commit(self): + + @ndb.transactional + def txn(): + with mock.patch.dict( + os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true'}): + taskqueue.add(url='/w', transactional=True) + self.mock_client.create_task.assert_not_called() + + txn() + + self.mock_client.create_task.assert_called_once() + self.assertEqual(self._pending_entities(), []) + + def test_rolled_back_ndb_transaction_stages_nothing(self): + + @ndb.transactional + def txn(): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False) + raise ndb.Rollback() + + txn() + + self.mock_client.create_task.assert_not_called() + self.assertEqual(self._pending_entities(), []) + + def test_dispatch_skips_locked_and_deleted_entities(self): + locked_until = datetime.datetime.utcnow() + datetime.timedelta(minutes=1) + key = self._put_pending( + status=cloudtask_transactional._TX_TASK_STATUS_PROCESSING, + lock_expires=locked_until) + cloudtask_transactional._dispatch_pending_keys_now([key]) + self.mock_client.create_task.assert_not_called() + + datastore.Delete(key) + cloudtask_transactional._dispatch_pending_keys_now([key]) + self.mock_client.create_task.assert_not_called() + self.assertEqual(self._pending_entities(), []) + + def test_failed_dispatch_does_not_recreate_entity_deleted_meanwhile(self): + key = self._put_pending() + + def create_task(request): + del request + datastore.Delete(key) # Another dispatcher finished this task. + raise google_exceptions.ServiceUnavailable('try again') + + self.mock_client.create_task.side_effect = create_task + cloudtask_transactional._dispatch_pending_keys_now([key]) + self.assertEqual(self._pending_entities(), []) + + def test_failed_dispatch_is_retried_then_marked_failed(self): + key = self._put_pending() + self.mock_client.create_task.side_effect = ( + google_exceptions.ServiceUnavailable('down')) + for _ in range(cloudtask_transactional._SWEEPER_MAX_RETRIES): + cloudtask_transactional._dispatch_pending_keys_now([key]) + entity = datastore.Get(key) + self.assertEqual( + entity['status'], cloudtask_transactional._TX_TASK_STATUS_FAILED) + self.assertIsNotNone(entity['failed_at']) + + def test_sweep_deletes_failed_entities_after_retention(self): + now = datetime.datetime.utcnow() + failed = cloudtask_transactional._TX_TASK_STATUS_FAILED + old = self._put_pending( + status=failed, failed_at=now - datetime.timedelta(days=8)) + recent = self._put_pending( + status=failed, failed_at=now - datetime.timedelta(days=1)) + + cloudtask_transactional.sweep() + + remaining = [e.key() for e in self._pending_entities()] + self.assertEqual(remaining, [recent]) + self.assertNotIn(old, remaining) + self.mock_client.create_task.assert_not_called() + + + def test_xg_transaction_with_entity_and_task_dispatches_after_commit(self): + # Staged tasks are root entities, so a transaction that also writes an + # application entity must be cross-group, as documented. + + def db_txn(): + datastore.Put(datastore.Entity('Order')) + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False) + + datastore.RunInTransactionOptions( + datastore.CreateTransactionOptions(xg=True), db_txn) + + class Order(ndb.Model): + pass + + @ndb.transactional(xg=True) + def ndb_txn(): + Order().put() + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False) + + ndb_txn() + + self.assertEqual(self.mock_client.create_task.call_count, 2) + self.assertEqual(self._pending_entities(), []) + self.assertEqual(len(list(datastore.Query('Order').Run())), 2) + + def non_xg_txn(): + datastore.Put(datastore.Entity('Order')) + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False) + + with self.assertRaises(datastore_errors.BadRequestError): + datastore.RunInTransaction(non_xg_txn) + self.assertEqual(self.mock_client.create_task.call_count, 2) + + def test_more_than_five_tasks_in_a_transaction_are_rejected(self): + max_tasks = cloudtask_transactional._MAX_TASKS_PER_TRANSACTION + + def add(count): + cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w') for _ in range(count)], + multiple=True) + + # Each staged task is its own entity group, so these transactions are XG. + xg = datastore.CreateTransactionOptions(xg=True) + + # The limit applies across calls in one transaction, like legacy. + def db_txn(): + add(3) + add(3) + + @ndb.transactional(xg=True) + def ndb_txn(): + add(max_tasks + 1) + + for txn in (lambda: datastore.RunInTransactionOptions(xg, db_txn), + ndb_txn): + with self.assertRaises(taskqueue.DatastoreError) as cm: + txn() + self.assertIsInstance(cm.exception, datastore_errors.BadRequestError) + self.mock_client.create_task.assert_not_called() + self.assertEqual(self._pending_entities(), []) + + def five_calls(): + for _ in range(max_tasks): + add(1) + + datastore.RunInTransactionOptions(xg, five_calls) + self.assertEqual(self.mock_client.create_task.call_count, max_tasks) + self.assertEqual(self._pending_entities(), []) + + @mock.patch.object(cloudtask_transactional, '_SWEEPER_BATCH_SIZE', 2) + def test_sweep_reads_a_bounded_batch_of_each_status(self): + now = datetime.datetime.utcnow() + old = now - datetime.timedelta(minutes=5) + failed = cloudtask_transactional._TX_TASK_STATUS_FAILED + pending = cloudtask_transactional._TX_TASK_STATUS_PENDING + # Failed entities kept for inspection must not crowd out pending ones. + for _ in range(3): + self._put_pending(status=failed, failed_at=now) + for _ in range(3): + self._put_pending(created=old) + + cloudtask_transactional.sweep() + + self.assertEqual(self.mock_client.create_task.call_count, 2) + statuses = sorted(e['status'] for e in self._pending_entities()) + self.assertEqual(statuses, [failed] * 3 + [pending]) + + cloudtask_transactional.sweep() # The next run takes the rest. + + self.assertEqual(self.mock_client.create_task.call_count, 3) + statuses = [e['status'] for e in self._pending_entities()] + self.assertEqual(statuses, [failed] * 3) + + +if __name__ == '__main__': + unittest.main() From 75fd652105691703ecc3ffb0415aa246e28ae41a Mon Sep 17 00:00:00 2001 From: Riddhi Shivhare Date: Thu, 8 Oct 2026 08:48:02 +0000 Subject: [PATCH 2/4] fix(taskqueue): reject batch add >100 and delete >1000 instead of chunking --- .../appengine/api/taskqueue/cloudtask.py | 267 ++++++++---------- .../appengine/api/taskqueue/taskqueue.py | 5 + .../appengine/api/taskqueue/cloudtask_test.py | 36 +-- 3 files changed, 138 insertions(+), 170 deletions(-) diff --git a/src/google/appengine/api/taskqueue/cloudtask.py b/src/google/appengine/api/taskqueue/cloudtask.py index d2fb62d..a5769dd 100644 --- a/src/google/appengine/api/taskqueue/cloudtask.py +++ b/src/google/appengine/api/taskqueue/cloudtask.py @@ -16,7 +16,6 @@ import base64 from concurrent import futures -import contextvars import datetime import functools import http @@ -175,6 +174,11 @@ def is_cloudtask_push_queue_enabled(): @_translate_errors def create_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): """Creates one or more tasks using Cloud Tasks API (supporting BatchCreateTasks).""" + if len(tasks) > _BATCH_CREATE_TASKS_MAX_SIZE: + raise taskqueue.TooManyTasksError( + 'No more than %d tasks can be added in a single call' + % _BATCH_CREATE_TASKS_MAX_SIZE + ) for task in tasks: if task.was_enqueued: raise taskqueue.BadTaskStateError('The task has already been enqueued.') @@ -190,6 +194,13 @@ def create_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): @_translate_errors def delete_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): """Deletes tasks from a queue using Cloud Tasks Client SDK (supporting BatchDeleteTasks).""" + if not tasks: + return [] if multiple else None + if len(tasks) > _BATCH_DELETE_TASKS_MAX_SIZE: + raise taskqueue.TooManyTasksError( + 'No more than %d tasks can be deleted in a single call' + % _BATCH_DELETE_TASKS_MAX_SIZE + ) deadline_at = _deadline_at(deadline) client = _get_client() project = _get_project_id() @@ -212,65 +223,42 @@ def delete_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): ) task_names_set.add(task.name) - def _delete_batch_chunk(batch): - task_names = [ - client.task_path(project, region, queue_name, t.name) for t in batch - ] - op = client.batch_delete_tasks( - request={'parent': parent, 'names': task_names}, - **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) - ) - _extract_lro_response(op) - metadata = getattr(op, 'metadata', None) - failed_requests = ( - getattr( - metadata, - 'failed_requests', - getattr(metadata, 'failedRequests', None), - ) - if metadata - else None - ) - - chunk_exception = None - for idx, t in enumerate(batch): - error_status = _get_failed_request(failed_requests, idx) - if error_status: - code = getattr(error_status, 'code', None) - tq_code = _map_rest_code_to_tq_code(code) - if tq_code in [ - taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK, - taskqueue_service_pb2.TaskQueueServiceError.TOMBSTONED_TASK, - ]: - t._Task__deleted = False - elif chunk_exception is None: - chunk_exception = taskqueue._TranslateError( - tq_code, getattr(error_status, 'message', '') - ) - else: - t._Task__deleted = True - - return chunk_exception - - chunks = [ - tasks[i : i + _BATCH_DELETE_TASKS_MAX_SIZE] - for i in range(0, len(tasks), _BATCH_DELETE_TASKS_MAX_SIZE) + task_names = [ + client.task_path(project, region, queue_name, t.name) for t in tasks ] + op = client.batch_delete_tasks( + request={'parent': parent, 'names': task_names}, + **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) + ) + _extract_lro_response(op) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', None), + ) + if metadata + else None + ) + exception = None - if len(chunks) == 1: - exception = _delete_batch_chunk(chunks[0]) - elif chunks: - chunk_futures = [ - _THREAD_POOL.submit( - contextvars.copy_context().run, _delete_batch_chunk, batch + for idx, t in enumerate(tasks): + error_status = _get_failed_request(failed_requests, idx) + if error_status: + code = getattr(error_status, 'code', None) + tq_code = _map_rest_code_to_tq_code(code) + if tq_code in [ + taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK, + taskqueue_service_pb2.TaskQueueServiceError.TOMBSTONED_TASK, + ]: + t._Task__deleted = False + elif exception is None: + exception = taskqueue._TranslateError( + tq_code, getattr(error_status, 'message', '') ) - for batch in chunks - ] - futures.wait(chunk_futures) - for fut in chunk_futures: - chunk_exc = fut.result() - if exception is None and chunk_exc is not None: - exception = chunk_exc + else: + t._Task__deleted = True if exception is not None: raise exception @@ -772,7 +760,7 @@ def _set_preferred_exception(current, candidate): def _create_batch_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline_at=None): - """Helper to create tasks in batches using CloudTasksClient BatchCreateTasks API.""" + """Helper to create tasks in a batch using CloudTasksClient BatchCreateTasks API.""" client = _get_client() project = _get_project_id() region = _get_region() @@ -789,104 +777,79 @@ def _create_batch_tasks_in_cloud_tasks(queue_name, tasks, multiple, ) task_names.add(task.name) - def _create_batch_chunk(batch): - requests_payload = [ - { - 'parent': parent, - 'task': _build_ct_task_payload( - queue_name, t, client, project, region - ), - } - for t in batch - ] - try: - op = client.batch_create_tasks( - request={'parent': parent, 'requests': requests_payload}, - **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) - ) - metadata = getattr(op, 'metadata', None) - failed_requests = ( - getattr( - metadata, - 'failed_requests', - getattr(metadata, 'failedRequests', None), - ) - if metadata - else None - ) - response = _extract_lro_response(op) - response_tasks = getattr(response, 'tasks', []) if response else [] - - res_iter = iter(response_tasks) - chunk_created = [] - chunk_exception = None - missing_task_idx = None - - for idx, t in enumerate(batch): - error_status = _get_failed_request(failed_requests, idx) - if error_status: - code = getattr(error_status, 'code', None) - tq_code = _map_rest_code_to_tq_code(code) - err = taskqueue._TranslateError( - tq_code, getattr(error_status, 'message', '') - ) - chunk_exception = _set_preferred_exception(chunk_exception, err) - else: - try: - res_task = next(res_iter) - task_id = ( - res_task.name.split('/')[-1] - if hasattr(res_task, 'name') - else res_task['name'].split('/')[-1] - ) - t._Task__name = task_id - t._Task__queue_name = queue_name - t._Task__enqueued = True - chunk_created.append(t) - except StopIteration: - if missing_task_idx is None: - missing_task_idx = idx - - if chunk_exception is None and missing_task_idx is not None: - chunk_exception = taskqueue.InternalError( - 'BatchCreateTasks response is missing created task at index %d' - % missing_task_idx + requests_payload = [ + { + 'parent': parent, + 'task': _build_ct_task_payload( + queue_name, t, client, project, region + ), + } + for t in tasks + ] + try: + op = client.batch_create_tasks( + request={'parent': parent, 'requests': requests_payload}, + **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) + ) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', None), ) + if metadata + else None + ) + response = _extract_lro_response(op) + response_tasks = getattr(response, 'tasks', []) if response else [] - return chunk_created, chunk_exception - except google_exceptions.Aborted: - # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. - raise - except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: - raise taskqueue.TaskAlreadyExistsError(str(e)) - except google_exceptions.NotFound as e: - raise taskqueue.UnknownQueueError(str(e)) - except google_exceptions.BadRequest as e: - if 'Queue does not exist' in str(e): - raise taskqueue.UnknownQueueError(str(e)) - raise e + res_iter = iter(response_tasks) + created_tasks = [] + exception = None + missing_task_idx = None - chunks = [ - tasks[i : i + _BATCH_CREATE_TASKS_MAX_SIZE] - for i in range(0, len(tasks), _BATCH_CREATE_TASKS_MAX_SIZE) - ] - created_tasks = [] - exception = None - if len(chunks) == 1: - chunk_created, exception = _create_batch_chunk(chunks[0]) - created_tasks.extend(chunk_created) - elif chunks: - chunk_futures = [ - _THREAD_POOL.submit( - contextvars.copy_context().run, _create_batch_chunk, batch + for idx, t in enumerate(tasks): + error_status = _get_failed_request(failed_requests, idx) + if error_status: + code = getattr(error_status, 'code', None) + tq_code = _map_rest_code_to_tq_code(code) + err = taskqueue._TranslateError( + tq_code, getattr(error_status, 'message', '') ) - for batch in chunks - ] - futures.wait(chunk_futures) - for fut in chunk_futures: - chunk_created, chunk_exc = fut.result() - created_tasks.extend(chunk_created) - exception = _set_preferred_exception(exception, chunk_exc) + exception = _set_preferred_exception(exception, err) + else: + try: + res_task = next(res_iter) + task_id = ( + res_task.name.split('/')[-1] + if hasattr(res_task, 'name') + else res_task['name'].split('/')[-1] + ) + t._Task__name = task_id + t._Task__queue_name = queue_name + t._Task__enqueued = True + created_tasks.append(t) + except StopIteration: + if missing_task_idx is None: + missing_task_idx = idx + + if exception is None and missing_task_idx is not None: + exception = taskqueue.InternalError( + 'BatchCreateTasks response is missing created task at index %d' + % missing_task_idx + ) + except google_exceptions.Aborted: + # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. + raise + except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: + raise taskqueue.TaskAlreadyExistsError(str(e)) + except google_exceptions.NotFound as e: + raise taskqueue.UnknownQueueError(str(e)) + except google_exceptions.BadRequest as e: + if 'Queue does not exist' in str(e): + raise taskqueue.UnknownQueueError(str(e)) + raise e if exception is not None: raise exception diff --git a/src/google/appengine/api/taskqueue/taskqueue.py b/src/google/appengine/api/taskqueue/taskqueue.py index a4f005b..b264750 100755 --- a/src/google/appengine/api/taskqueue/taskqueue.py +++ b/src/google/appengine/api/taskqueue/taskqueue.py @@ -2159,6 +2159,11 @@ def add_async(self, task, transactional=False, rpc=None): raise InvalidTaskError( 'You cannot add both push and pull tasks in a single call.') + if len(tasks) > MAX_TASKS_PER_ADD: + raise TooManyTasksError( + 'No more than %d tasks can be added in a single call' % + MAX_TASKS_PER_ADD) + # Intercept for Cloud Tasks backend if (cloudtask.is_cloudtask_push_queue_enabled() and has_push_task diff --git a/tests/google/appengine/api/taskqueue/cloudtask_test.py b/tests/google/appengine/api/taskqueue/cloudtask_test.py index 91263ac..41bc9f6 100644 --- a/tests/google/appengine/api/taskqueue/cloudtask_test.py +++ b/tests/google/appengine/api/taskqueue/cloudtask_test.py @@ -378,28 +378,20 @@ def test_create_batch_tasks_with_real_lro_operation(self, mock_client_cls): with self.assertRaises(taskqueue.BadTaskStateError): cloudtask.create_tasks_in_cloud_tasks('default', [t1], multiple=False) - # Multi-chunk batch execution (>100 tasks) preserves order and names. - def fake_batch_create(request, **_): - reqs = request['requests'] - return _make_batch_create_operation( - task_names=[ - f"projects/my-proj/locations/us-central1/queues/default/tasks/{r['task']['app_engine_http_request']['relative_uri'].lstrip('/')}" - for r in reqs - ] + # Exceeding MAX_TASKS_PER_ADD (100) raises TooManyTasksError. + many_tasks = [taskqueue.Task(url=f'/t-{i}') for i in range(101)] + with self.assertRaises(taskqueue.TooManyTasksError): + cloudtask.create_tasks_in_cloud_tasks( + 'default', many_tasks, multiple=True ) - - self.mock_client.batch_create_tasks.reset_mock() - self.mock_client.batch_create_tasks.side_effect = fake_batch_create - many_tasks = [taskqueue.Task(url=f'/t-{i}') for i in range(250)] - many_res = cloudtask.create_tasks_in_cloud_tasks( - 'default', many_tasks, multiple=True - ) - self.assertEqual(self.mock_client.batch_create_tasks.call_count, 3) - self.assertEqual([t.name for t in many_res], [f't-{i}' for i in range(250)]) + with mock.patch.dict( + os.environ, {cloudtask.ENV_USE_CLOUDTASK_PUSH_QUEUE: 'true'} + ): + with self.assertRaises(taskqueue.TooManyTasksError): + taskqueue.Queue('default').add(many_tasks) # Missing task in BatchCreateTasks response raises InternalError instead of # silently leaving task.name as None. - self.mock_client.batch_create_tasks.side_effect = None self.mock_client.batch_create_tasks.return_value = ( _make_batch_create_operation(task_names=[]) ) @@ -479,6 +471,14 @@ def test_delete_tasks_in_cloud_tasks(self, mock_client_cls): 'default', [taskqueue.Task(name='d'), taskqueue.Task(name='d')], multiple=True ) + # Exceeding _BATCH_DELETE_TASKS_MAX_SIZE (1000) raises TooManyTasksError. + with self.assertRaises(taskqueue.TooManyTasksError): + cloudtask.delete_tasks_in_cloud_tasks( + 'default', + [taskqueue.Task(name=f'del-{i}') for i in range(1001)], + multiple=True, + ) + # NOT_FOUND in failed_requests marks was_deleted=False without raising self.mock_client.batch_delete_tasks.return_value = _make_batch_delete_operation( failed_requests={ From d60df9dcba73214e14a125273157c1fd84a99c2d Mon Sep 17 00:00:00 2001 From: Riddhi Shivhare Date: Thu, 8 Oct 2026 09:34:54 +0000 Subject: [PATCH 3/4] fix(taskqueue): enforce done=True and timeout=0 on batch LRO extraction --- .../appengine/api/taskqueue/cloudtask.py | 10 +++++++++- .../appengine/api/taskqueue/cloudtask_test.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/google/appengine/api/taskqueue/cloudtask.py b/src/google/appengine/api/taskqueue/cloudtask.py index a5769dd..80f914a 100644 --- a/src/google/appengine/api/taskqueue/cloudtask.py +++ b/src/google/appengine/api/taskqueue/cloudtask.py @@ -739,11 +739,19 @@ def _create_single_task_in_cloud_tasks(queue_name, task, multiple, def _extract_lro_response(op): """Extracts the unpacked response message from a synchronous LRO.""" + raw_op = getattr(op, 'operation', None) + if raw_op is not None and not getattr(raw_op, 'done', True): + raise taskqueue.InternalError( + 'Cloud Tasks batch operation returned done=False' + ) if hasattr(op, 'result') and callable(op.result): if (not isinstance(op, futures.Future) and getattr(op, 'response', None) is not None): return op.response - return op.result() + try: + return op.result(timeout=0) + except TypeError: + return op.result() return getattr(op, 'response', None) diff --git a/tests/google/appengine/api/taskqueue/cloudtask_test.py b/tests/google/appengine/api/taskqueue/cloudtask_test.py index 41bc9f6..dcd8638 100644 --- a/tests/google/appengine/api/taskqueue/cloudtask_test.py +++ b/tests/google/appengine/api/taskqueue/cloudtask_test.py @@ -402,6 +402,25 @@ def test_create_batch_tasks_with_real_lro_operation(self, mock_client_cls): multiple=True, ) + # An incomplete Operation (done=False) raises InternalError immediately + # without polling Operations.GetOperation. + refresh_mock = mock.Mock() + incomplete_op = operation.Operation( + operations_pb2.Operation(name='operations/incomplete', done=False), + refresh=refresh_mock, + cancel=mock.Mock(), + result_type=tasks_v2.BatchCreateTasksResponse, + metadata_type=tasks_v2.BatchCreateTasksMetadata, + ) + self.mock_client.batch_create_tasks.return_value = incomplete_op + with self.assertRaises(taskqueue.InternalError): + cloudtask.create_tasks_in_cloud_tasks( + 'default', + [taskqueue.Task(url='/1'), taskqueue.Task(url='/2')], + multiple=True, + ) + refresh_mock.assert_not_called() + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) @mock.patch('google.cloud.tasks_v2.CloudTasksClient') def test_create_batch_tasks_partial_failure_and_duplicate_names(self, mock_client_cls): From 3cf36092c62554ba639e2bdba209aa54a97c4df6 Mon Sep 17 00:00:00 2001 From: Riddhi Shivhare Date: Thu, 8 Oct 2026 09:46:12 +0000 Subject: [PATCH 4/4] Fix LRO delete metadata extraction order, default namespace for pending tasks, and TombstonedTaskError translation --- .../appengine/api/taskqueue/cloudtask.py | 172 ++++++++---------- .../api/taskqueue/cloudtask_transactional.py | 6 +- .../appengine/api/taskqueue/cloudtask_test.py | 77 +++++++- 3 files changed, 157 insertions(+), 98 deletions(-) diff --git a/src/google/appengine/api/taskqueue/cloudtask.py b/src/google/appengine/api/taskqueue/cloudtask.py index 80f914a..17e8a43 100644 --- a/src/google/appengine/api/taskqueue/cloudtask.py +++ b/src/google/appengine/api/taskqueue/cloudtask.py @@ -125,6 +125,9 @@ def _to_taskqueue_error(error): return taskqueue.TaskAlreadyExistsError(detail) if isinstance(error, google_exceptions.NotFound): return taskqueue.UnknownQueueError(detail) + if (isinstance(error, google_exceptions.BadRequest) + and 'queue does not exist' in detail.lower()): + return taskqueue.UnknownQueueError(detail) if isinstance(error, (google_exceptions.PermissionDenied, google_exceptions.Unauthenticated)): return taskqueue.PermissionDeniedError(detail) @@ -230,7 +233,6 @@ def delete_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): request={'parent': parent, 'names': task_names}, **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) ) - _extract_lro_response(op) metadata = getattr(op, 'metadata', None) failed_requests = ( getattr( @@ -241,22 +243,22 @@ def delete_tasks_in_cloud_tasks(queue_name, tasks, multiple, deadline=None): if metadata else None ) + _extract_lro_response(op, has_failed_requests=bool(failed_requests)) exception = None for idx, t in enumerate(tasks): error_status = _get_failed_request(failed_requests, idx) if error_status: code = getattr(error_status, 'code', None) - tq_code = _map_rest_code_to_tq_code(code) + msg = getattr(error_status, 'message', '') + tq_code = _map_rest_code_to_tq_code(code, msg) if tq_code in [ taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK, taskqueue_service_pb2.TaskQueueServiceError.TOMBSTONED_TASK, ]: t._Task__deleted = False elif exception is None: - exception = taskqueue._TranslateError( - tq_code, getattr(error_status, 'message', '') - ) + exception = taskqueue._TranslateError(tq_code, msg) else: t._Task__deleted = True @@ -710,34 +712,20 @@ def _create_single_task_in_cloud_tasks(queue_name, task, multiple, parent = client.queue_path(project, region, queue_name) ct_task = _build_ct_task_payload(queue_name, task, client, project, region) - try: - response_task = client.create_task( - request={'parent': parent, 'task': ct_task}, - **_timeout_kwargs(deadline_at)) - task_id = response_task.name.split('/')[-1] - task._Task__name = task_id - task._Task__queue_name = queue_name - task._Task__enqueued = True - if multiple: - return [task] - else: - return task - except google_exceptions.Aborted: - # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. - raise - except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: - raise taskqueue.TaskAlreadyExistsError(str(e)) - except google_exceptions.NotFound as e: - raise taskqueue.UnknownQueueError(str(e)) - except google_exceptions.BadRequest as e: - if 'Queue does not exist' in str(e): - raise taskqueue.UnknownQueueError(str(e)) - raise e - except Exception as e: - raise e + response_task = client.create_task( + request={'parent': parent, 'task': ct_task}, + **_timeout_kwargs(deadline_at)) + task_id = response_task.name.split('/')[-1] + task._Task__name = task_id + task._Task__queue_name = queue_name + task._Task__enqueued = True + if multiple: + return [task] + else: + return task -def _extract_lro_response(op): +def _extract_lro_response(op, has_failed_requests=False): """Extracts the unpacked response message from a synchronous LRO.""" raw_op = getattr(op, 'operation', None) if raw_op is not None and not getattr(raw_op, 'done', True): @@ -752,6 +740,10 @@ def _extract_lro_response(op): return op.result(timeout=0) except TypeError: return op.result() + except Exception: # pylint: disable=broad-except + if has_failed_requests: + return None + raise return getattr(op, 'response', None) @@ -794,70 +786,57 @@ def _create_batch_tasks_in_cloud_tasks(queue_name, tasks, multiple, } for t in tasks ] - try: - op = client.batch_create_tasks( - request={'parent': parent, 'requests': requests_payload}, - **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) - ) - metadata = getattr(op, 'metadata', None) - failed_requests = ( - getattr( - metadata, - 'failed_requests', - getattr(metadata, 'failedRequests', None), + op = client.batch_create_tasks( + request={'parent': parent, 'requests': requests_payload}, + **_timeout_kwargs(deadline_at, _DEFAULT_RPC_TIMEOUT_SECONDS) + ) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', None), + ) + if metadata + else None + ) + response = _extract_lro_response(op, has_failed_requests=bool(failed_requests)) + response_tasks = getattr(response, 'tasks', []) if response else [] + + res_iter = iter(response_tasks) + created_tasks = [] + exception = None + missing_task_idx = None + + for idx, t in enumerate(tasks): + error_status = _get_failed_request(failed_requests, idx) + if error_status: + code = getattr(error_status, 'code', None) + msg = getattr(error_status, 'message', '') + tq_code = _map_rest_code_to_tq_code(code, msg) + err = taskqueue._TranslateError(tq_code, msg) + exception = _set_preferred_exception(exception, err) + else: + try: + res_task = next(res_iter) + task_id = ( + res_task.name.split('/')[-1] + if hasattr(res_task, 'name') + else res_task['name'].split('/')[-1] ) - if metadata - else None + t._Task__name = task_id + t._Task__queue_name = queue_name + t._Task__enqueued = True + created_tasks.append(t) + except StopIteration: + if missing_task_idx is None: + missing_task_idx = idx + + if exception is None and missing_task_idx is not None: + exception = taskqueue.InternalError( + 'BatchCreateTasks response is missing created task at index %d' + % missing_task_idx ) - response = _extract_lro_response(op) - response_tasks = getattr(response, 'tasks', []) if response else [] - - res_iter = iter(response_tasks) - created_tasks = [] - exception = None - missing_task_idx = None - - for idx, t in enumerate(tasks): - error_status = _get_failed_request(failed_requests, idx) - if error_status: - code = getattr(error_status, 'code', None) - tq_code = _map_rest_code_to_tq_code(code) - err = taskqueue._TranslateError( - tq_code, getattr(error_status, 'message', '') - ) - exception = _set_preferred_exception(exception, err) - else: - try: - res_task = next(res_iter) - task_id = ( - res_task.name.split('/')[-1] - if hasattr(res_task, 'name') - else res_task['name'].split('/')[-1] - ) - t._Task__name = task_id - t._Task__queue_name = queue_name - t._Task__enqueued = True - created_tasks.append(t) - except StopIteration: - if missing_task_idx is None: - missing_task_idx = idx - - if exception is None and missing_task_idx is not None: - exception = taskqueue.InternalError( - 'BatchCreateTasks response is missing created task at index %d' - % missing_task_idx - ) - except google_exceptions.Aborted: - # Aborted subclasses Conflict (HTTP 409) but does not mean the task exists. - raise - except (google_exceptions.AlreadyExists, google_exceptions.Conflict) as e: - raise taskqueue.TaskAlreadyExistsError(str(e)) - except google_exceptions.NotFound as e: - raise taskqueue.UnknownQueueError(str(e)) - except google_exceptions.BadRequest as e: - if 'Queue does not exist' in str(e): - raise taskqueue.UnknownQueueError(str(e)) - raise e if exception is not None: raise exception @@ -882,8 +861,13 @@ def _get_failed_request(failed_requests, idx): return None -def _map_rest_code_to_tq_code(code): +def _map_rest_code_to_tq_code(code, message=''): """Maps gRPC / HTTP error status codes to legacy TaskQueue error enum codes.""" + if message: + match = _EXECUTOR_ERROR_RE.search(message) + error_codes = taskqueue_service_pb2.TaskQueueServiceError.ErrorCode + if match and match.group(1) in error_codes.keys(): + return error_codes.Value(match.group(1)) if code in [code_pb2.NOT_FOUND, http.HTTPStatus.NOT_FOUND]: return taskqueue_service_pb2.TaskQueueServiceError.UNKNOWN_TASK if code in [code_pb2.INVALID_ARGUMENT, http.HTTPStatus.BAD_REQUEST]: diff --git a/src/google/appengine/api/taskqueue/cloudtask_transactional.py b/src/google/appengine/api/taskqueue/cloudtask_transactional.py index 602c831..fffc336 100644 --- a/src/google/appengine/api/taskqueue/cloudtask_transactional.py +++ b/src/google/appengine/api/taskqueue/cloudtask_transactional.py @@ -115,7 +115,7 @@ def add_transactional_tasks(queue_name, tasks, multiple): queue_name, task ) - entity = datastore.Entity(_PENDING_TASK_KIND) + entity = datastore.Entity(_PENDING_TASK_KIND, namespace='') entity['queue_name'] = queue_name entity['cloud_task_name'] = generated_name # Text is unindexed; indexed strings are limited to 1500 bytes and task @@ -217,7 +217,9 @@ def sweep(): _TX_TASK_STATUS_FAILED): try: with _use_default_datastore_adapter(): - query = datastore.Query(_PENDING_TASK_KIND, {'status =': status}) + query = datastore.Query( + _PENDING_TASK_KIND, {'status =': status}, namespace='' + ) # Fetch inside the adapter context; Run() returns a lazy iterator. entities.extend(query.Run(limit=_SWEEPER_BATCH_SIZE)) except Exception as e: diff --git a/tests/google/appengine/api/taskqueue/cloudtask_test.py b/tests/google/appengine/api/taskqueue/cloudtask_test.py index dcd8638..379797a 100644 --- a/tests/google/appengine/api/taskqueue/cloudtask_test.py +++ b/tests/google/appengine/api/taskqueue/cloudtask_test.py @@ -34,6 +34,7 @@ from google.appengine.api import datastore from google.appengine.api import datastore_errors from google.appengine.api import datastore_types +from google.appengine.api import namespace_manager from google.appengine.api.taskqueue import cloudtask from google.appengine.api.taskqueue import cloudtask_transactional from google.appengine.api.taskqueue import taskqueue @@ -67,12 +68,15 @@ def _make_batch_create_operation(task_names=None, failed_requests=None): ) -def _make_batch_delete_operation(failed_requests=None): +def _make_batch_delete_operation(failed_requests=None, op_error=None): meta_pb = tasks_v2.BatchDeleteTasksMetadata.pb()( failed_requests=failed_requests or {} ) op_pb = operations_pb2.Operation(name='operations/batch-delete-1', done=True) - op_pb.response.Pack(empty_pb2.Empty()) + if op_error is not None: + op_pb.error.CopyFrom(op_error) + else: + op_pb.response.Pack(empty_pb2.Empty()) op_pb.metadata.Pack(meta_pb) return operation.from_gapic( op_pb, @@ -342,6 +346,14 @@ def test_create_single_task_success_and_errors(self, mock_client_cls): with self.assertRaises(taskqueue.TaskAlreadyExistsError): cloudtask.create_tasks_in_cloud_tasks('default', [taskqueue.Task(url='/worker')], multiple=False) + # AlreadyExists with ExecutorServiceError::TOMBSTONED_TASK -> TombstonedTaskError + self.mock_client.create_task.side_effect = google_exceptions.AlreadyExists( + 'Requested entity already exists [detail: "[ORIGINAL ERROR] ' + 'ExecutorError::1015: apphosting::ExecutorServiceError::TOMBSTONED_TASK"]' + ) + with self.assertRaises(taskqueue.TombstonedTaskError): + cloudtask.create_tasks_in_cloud_tasks('default', [taskqueue.Task(url='/worker')], multiple=False) + # NotFound -> UnknownQueueError self.mock_client.create_task.side_effect = google_exceptions.NotFound('missing queue') with self.assertRaises(taskqueue.UnknownQueueError): @@ -521,6 +533,26 @@ def test_delete_tasks_in_cloud_tasks(self, mock_client_cls): 'default', [taskqueue.Task(name='denied-task')], multiple=False ) + # Operation-level error alongside per-task NOT_FOUND in failed_requests still + # processes failed_requests and marks NOT_FOUND tasks as was_deleted=False. + self.mock_client.batch_delete_tasks.return_value = ( + _make_batch_delete_operation( + failed_requests={ + 0: status_pb2.Status( + code=code_pb2.NOT_FOUND, message='Not found' + ) + }, + op_error=status_pb2.Status( + code=code_pb2.NOT_FOUND, message='1 request failed' + ), + ) + ) + t_missing_with_op_err = taskqueue.Task(name='missing-with-op-err') + cloudtask.delete_tasks_in_cloud_tasks( + 'default', [t_missing_with_op_err], multiple=False + ) + self.assertFalse(t_missing_with_op_err.was_deleted) + @mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'my-proj', 'LOCATION_ID': 'us-central1'}) @mock.patch('google.cloud.tasks_v2.CloudTasksClient') def test_purge_queue_in_cloud_tasks(self, mock_client_cls): @@ -793,6 +825,10 @@ def test_api_errors_are_raised_as_taskqueue_errors( (google_exceptions.FailedPrecondition( 'Queue does not exist. apphosting::ExecutorServiceError::' 'UNKNOWN_QUEUE'), taskqueue.UnknownQueueError), + (google_exceptions.AlreadyExists( + 'Requested entity already exists [detail: "[ORIGINAL ERROR] ' + 'ExecutorError::1015: apphosting::ExecutorServiceError::' + 'TOMBSTONED_TASK"]'), taskqueue.TombstonedTaskError), ] calls = { 'create_task': lambda: cloudtask.create_tasks_in_cloud_tasks( @@ -1286,6 +1322,43 @@ def test_sweep_reads_a_bounded_batch_of_each_status(self): statuses = [e['status'] for e in self._pending_entities()] self.assertEqual(statuses, [failed] * 3) + def test_transactional_tasks_under_non_default_namespace_use_empty_namespace( + self, + ): + old = datetime.datetime.utcnow() - datetime.timedelta(minutes=5) + namespace_manager.set_namespace('tenant-a') + self.addCleanup(lambda: namespace_manager.set_namespace('')) + + # Simulate a transaction whose post-commit hook failed so the sweeper must + # find the staged entity from the default ('') namespace. + with mock.patch.object( + cloudtask_transactional, '_dispatch_pending_keys_now' + ): + datastore.RunInTransaction( + lambda: cloudtask_transactional.add_transactional_tasks( + 'default', [taskqueue.Task(url='/w')], multiple=False + ) + ) + + # Even while 'tenant-a' is active, the staged entity is stored in the '' + # namespace and visible when querying namespace=''. + staged = list( + datastore.Query( + cloudtask_transactional._PENDING_TASK_KIND, namespace='' + ).Run() + ) + self.assertEqual(len(staged), 1) + self.assertEqual(staged[0].namespace(), '') + staged[0]['created'] = old + datastore.Put(staged[0]) + + # Sweeper invoked under default namespace ('') or another namespace still + # finds and dispatches the entity. + namespace_manager.set_namespace('') + cloudtask_transactional.sweep() + self.mock_client.create_task.assert_called_once() + self.assertEqual(self._pending_entities(), []) + if __name__ == '__main__': unittest.main()