def add(url, name=None, queue_name='default', transactional=False, ignore_already=False, fail_fast=False, headers=None, **kwds): """Adds a task to a TaskQueue.""" is_in_tx = transactional and in_transaction() transactional = transactional or is_in_tx if not name and not is_in_tx: name = uuid.uuid4().hex headers = headers or {} if fail_fast: headers['X-AppEngine-FailFast'] = 'true' if not url.startswith('/'): url = '/' + url task = taskqueue.Task(url=url, name=name, headers=headers, **kwds) try: task.add(queue_name, transactional=transactional) except (taskqueue.TaskAlreadyExistsError, taskqueue.TombstonedTaskError): if not ignore_already: raise return task
def _get_pending_auth_db_transaction(): """Used internally to keep track of changes done in the transaction. Returns: Instance of _AuthDBTransaction (stored in the transaction context). """ # Use transaction context to store the object. Note that each transaction # retry gets its own new transaction context which is what we need, # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6). assert ndb.in_transaction() ctx = ndb.get_context() txn = getattr(ctx, '_auth_db_transaction', None) if txn: return txn # Prepare next AuthReplicationState (auth_db_rev +1). state = replication_state_key().get() if not state: primary_id = app_identity.get_application_id() if is_primary() else None state = AuthReplicationState( key=replication_state_key(), primary_id=primary_id, auth_db_rev=0) # Assert Primary or Standalone. Replicas can't increment auth db revision. if not is_primary() and state.primary_id: raise ValueError('Can\'t modify Auth DB on Replica') state.auth_db_rev += 1 state.modified_ts = utils.utcnow() # Store the state in the transaction context. Used in replicate_auth_db(...) # later. txn = _AuthDBTransaction(state) ctx._auth_db_transaction = txn return txn
def replicate_auth_db(): """Increments auth_db_rev, updates historical log, triggers replication. Must be called once from inside a transaction (right before exiting it). Should only be called for services in Standalone or Primary modes. Will raise ValueError if called on Replica. When called for service in Standalone mode, will update auth_db_rev but won't kick any replication. For services in Primary mode will also initiate replication by calling callback set in 'configure_as_primary'. The callback usually transactionally enqueues a task (to gracefully handle transaction rollbacks). WARNING: This function relies on a valid transaction context. NDB hooks and asynchronous operations are known to be buggy in this regard: NDB hook for an async operation in a transaction may be called with a wrong context (main event loop context instead of transaction context). One way to work around that is to monkey patch NDB (as done here: https://goo.gl/1yASjL). Another is to not use hooks at all. There's no way to differentiate between sync and async modes of an NDB operation from inside a hook. And without a strict assert it's very easy to forget about "Do not use put_async" warning. For that reason _post_put_hook is NOT used and replicate_auth_db() should be called explicitly whenever relevant part of root_key() entity group is updated. Returns: New AuthDB revision number. """ assert ndb.in_transaction() txn = _get_pending_auth_db_transaction() txn.commit() if is_primary(): _replication_callback(txn.replication_state) return txn.replication_state.auth_db_rev
def generate_id(cls): """ Generate a unique id. """ if ndb.in_transaction(): tiny_id = tinyid.TinyIDGenerator(namespace='P').generate_tinyid(run_in_transaction=False).upper() else: tiny_id = tinyid.TinyIDGenerator(namespace='P').generate_tinyid(run_in_transaction=True).upper() return '%s-%s' % ('P', tiny_id)
def _update_oauth_config(rev, conf): assert ndb.in_transaction(), 'Must be called in AuthDB transaction' existing = model.root_key().get() existing_as_dict = { 'oauth_client_id': existing.oauth_client_id, 'oauth_client_secret': existing.oauth_client_secret, 'oauth_additional_client_ids': list(existing.oauth_additional_client_ids), 'token_server_url': existing.token_server_url, } new_as_dict = { 'oauth_client_id': conf.primary_client_id, 'oauth_client_secret': conf.primary_client_secret, 'oauth_additional_client_ids': list(conf.client_ids), 'token_server_url': conf.token_server_url, } if new_as_dict == existing_as_dict: return False existing.populate(**new_as_dict) existing.record_revision(modified_by=model.get_service_self_identity(), modified_ts=utils.utcnow(), comment='Importing oauth.cfg at rev %s' % rev.revision) existing.put() return True
def _maybe_pubsub_notify_via_tq(result_summary, request): """Examines result_summary and enqueues a task to send PubSub message. Must be called within a transaction. Raises CommitError on errors (to abort the transaction). """ assert ndb.in_transaction() assert isinstance(result_summary, task_result.TaskResultSummary), result_summary assert isinstance(request, task_request.TaskRequest), request if (result_summary.state in task_result.State.STATES_NOT_RUNNING and request.pubsub_topic): task_id = task_pack.pack_result_summary_key(result_summary.key) ok = utils.enqueue_task(url='/internal/taskqueue/pubsub/%s' % task_id, queue_name='pubsub', transactional=True, payload=utils.encode_to_json({ 'task_id': task_id, 'topic': request.pubsub_topic, 'auth_token': request.pubsub_auth_token, 'userdata': request.pubsub_userdata, })) if not ok: raise datastore_utils.CommitError( 'Failed to enqueue task queue task')
def _maybe_pubsub_notify_now(result_summary, request): """Examines result_summary and sends task completion PubSub message. Does it only if result_summary indicates a task in some finished state and the request is specifying pubsub topic. Returns False to trigger the retry (on transient errors), or True if retry is not needed (e.g. messages was sent successfully or fatal error happened). """ assert not ndb.in_transaction() assert isinstance(result_summary, task_result.TaskResultSummary), result_summary assert isinstance(request, task_request.TaskRequest), request if (result_summary.state in task_result.State.STATES_NOT_RUNNING and request.pubsub_topic): task_id = task_pack.pack_result_summary_key(result_summary.key) try: _pubsub_notify(task_id, request.pubsub_topic, request.pubsub_auth_token, request.pubsub_userdata) except pubsub.TransientError: logging.exception( 'Transient error when sending PubSub notification') return False except pubsub.Error: logging.exception('Fatal error when sending PubSub notification') return True # do not retry it return True
def cancel_incomplete_steps_async(cls, build_id, end_ts): """Marks incomplete steps as canceled in the Datastore, if any.""" assert end_ts.seconds assert ndb.in_transaction() entity = yield cls.key_for(ndb.Key(Build, build_id)).get_async() if not entity: return container = build_pb2.Build() entity.read_steps(container) changed = False for s in container.steps: if not is_terminal_status(s.status): s.status = common_pb2.CANCELED s.end_time.CopyFrom(end_ts) if s.summary_markdown: # pragma: no branch s.summary_markdown += '\n' s.summary_markdown += ( 'step was canceled because it did not end before build ended' ) changed = True if changed: # pragma: no branch entity.write_steps(container) yield entity.put_async()
def _maybe_pubsub_notify_via_tq(result_summary, request): """Examines result_summary and enqueues a task to send PubSub message. Must be called within a transaction. Raises CommitError on errors (to abort the transaction). """ assert ndb.in_transaction() assert isinstance(result_summary, task_result.TaskResultSummary), result_summary assert isinstance(request, task_request.TaskRequest), request if result_summary.state in task_result.State.STATES_NOT_RUNNING and request.pubsub_topic: task_id = task_pack.pack_result_summary_key(result_summary.key) ok = utils.enqueue_task( url="/internal/taskqueue/pubsub/%s" % task_id, queue_name="pubsub", transactional=True, payload=utils.encode_to_json( { "task_id": task_id, "topic": request.pubsub_topic, "auth_token": request.pubsub_auth_token, "userdata": request.pubsub_userdata, } ), ) if not ok: raise datastore_utils.CommitError("Failed to enqueue task queue task")
def _CheckVotingAllowed(self): """Check whether the voting on the blockable is permitted. **NOTE** This method is a noop outside of a transaction (i.e. outside of _TransactionalVoting) EXCEPT for SantaBundles. This behavior is intended to accommodate the SantaBundle._HasFlagged* checks which can touch more than 25 entities. For SantaBundles, IsVotingAllowed is run once in its entirety prior to voting and again in each attempt of _TransactionalVoting method without the _HasFlagged* checks. Raises: OperationNotAllowed: The user may not vote on the blockable due to one of the VOTING_PROHIBITED_REASONS. """ if isinstance(self.blockable, santa.SantaBundle): allowed, reason = self.blockable.IsVotingAllowed( current_user=self.user, enable_flagged_checks=not ndb.in_transaction()) if not allowed: message = 'Voting on this Blockable is not allowed (%s)' % reason logging.warning(message) raise OperationNotAllowed(message) else: super(SantaBallotBox, self)._CheckVotingAllowed()
def _post_put_hook(self, future): from plugins.tff_backend.dal.node_orders import index_node_order if ndb.in_transaction(): from google.appengine.ext import deferred deferred.defer(index_node_order, self, _transactional=True) else: index_node_order(self)
def set_from_run_result(self, run_result, request): """Copies all the relevant properties from a TaskRunResult into this TaskResultSummary. If the task completed, succeeded and is idempotent, self.properties_hash is set. """ assert ndb.in_transaction() assert isinstance(request, task_request.TaskRequest), request assert isinstance(run_result, TaskRunResult), run_result for property_name in _TaskResultCommon._properties_fixed(): setattr(self, property_name, getattr(run_result, property_name)) # Include explicit support for 'state' and 'try_number'. TaskRunResult.state # is a ComputedProperty so it can't be copied as-is, and try_number is a # generated property. # pylint: disable=W0201 self.state = run_result.state self.try_number = run_result.try_number while len(self.costs_usd) < run_result.try_number: self.costs_usd.append(0.) self.costs_usd[run_result.try_number - 1] = run_result.cost_usd # Update the automatic tags, removing the ones from the other # TaskProperties. t = request.task_slice(run_result.current_task_slice or 0) if run_result.current_task_slice != self.current_task_slice: self.tags = task_request.get_automatic_tags( request, run_result.current_task_slice) if (self.state == State.COMPLETED and not self.failure and not self.internal_failure and t.properties.idempotent and not self.deduped_from): # Signal the results are valid and can be reused. If the request has a # SecretBytes, it is GET, which is a performance concern. self.properties_hash = t.properties_hash()
def DeferCreate(cls, **row_params): if not settings.ENV.ENABLE_BIGQUERY_STREAMING: return deferred.defer(cls.Create, _queue=constants.TASK_QUEUE.BQ_PERSISTENCE, _transactional=ndb.in_transaction(), **row_params)
def has_capacity(dimensions): """Returns True if there's a reasonable chance for this task request dimensions set to be serviced by a bot alive. First look at the task queues, then look into the datastore to figure this out. """ assert not ndb.in_transaction() # Look at the fast path. cap = task_queues.probably_has_capacity(dimensions) if cap is not None: return cap # Do a query. That's slower and it's eventually consistent. q = BotInfo.query() flat = task_queues.dimensions_to_flat(dimensions) for f in flat: q = q.filter(BotInfo.dimensions_flat == f) if q.count(limit=1): logging.info('Found capacity via BotInfo: %s', flat) # Add it to the quick cache to improve performance. task_queues.set_has_capacity(dimensions) return True logging.error('HAS NO CAPACITY: %s', flat) # TODO(maruel): https://crbug.com/839173 return _FAKE_CAPACITY
def _get_pending_auth_db_transaction(): """Used internally to keep track of changes done in the transaction. Returns: Instance of _AuthDBTransaction (stored in the transaction context). """ # Use transaction context to store the object. Note that each transaction # retry gets its own new transaction context which is what we need, # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6). assert ndb.in_transaction() ctx = ndb.get_context() txn = getattr(ctx, "_auth_db_transaction", None) if txn: return txn # Prepare next AuthReplicationState (auth_db_rev +1). state = replication_state_key().get() if not state: primary_id = app_identity.get_application_id() if is_primary() else None state = AuthReplicationState(key=replication_state_key(), primary_id=primary_id, auth_db_rev=0) # Assert Primary or Standalone. Replicas can't increment auth db revision. if not is_primary() and state.primary_id: raise ValueError("Can't modify Auth DB on Replica") state.auth_db_rev += 1 state.modified_ts = utils.utcnow() # Store the state in the transaction context. Used in replicate_auth_db(...) # later. txn = _AuthDBTransaction(state) ctx._auth_db_transaction = txn return txn
def send_firebase_update(user_id, data): if not is_firebase_enabled(): logging.warn('Not sending channel update, firebase is not configured properly') return if ndb.in_transaction(): on_trans_committed(_send_firebase_update, user_id, data) else: try_or_defer(_send_firebase_update, user_id, data)
def replicate_auth_db(): """Increments auth_db_rev by one. It is a signal that Auth DB should be replicated to Replicas. If called from inside a transaction, it inherits it and updates auth_db_rev only once (even if called multiple times during that transaction). Should only be called for services in Standalone or Primary modes. Will raise ValueError if called on Replica. When called for service in Standalone mode, will update auth_db_rev but won't kick any replication. For services in Primary mode will also initiate replication by calling callback set in 'configure_as_primary'. WARNING: This function relies on a valid transaction context. NDB hooks and asynchronous operations are known to be buggy in this regard: NDB hook for an async operation in a transaction may be called with a wrong context (main event loop context instead of transaction context). One way to work around that is to monkey patch NDB (as done here: https://goo.gl/1yASjL). Another is to not use hooks at all. There's no way to differentiate between sync and async modes of an NDB operation from inside a hook. And without a strict assert it's very easy to forget about "Do not use put_async" warning. For that reason _post_put_hook is NOT used and replicate_auth_db() should be called explicitly whenever relevant part of root_key() entity group is updated. """ def increment_revision_and_update_replicas(): """Does the actual job, called inside a transaction.""" # Update auth_db_rev. replication_state_key() is in same group as root_key. state = replication_state_key().get() if not state: primary_id = app_identity.get_application_id() if is_primary() else None state = AuthReplicationState( key=replication_state_key(), primary_id=primary_id, auth_db_rev=0) # Assert Primary or Standalone. Replicas can't increment auth db revision. if not is_primary() and state.primary_id: raise ValueError('Can\'t modify Auth DB on Replica') state.auth_db_rev += 1 state.modified_ts = utils.utcnow() state.put() # Only Primary does active replication. if is_primary(): _replication_callback(state) # If not in a transaction, start a new one. if not ndb.in_transaction(): ndb.transaction(increment_revision_and_update_replicas) return # If in a transaction, use transaction context to store "already did this" # flag. Note that each transaction retry gets its own new transaction context, # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6). ctx = ndb.get_context() if not getattr(ctx, '_auth_db_inc_called', False): increment_revision_and_update_replicas() ctx._auth_db_inc_called = True
def _trigger_invalidate_cache(self): def invalidate_cache(): self.invalidateCache() logging.info('%s: Cache invalidated', self.__class__.__name__) if ndb.in_transaction() and self.on_trans_committed: self.on_trans_committed(invalidate_cache) else: invalidate_cache()
def _post_put_hook(self, future): from plugins.tff_backend.dal.investment_agreements import index_investment_agreement if ndb.in_transaction(): from google.appengine.ext import deferred deferred.defer(index_investment_agreement, self, _transactional=True) else: index_investment_agreement(self)
def _get_query_BotTaskDimensions_keys(task_dimensions_flat): """Returns a BotTaskDimensions ndb.Key ndb.QueryIterator for the bots that corresponds to these task request dimensions. """ assert not ndb.in_transaction() q = BotDimensions.query() for d in task_dimensions_flat: q = q.filter(BotDimensions.dimensions_flat == d) return q.iter(batch_size=100, keys_only=True, deadline=15)
def has_capacity(dimensions): """Returns True if there's a reasonable chance for this task request dimensions set to be serviced by a bot alive. First look at the task queues, then look into the datastore to figure this out. """ assert not ndb.in_transaction() # Look at the fast path. cap = task_queues.probably_has_capacity(dimensions) if cap is not None: return cap # Add it to the 'quick cache' to improve performance. This cache is kept for # the same duration as how long bots are considered still alive without a # ping. Useful if there's a single bot in the fleet for these dimensions and # it takes a long time to reboot. This is the case with Android with slow # initialization and some baremetal bots (thanks SCSI firmware!). seconds = config.settings().bot_death_timeout_secs @ndb.tasklet def run_query(flat): # Do a query. That's slower and it's eventually consistent. q = BotInfo.query() for f in flat: q = q.filter(BotInfo.dimensions_flat == f) num = yield q.count_async(limit=1) if num: logging.info('Found capacity via BotInfo: %s', flat) raise ndb.Return(True) # Search a bit harder. In this case, we're looking for BotEvent which would # be a bot that used to exist recently. cutoff = utils.utcnow() - datetime.timedelta(seconds=seconds) q = BotEvent.query(BotEvent.ts > cutoff) for f in flat: q = q.filter(BotEvent.dimensions_flat == f) num = yield q.count_async(limit=1) if num: logging.info('Found capacity via BotEvent: %s', flat) raise ndb.Return(True) raise ndb.Return(False) futures = [ run_query(f) for f in task_queues.expand_dimensions_to_flats(dimensions) ] ndb.tasklets.Future.wait_all(futures) if any(f.get_result() for f in futures): task_queues.set_has_capacity(dimensions, seconds) return True logging.warning('HAS NO CAPACITY: %s', dimensions) return False
def send_collection_message(sik, service_identity, email, app_id, message): json_rpc_id = guid() deferred.defer(_send_collection_message, sik, service_identity, email, app_id, message, json_rpc_id, _transactional=ndb.in_transaction())
def update_user_data(sik, service_identity, email, app_id, address, notifications, collections): deferred.defer(_update_user_data, sik, service_identity, email, app_id, address, notifications, collections, _transactional=ndb.in_transaction())
def on_build_completed(build): # pragma: no cover assert not ndb.in_transaction() logging.info( 'Build %s was completed by %s. Status: %s. Result: %s', build.key.id(), auth.get_current_identity().to_bytes(), build.status, build.result) metrics.inc_completed_builds(build) metrics.add_build_cycle_duration(build) if build.start_time: metrics.add_build_run_duration(build)
def send_message_and_email(app_user, message, subject): human_user, app_id = get_app_user_tuple(app_user) member = MemberTO(member=human_user.email(), app_id=app_id, alert_flags=0) deferred.defer(send_rogerthat_message, member, message, _transactional=ndb.in_transaction()) if not DEBUG: iyo_username = get_iyo_username(app_user) message += '\n\nKind regards,\nThe ThreeFold Team' if iyo_username is None: logging.error( 'Could not find itsyou.online username for app_user %s, not sending intercom email' '\nSubject: %s\nMessage:%s', app_user, subject, message) else: deferred.defer(send_intercom_email, iyo_username, subject, message, _transactional=ndb.in_transaction())
def on_build_completed(build): # pragma: no cover assert not ndb.in_transaction() logging.info( 'Build %s was completed by %s. Status: %s', build.key.id(), auth.get_current_identity().to_bytes(), common_pb2.Status.Name(build.proto.status), ) metrics.inc_completed_builds(build) metrics.add_build_cycle_duration(build) if build.proto.HasField('start_time'): metrics.add_build_run_duration(build)
def DeferLookupMetric(blockable_id, reason, queue=constants.TASK_QUEUE.METRICS): """Defer a task to collect binary health analysis.""" if not settings.ENABLE_BINARY_ANALYSIS_PRECACHING: return deferred.defer(CollectLookup, blockable_id, reason, _queue=queue, _transactional=ndb.in_transaction())
def enqueue_bq_export_async(build): """Enqueues a pull task to export a completed build to BigQuery.""" assert ndb.in_transaction() assert build assert build.is_ended task_def = { 'method': 'PULL', 'payload': { 'id': build.key.id() }, } return tq.enqueue_async('bq-export', [task_def])
def WrappedFunc(*args, **kwargs): """Function wrapper to group transactional defers into single tasks.""" ctx_id = id(ndb.get_context()) _DELAYED_TASKS.setdefault(ctx_id, collections.defaultdict(list)) try: result = func(*args, **kwargs) stop_iter = None if isinstance(result, types.GeneratorType): # Emulate the NDB event loop. try: while True: response = result.send((yield result.next())) while isinstance(response, ndb.Future): response = result.send((yield response)) except StopIteration as e: stop_iter = e else: ret = ndb.Future() ret.set_result(result) yield ret # Flush all pending async tasks. # NOTE: This isn't strictly speaking necessary but protects # against cases where the user forgets to yield or manually wait on all # async operations. These would normally get executed when the transaction # concludes but that will be after this defer decorator ends. ndb.get_context().flush().wait() task_maps = _DELAYED_TASKS[ctx_id] assert len(task_maps) <= 5, ( '%s taskqueues used in a transaction (max=5)' % len(task_maps)) assert ndb.in_transaction(), ( 'Transactional defer grouping is permitted outside a transaction. ' '(Ensure that this decorator is below/inside any ' 'transactional decorators.)') for queue, task_list in task_maps.iteritems(): deferred.defer(_IndirectTaskEnqueue, queue, *task_list, _queue=queue, _transactional=True, _headers={_COMMIT_KEY: True}) if stop_iter is not None: raise stop_iter # pylint: disable=raising-bad-type finally: _DELAYED_TASKS[ctx_id].clear()
def _delete_instance(instance_key, instance_group_manager): """Attempts to delete the given Instance. Args: instance_key: ndb.Key for a models.Instance entity. instance_group_manager: models.InstanceGroupManager. """ assert ndb.in_transaction() for i, key in enumerate(instance_group_manager.instances): if key.id() == instance_key.id(): instance_group_manager.instances.pop(i) instance_group_manager.put() key.delete() return
def enqueue_callback_task_if_needed(build): assert ndb.in_transaction() assert build if build.pubsub_callback: # pragma: no branch deferred.defer( _publish_pubsub_message, build.key.id(), build.pubsub_callback.topic, build.pubsub_callback.user_data, build.pubsub_callback.auth_token, _transactional=True, _retry_options=taskqueue.TaskRetryOptions( task_age_limit=model.BUILD_TIMEOUT.total_seconds()), )
def on_trans_committed(func, *args, **kwargs): """ Executes func when the transaction the function is run in has completed. Args: func: Function to execute *args: Positional arguments for func **kwargs: Keyword arguments for func Notes: Does not return the function's return value. """ azzert(ndb.in_transaction()) post_transaction_actions.append(True, func, *args, **kwargs)
def _batch_operations(self, func): """ Setups a memory cache and a transaction to execute the operation specified in the function |func|. Performing multiple operations on the tree in |func| allows for optimal caching of nodes. _batch_operations() can also be nested, with the nested calls having no effect on the cache. If |func| is finished, all changes to the tree will be flushed to the datastore. All operations must be part of a call to _batch_operations, as it sets up caches that are used in most calls. """ def txn(): first_batch_call = not all([ hasattr(self, "_nodes_to_put"), hasattr(self, "_indices_to_put"), hasattr(self, "_identifier_cache"), hasattr(self, "_keys_to_delete") ]) if first_batch_call: self._nodes_to_put = dict() self._indices_to_put = dict() self._identifier_cache = dict() self._keys_to_delete = set() try: results = func() if first_batch_call and any([ self._nodes_to_put, self._indices_to_put, self._keys_to_delete ]): futures = ndb.delete_multi_async(self._keys_to_delete) ndb.put_multi( chain(self._nodes_to_put.itervalues(), self._indices_to_put.itervalues())) [future.get_result() for future in futures] finally: if first_batch_call: del self._nodes_to_put del self._indices_to_put del self._identifier_cache del self._keys_to_delete return results if ndb.in_transaction(): return txn() else: return ndb.transaction(txn)
def _update_security_config(root, _rev, conf): assert ndb.in_transaction(), 'Must be called in AuthDB transaction' assert isinstance(root, model.AuthGlobalConfig), root # Any changes? Compare semantically, not as byte blobs, since it is not # guaranteed that the byte blob serialization is stable. existing = security_config_pb2.SecurityConfig() if root.security_config: existing.MergeFromString(root.security_config) if existing == conf: return False # Note: this byte blob will be pushed to all service as is. root.security_config = conf.SerializeToString() return True
def get_versioned_most_recent_with_root(cls, root_key): """Returns the most recent instance of a versioned entity and the root entity. Getting the root entity is needed to get the current index. """ # Using a cls.query(ancestor=root_key).get() would work too but is less # efficient since it can't be cached by ndb's cache. assert not ndb.in_transaction() assert issubclass(cls, ndb.Model), cls assert root_key is None or isinstance(root_key, ndb.Key), root_key root = root_key.get() if not root or not root.current: return None, None return root, ndb.Key(cls, root.current, parent=root_key).get()
def save_utility_bill(url, profile_key): from plugins.tff_backend.bizz.user import store_kyc_in_user_data result = urlfetch.fetch(url) # type: urlfetch._URLFetchResult if result.status_code != 200: raise Exception('Invalid status %s %s' % (result.status_code, result.content)) profile = profile_key.get() # type: TffProfile content_type = result.headers.get('Content-Type', 'image/jpeg') filename = 'users/%s/utility_bill.jpeg' % profile.username profile.kyc.utility_bill_url = upload_to_gcs(filename, result.content, content_type) profile.put() deferred.defer(store_kyc_in_user_data, profile.app_user, _transactional=ndb.in_transaction())
def enqueue_notifications_async(build): assert ndb.in_transaction() assert build def mktask(mode): return dict( url='/internal/task/buildbucket/notify/%d' % build.key.id(), payload=dict(id=build.key.id(), mode=mode), retry_options=dict(task_age_limit=model.BUILD_TIMEOUT.total_seconds()), ) tasks = [mktask('global')] if build.pubsub_callback: # pragma: no branch tasks.append(mktask('callback')) return tq.enqueue_async('backend-default', tasks)
def GetVersion(cls, version=None): """Returns a version of the entity, the latest if version=None.""" assert not ndb.in_transaction() root_key = cls._GetRootKey() root = root_key.get() if not root or not root.current: return None if version is None: version = root.current elif version < 1: # Return None for versions < 1, which causes exceptions in ndb.Key() return None return ndb.Key(cls, version, parent=root_key).get()
def _enqueue_callback_task_if_needed(build): assert ndb.in_transaction() assert build if not build.callback: return task = taskqueue.Task( url=build.callback.url, headers=build.callback.headers, payload=json.dumps({ 'build_id': build.key.id(), }), ) add_kwargs = {} if build.callback.queue_name: # pragma: no branch add_kwargs['queue_name'] = build.callback.queue_name task.add(transactional=True, **add_kwargs)
def get_current_user(request): """Returns UserProfile of the current user or None. Args: request: webapp2.Request object with the current request. """ assert not ndb.in_transaction(), "Do not call get_current_user() in a transaction" session = get_open_session(request.cookies.get(COOKIE_NAME)) if not session: return None return UserProfile( sub=session.key.parent().id(), email=session.email, name=session.name, picture=session.picture.encode("ascii") if session.picture else None, )
def _clearCache(cls, models): """ Makes a deferred call to clear cache. Needs to save _affected_references and dirty flag """ all_affected_references = [] for model in models: if getattr(model, 'dirty', False) and hasattr(model, '_affected_references'): all_affected_references.append(model._affected_references) if all_affected_references != []: deferred.defer( cls._clearCacheDeferred, all_affected_references, _queue='cache-clearing', _transactional=ndb.in_transaction(), _target='backend-tasks')
def _batch_operations(self, func): """ Setups a memory cache and a transaction to execute the operation specified in the function |func|. Performing multiple operations on the tree in |func| allows for optimal caching of nodes. _batch_operations() can also be nested, with the nested calls having no effect on the cache. If |func| is finished, all changes to the tree will be flushed to the datastore. All operations must be part of a call to _batch_operations, as it sets up caches that are used in most calls. """ def txn(): first_batch_call = not all([hasattr(self, "_nodes_to_put"), hasattr(self, "_indices_to_put"), hasattr(self, "_identifier_cache"), hasattr(self, "_keys_to_delete")]) if first_batch_call: self._nodes_to_put = dict() self._indices_to_put = dict() self._identifier_cache = dict() self._keys_to_delete = set() try: results = func() if first_batch_call and any([self._nodes_to_put, self._indices_to_put, self._keys_to_delete]): futures = ndb.delete_multi_async(self._keys_to_delete) ndb.put_multi(chain(self._nodes_to_put.itervalues(), self._indices_to_put.itervalues())) [future.get_result() for future in futures] finally: if first_batch_call: del self._nodes_to_put del self._indices_to_put del self._identifier_cache del self._keys_to_delete return results if ndb.in_transaction(): return txn() else: return ndb.transaction(txn)
def _pre_put_hook(self): """Use extra validation that cannot be validated throught 'validator'.""" super(_TaskResultCommon, self)._pre_put_hook() assert ndb.in_transaction(), ( 'Saving %s outside of transaction' % self.__class__.__name__) if self.state == State.EXPIRED: if self.failure or self.exit_codes: raise datastore_errors.BadValueError( 'Unexpected State, a task can\'t fail if it hasn\'t started yet') if self.state == State.TIMED_OUT and not self.failure: raise datastore_errors.BadValueError('Timeout implies task failure') if not self.modified_ts: raise datastore_errors.BadValueError('Must update .modified_ts') self.children_task_ids = sorted( set(self.children_task_ids), key=lambda x: int(x, 16))
def store_new_version_async(entity, root_cls, extra=None): """Stores a new version of the instance. entity.key is updated to the key used to store the entity. Only the parent key needs to be set. E.g. Entity(parent=ndb.Key(ParentCls, ParentId), ...) or entity.key = ndb.Key(Entry, None, ParentCls, ParentId). If there was no root entity in the DB, one is created by calling root_cls(). Fetch for root entity is not done in a transaction, so this function is unsafe w.r.t. root content. Arguments: entity: ndb.Model entity to append in the DB. root_cls: class returned by get_versioned_root_model(). extra: extraneous entities to put in the transaction. They must all be in the same entity group. Returns: tuple(root, entity) with the two entities that were PUT in the db. """ assert not ndb.in_transaction() assert isinstance(entity, ndb.Model), entity assert entity.key and entity.key.parent(), 'entity.key.parent() must be set.' # Access to a protected member _XX of a client class - pylint: disable=W0212 assert root_cls._properties.keys() == ['current'], ( 'This function is unsafe for root entity, use store_new_version_safe ' 'which is not yet implemented') root_key = entity.key.parent() root = (yield root_key.get_async()) or root_cls(key=root_key) root.current = root.current or HIGH_KEY_ID flat = list(entity.key.flat()) flat[-1] = root.current entity.key = ndb.Key(flat=flat) def _new_key_minus_one_current(): flat[-1] -= 1 root.current = flat[-1] return ndb.Key(flat=flat) extra = (extra or [])[:] extra.append(root) result = yield insert_async(entity, _new_key_minus_one_current, extra=extra) raise ndb.Return(result)
def log_service_activity(service_user, rpc_id, type_, status, function, request, response, error_code=0, error_message=None): request = _limit_request_data(request, function) if ndb.in_transaction(): rpc.rpc_items.append( NdbServiceLog(parent=ndb.Key(u'ServiceLogParent', rpc_id), user=service_user, type=type_, status=status, function=function, request=request, response=response, timestamp=int(time.time() * 1000), error_code=error_code, error_message=error_message).put_async(), _log_service_activity_deferred, service_user, rpc_id, type_, status, function, request, response, error_code, error_message) else: rpc.rpc_items.append( db.put_async( ServiceLog(parent=db.Key.from_path(u'ServiceLogParent', rpc_id), user=service_user, type=type_, status=status, function=function, request=request, response=response, timestamp=int(time.time() * 1000), error_code=error_code, error_message=error_message)), _log_service_activity_deferred, service_user, rpc_id, type_, status, function, request, response, error_code, error_message)
def _clearCache(cls, models): """ Makes a deferred call to clear cache. Needs to save _affected_references and dirty flag """ if not tba_config.CONFIG['database_query_cache'] and not tba_config.CONFIG['response_cache']: return all_affected_references = [] for model in models: if getattr(model, 'dirty', False) and hasattr(model, '_affected_references'): all_affected_references.append(model._affected_references) if all_affected_references != []: deferred.defer( cls._clearCacheDeferred, all_affected_references, _queue='cache-clearing', _transactional=ndb.in_transaction(), _target='backend-tasks')
def enqueue_process_change_task(auth_db_rev): """Transactionally adds a call to 'process_change' to the task queue. Pins the task to currently executing version of BACKEND_MODULE module (defined in config.py). Added as AuthDB commit callback in get_backend_routes() below. """ assert ndb.in_transaction() conf = config.ensure_configured() try: # Pin the task to the module and version. taskqueue.add( url="/internal/auth/taskqueue/process-change/%d" % auth_db_rev, queue_name=conf.PROCESS_CHANGE_TASK_QUEUE, headers={"Host": modules.get_hostname(module=conf.BACKEND_MODULE)}, transactional=True, ) except Exception as e: logging.error('Problem adding "process-change" task to the task queue (%s): %s', e.__class__.__name__, e) raise
def get_or_create(cls, name, minimum_degree, parent=None): """ Gets the BTree with the given |name|. If this function is called from a transaction, then the tree is directly retrieved from the Datastore. If not, then first memcache is tried, and then the datastore. This is save, as the BTree model entities are immutable and can be safely memcached. If the tree does not exist yet, then a new transaction is started to create the tree wth the provided |degree|. If a transaction is already in progress, then this transaction will be used to create the new tree. Note that this can produce errors if the entity groups do not match (and no cross-group transactions are used). Args: name: The key name of the BTree that is retrieved or otherwise inserted to the Datastore. Can be an integer or a string. minimum_degree: The degree of the tree if it is created. Must be at least 2. See comments at the top of this module for guidance on choosing the right degree. parent: An optional ndb.Key tbat is the key of the parent entity for this BTree. """ key = ndb.Key(cls, name, parent=parent) def txn(): tree = key.get() if tree is None: tree = cls.create(name, minimum_degree, parent=parent) return tree if ndb.in_transaction(): tree = txn() else: # Not in a transaction, try memcache, then datastore. tree = key.get() if tree is None: tree = ndb.transaction(txn) return tree
def _update_oauth_config(rev, conf): assert ndb.in_transaction(), 'Must be called in AuthDB transaction' existing = model.root_key().get() existing_as_dict = { 'oauth_client_id': existing.oauth_client_id, 'oauth_client_secret': existing.oauth_client_secret, 'oauth_additional_client_ids': list(existing.oauth_additional_client_ids), } new_as_dict = { 'oauth_client_id': conf.primary_client_id, 'oauth_client_secret': conf.primary_client_secret, 'oauth_additional_client_ids': list(conf.client_ids), } if new_as_dict == existing_as_dict: return False existing.populate(**new_as_dict) existing.record_revision( modified_by=model.get_service_self_identity(), modified_ts=utils.utcnow(), comment='Importing oauth.cfg at rev %s' % rev.revision) existing.put() return True
def set_lookup_cache(task_key, is_available_to_schedule): """Updates the quick lookup cache to mark an item as available or not. This cache is a blacklist of items that are already reaped, so it is not worth trying to reap it with a DB transaction. This saves on DB contention when a high number (>1000) of concurrent bots with similar dimension are reaping tasks simultaneously. In this case, there is a high likelihood that multiple concurrent HTTP handlers are trying to reap the exact same task simultaneously. This blacklist helps reduce the contention. """ # Set the expiration time for items in the negative cache as 2 minutes. This # copes with significant index inconsistency but do not clog the memcache # server with unneeded keys. cache_lifetime = 120 assert not ndb.in_transaction() key = _memcache_to_run_key(task_key) if is_available_to_schedule: # The item is now available, so remove it from memcache. memcache.delete(key, namespace='task_to_run') else: memcache.set(key, True, time=cache_lifetime, namespace='task_to_run')
def createUser(username, host_for=None): """Creates a new User entity for the specified username for the currently logged in account. Please note that there should be one-one relationship between Google Accounts and User entities. This function, however, does not check if User entity does not exist for the account. Therefore, the callers should try to make sure that this function will not create a duplicate User entity. This function will raise an error, if it is not called from within a transaction. Args: username: A string containing username. host_for: A list of program keys for which the user has a program administrator role. Returns: RichBool whose value is set to True if user has been successfully created. In that case, extra part points to the newly created user entity. Otherwise, RichBool whose value is set to False and extra part is a string that represents the reason why the action could not be completed. """ if not ndb.in_transaction(): raise RuntimeError('This function must be called from within a transaction') account = users_api.get_current_user() if not account: return rich_bool.RichBool(False, _ACCOUNT_NOT_LOGGED_IN) elif user_model.User.get_by_id(username): # there is already a user with the specified username return rich_bool.RichBool(False, _USER_EXISTS_FOR_USERNAME % username) else: host_for = host_for or [] user = user_model.User( id=username, account_id=account.user_id(), host_for=host_for) user.put() return rich_bool.RichBool(True, user)
def _maybe_pubsub_notify_now(result_summary, request): """Examines result_summary and sends task completion PubSub message. Does it only if result_summary indicates a task in some finished state and the request is specifying pubsub topic. Returns False to trigger the retry (on transient errors), or True if retry is not needed (e.g. messages was sent successfully or fatal error happened). """ assert not ndb.in_transaction() assert isinstance(result_summary, task_result.TaskResultSummary), result_summary assert isinstance(request, task_request.TaskRequest), request if result_summary.state in task_result.State.STATES_NOT_RUNNING and request.pubsub_topic: task_id = task_pack.pack_result_summary_key(result_summary.key) try: _pubsub_notify(task_id, request.pubsub_topic, request.pubsub_auth_token, request.pubsub_userdata) except pubsub.TransientError: logging.exception("Transient error when sending PubSub notification") return False except pubsub.Error: logging.exception("Fatal error when sending PubSub notification") return True # do not retry it return True
def getSpawnMailTaskTxn(context, parent=None, transactional=True): """Spawns a new Task that sends out an email with the given dictionary.""" if not (context.get('to') or context.get('bcc')): context['body'] = context.get('body', '')[:10] logging.debug("Not sending email: '%s'", context) # no-one cares :( return lambda: None # TODO(daniel): drop this when DB models are not used anymore if not parent or isinstance(parent, db.Model): mail_entity = db_email_model.Email( context=json.dumps(context), parent=parent) transactional = ndb.in_transaction() else: mail_entity = ndb_email_model.Email( parent=parent.key, context=json.dumps(context)) transactional = db.is_in_transaction() def txn(): """Transaction to ensure that a task get enqueued for each mail stored. """ mail_entity.put() if isinstance(mail_entity, db.Model): mail_entity_key = mail_entity.key() else: mail_entity_key = mail_entity.key.urlsafe() task_params = {'mail_key': str(mail_entity_key)} # Setting a countdown because the mail_entity might not be stored to # all the replicas yet. new_task = taskqueue.Task(params=task_params, url=SEND_MAIL_URL, countdown=5) new_task.add(queue_name='mail', transactional=transactional) return txn
def insert(entity, new_key_callback=None, extra=None): """Inserts an entity in the DB and guarantees creation. Similar in principle to ndb.Model.get_or_insert() except that it only succeeds when the entity was not already present. As such, this always requires a transaction. Optionally retries with a new key if |new_key_callback| is provided. Arguments: entity: entity to save, it should have its .key already set accordingly. The .key property will be mutated, even if the function fails. It is highly preferable to have a root entity so the transaction can be done safely. new_key_callback: function to generates a new key if the previous key was already taken. If this function returns None, the execution is aborted. If this parameter is None, insertion is only tried once. extra: additional entities to store simultaneously. For example a bookeeping entity that must be updated simultaneously along |entity|. All the entities must be inside the same entity group. This function is not safe w.r.t. `extra`, entities in this list will overwrite entities already in the DB. Returns: ndb.Key of the newly saved entity or None if the entity was already present in the db. """ assert not ndb.in_transaction() assert entity.key.id(), entity.key entities = [entity] if extra: entities.extend(extra) root = entity.key.pairs()[0] assert all(i.key and i.key.pairs()[0] == root for i in extra), extra if not new_key_callback: new_key_callback = lambda: None def run(): if entities[0].key.get(): # The entity exists, abort. return False ndb.put_multi(entities) return True # TODO(maruel): Run a severe load test and count the number of retries. while True: # First iterate outside the transaction in case the first entity key number # selected is already used. while entity.key and entity.key.id() and entity.key.get(): entity.key = new_key_callback() if not entity.key or not entity.key.id(): break try: if txn.transaction(run, retries=0): break except txn.CommitError: # Retry with the same key. pass else: # Entity existed. Get the next key. entity.key = new_key_callback() return entity.key
def replication_callback(auth_state): assert ndb.in_transaction() trigger_replication(auth_state.auth_db_rev, transactional=True)