import base64 import datetime import random import re import string import sys from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER from .constants import FULL_QUEUE_NAME_RE from .protocols import taskqueue_service_pb2 sys.path.append(APPSCALE_PYTHON_APPSERVER) MAX_TAG_LENGTH = 500 MAX_TASK_NAME_LENGTH = 500 # A regex rule for validating task names. TASK_NAME_PATTERN = r'^[a-zA-Z0-9_-]{1,%s}$' % MAX_TASK_NAME_LENGTH # A compiled regex rule for validating task names. TASK_NAME_RE = re.compile(TASK_NAME_PATTERN) # Validation rules for queue parameters. QUEUE_ATTRIBUTE_RULES = { 'id': lambda name: TASK_NAME_RE.match(name), 'queueName': lambda name: FULL_QUEUE_NAME_RE.match(name), 'tag': lambda tag: tag is None or len(tag) <= MAX_TAG_LENGTH } # All possible fields to include in a task's JSON representation. TASK_FIELDS = ('kind', 'queueName', 'id', 'enqueueTimestamp', 'leaseTimestamp', 'payloadBase64', 'retry_count', 'tag') def parse_timestamp(timestamp): """ Parses timestamps used for creating tasks. Args: timestamp: Either a datetime object or an integer specifying the number of microseconds since the epoch. Returns: A datetime object. """ if isinstance(timestamp, datetime.datetime): return timestamp else: return datetime.datetime.utcfromtimestamp(int(timestamp) / 1000000.0) class InvalidTaskInfo(Exception): pass class Task(object): """ Represents a task created by an App Engine application. """ # Attributes that may not be defined. OPTIONAL_ATTRS = ['queueName', 'enqueueTimestamp', 'leaseTimestamp', 'tag', 'payloadBase64'] def __init__(self, task_info): """ Create a Task object. Args: task_info: A dictionary containing task info. """ self.retry_count = 0 if 'payloadBase64' in task_info: encoded_payload = task_info['payloadBase64'] # Google's REST API adds missing padding. missing_padding = 4 - (len(encoded_payload) % 4) encoded_payload += '=' * missing_padding # This decode/encode step is performed in order to match Google's # behavior in cases where the given payload does not have the correct # padding. It can be removed if we start storing the payload as binary # blobs. The conversion from unicode string to byte string is needed # because urlsafe_b64decode chokes on some invalid base64 that Google # accepts. payload = base64.urlsafe_b64decode(encoded_payload.encode('utf8')) self.payloadBase64 = base64.urlsafe_b64encode(payload) if 'id' in task_info and task_info['id']: self.id = task_info['id'] else: self.id = ''.join(random.choice(string.ascii_lowercase) for _ in range(20)) if 'queueName' in task_info: self.queueName = task_info['queueName'] if 'tag' in task_info: self.tag = task_info['tag'] if 'retry_count' in task_info: self.retry_count = task_info['retry_count'] if 'leaseTimestamp' in task_info: self.leaseTimestamp = parse_timestamp(task_info['leaseTimestamp']) if 'enqueueTimestamp' in task_info: self.enqueueTimestamp = parse_timestamp(task_info['enqueueTimestamp']) self.validate_info() def validate_info(self): """ Make sure the existing attributes are valid. Raises: InvalidTaskInfo if one of the attribute fails validation. """ for attribute, rule in QUEUE_ATTRIBUTE_RULES.items(): try: value = getattr(self, attribute) except AttributeError: continue if not rule(value): raise InvalidTaskInfo( 'Invalid task info: {}={}'.format(attribute, value)) def get_eta(self): """ Returns the ETA for a task. Raises: InvalidTaskInfo if ETA information is not set. """ epoch = datetime.datetime.utcfromtimestamp(0) if hasattr(self, 'leaseTimestamp') and self.leaseTimestamp != epoch: return self.leaseTimestamp try: return self.enqueueTimestamp except AttributeError: raise InvalidTaskInfo('No ETA info for {}'.format(self)) def expired(self, max_retries): """ Checks whether or not a task has expired. Args: max_retries: An integer specifying the queue's task retry limit. Returns: A boolean indicating whether or not the task has expired. """ if self.retry_count < max_retries: return False if self.leaseTimestamp >= datetime.datetime.utcnow(): return False return True def __repr__(self): """ Generates a string representation of the task. Returns: A string representing the task. """ attributes = {'id': self.id} if hasattr(self, 'queueName'): attributes['queue'] = self.queueName values = ', '.join(['='.join([attr, val]) for attr, val in attributes.items()]) return ''.format(values) def json_safe_dict(self, fields=TASK_FIELDS): """ Generate a JSON-safe dictionary representation of the task. Args: fields: A list of fields to include in the response. Returns: A JSON-safe dictionary representing the task. """ task = {} if 'kind' in fields: task['kind'] = 'taskqueues#task' if 'id' in fields: task['id'] = self.id if 'retry_count' in fields: task['retry_count'] = self.retry_count # Timestamps are represented as microseconds since the epoch. epoch = datetime.datetime.utcfromtimestamp(0) for attribute in self.OPTIONAL_ATTRS: if attribute not in fields or not hasattr(self, attribute): continue value = getattr(self, attribute) if isinstance(value, datetime.datetime): # All numbers are represented as strings in the GCP ecosystem for # Javascript compatibility reasons. We convert to string so that # the response can be successfully parsed by Google API clients. value = str(int((value - epoch).total_seconds() * 1000000)) task[attribute] = value if attribute == 'payloadBase64': task[attribute] = task[attribute].decode('utf-8') return task def encode_lease_pb(self): """ Encode this task as a protocol buffer response. Returns: A TaskQueueQueryAndOwnTasksResponse.Task object. """ task_pb = \ taskqueue_service_pb2.TaskQueueQueryAndOwnTasksResponse().task.add() task_pb.task_name = self.id.encode('utf-8') epoch = datetime.datetime.utcfromtimestamp(0) task_pb.eta_usec = int((self.get_eta() - epoch).total_seconds() * 1000000) task_pb.retry_count = self.retry_count task_pb.body = base64.urlsafe_b64decode(self.payloadBase64) try: task_pb.tag = self.tag.encode('utf-8') except AttributeError: pass return task_pb