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..17e8a43 --- /dev/null +++ b/src/google/appengine/api/taskqueue/cloudtask.py @@ -0,0 +1,880 @@ +# 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 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.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) + 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).""" + 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.') + 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).""" + 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() + 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) + + 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) + ) + metadata = getattr(op, 'metadata', None) + failed_requests = ( + getattr( + metadata, + 'failed_requests', + getattr(metadata, 'failedRequests', 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) + 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, msg) + else: + t._Task__deleted = True + + 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) + + 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, 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): + 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 + try: + 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) + + +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 a batch 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) + + requests_payload = [ + { + 'parent': parent, + 'task': _build_ct_task_payload( + queue_name, t, client, project, region + ), + } + for t in tasks + ] + 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] + ) + 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 + ) + + 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, 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]: + # 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..fffc336 --- /dev/null +++ b/src/google/appengine/api/taskqueue/cloudtask_transactional.py @@ -0,0 +1,508 @@ +# 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, namespace='') + 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}, namespace='' + ) + # 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..b264750 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,26 @@ 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 + 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 +2221,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 +2522,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..379797a --- /dev/null +++ b/tests/google/appengine/api/taskqueue/cloudtask_test.py @@ -0,0 +1,1364 @@ +# 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 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 +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, 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) + 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, + 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) + + # 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): + 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) + + # 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 + ) + 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.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, + ) + + # 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): + 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 + ) + + # 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={ + 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 + ) + + # 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): + 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), + (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( + '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) + + 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()