Python google.appengine.ext.ndb.get_multi() Examples

def testSpiderBite(self):
    key_1 = test_utils.CreateSantaHost(
        client_mode=MONITOR).key
    key_2 = test_utils.CreateSantaHost(
        client_mode=MONITOR).key

    host_keys = [key_1, key_2]
    hosts = [key_1.get(), key_2.get()]

    with mock.patch.object(
        role_syncing.ndb, 'get_multi', return_value=hosts) as mock_get:
      with mock.patch.object(role_syncing.ndb, 'put_multi') as mock_put:
        role_syncing._SpiderBite(host_keys)

    self.assertEqual(1, mock_put.call_count)
    self.assertEqual(1, mock_get.call_count)
    self.assertEqual([key_1, key_2], mock_get.call_args[0][0])
    self.assertEqual([key_1.get(), key_2.get()], mock_put.call_args[0][0])
    self.assertEqual(key_1.get().client_mode, LOCKDOWN)
    self.assertEqual(key_2.get().client_mode, LOCKDOWN) 

def testChangeModeForGroup_MultiBatch(self, mock_ctor):

    users = [
        test_utils.CreateUser() for _ in xrange(role_syncing.BATCH_SIZE + 1)]
    hosts = [
        test_utils.CreateSantaHost(
            primary_user=user_utils.EmailToUsername(user.key.id()),
            client_mode=MONITOR)
        for user in users]
    mock_ctor.return_value.AllMembers.return_value = [
        user.key.id() for user in users]

    response = self.testapp.get('', headers={'X-AppEngine-Cron': 'true'})

    self.assertEqual(httplib.OK, response.status_int)

    self.assertTaskCount(constants.TASK_QUEUE.DEFAULT, 2)
    self.DrainTaskQueue(constants.TASK_QUEUE.DEFAULT)

    new_hosts = ndb.get_multi(host.key for host in hosts)
    self.assertTrue(all(host.client_mode == LOCKDOWN for host in new_hosts)) 

def get_stats(handler, resolution, now, num_items, as_dict):
  """Wrapper calls that returns items for the specified resolution.

  Arguments:
  - handler: Instance of StatisticsFramework.
  - resolution: One of 'days', 'hours' or 'minutes'
  - now: datetime.datetime or None.
  - num_items: Maximum number of items to return ending at 'now'.
  - as_dict: When True, preprocess the entities to convert them to_dict(). If
        False, returns the raw objects that needs to be handled manually.
  """
  mapping = {
    'days': _get_days_keys,
    'hours': _get_hours_keys,
    'minutes': _get_minutes_keys,
  }
  keys = mapping[resolution](handler, now, num_items)
  if as_dict:
    return [
      i.get_result() for i in _get_snapshot_as_dict_future(keys)
      if i.get_result()
    ]
  # Automatically skip missing entities.
  return [i for i in ndb.get_multi(keys) if i] 

def GetExemptionsForHosts(host_keys):
  """Retrieves all Exemptions corresponding to the specified Hosts.

  Args:
    host_keys: A list of NDB Host Keys.

  Returns:
    A dictionary mapping the given Host Keys to their corresponding Exemptions,
    or None if one doesn't exist.
  """
  # Compose the expected Exemption Keys for all Hosts.
  exm_keys = [ndb.Key(flat=k.flat() + ('Exemption', '1')) for k in host_keys]

  # Grab everything from Datastore.
  exms = ndb.get_multi(exm_keys)

  # Map Host Keys to the entities we got back.
  exm_dict = {exm.key.parent(): exm for exm in exms if exm}

  # Return a mapping of all Host Keys to their Exemptions, or None for those
  # without Exemptions.
  return {host_key: exm_dict.get(host_key) for host_key in host_keys} 

def get_stats(handler, resolution, now, num_items, as_dict):
  """Wrapper calls that returns items for the specified resolution.

  Arguments:
  - handler: Instance of StatisticsFramework.
  - resolution: One of 'days', 'hours' or 'minutes'
  - now: datetime.datetime or None.
  - num_items: Maximum number of items to return ending at 'now'.
  - as_dict: When True, preprocess the entities to convert them to_dict(). If
        False, returns the raw objects that needs to be handled manually.
  """
  mapping = {
    'days': _get_days_keys,
    'hours': _get_hours_keys,
    'minutes': _get_minutes_keys,
  }
  keys = mapping[resolution](handler, now, num_items)
  if as_dict:
    return [
      i.get_result() for i in _get_snapshot_as_dict_future(keys)
      if i.get_result()
    ]
  # Automatically skip missing entities.
  return [i for i in ndb.get_multi(keys) if i] 

def AdoptProjects(dst_user_id, src_user_id):
  """Transfer project ownership to a new user."""

  @ndb.transactional(xg=True)
  def _AdoptProject(project_key, dst_user_key, src_user_key):
    prj, dst_user, src_user = ndb.get_multi([project_key, dst_user_key,
                                             src_user_key])
    if project_key not in src_user.projects:
      # another concurrent request and transaction beat us to it
      return
    src_user.projects.remove(project_key)
    dst_user.projects.append(project_key)
    prj.owner = dst_user_key.id()
    prj.writers.remove(src_user_key.id())
    prj.writers.append(dst_user_key.id())
    ndb.put_multi([prj, dst_user, src_user])

  src_user = GetUser(src_user_id)
  if not src_user or not src_user.projects:
    return
  dst_user = GetOrCreateUser(dst_user_id)
  for project_key in src_user.projects:
    # slow, but transactionable
    _AdoptProject(project_key, dst_user.key, src_user.key) 

def make_reco():
    """Makes the final recommendations for customers. Receives as input all
    items a given customer interacted with such as browsed items, added to 
    basket and purchased ones. Returns a list of top selected recommendations.
    """
    t0 = time.time()
    scores = base_utils.process_input_items(request.args)
    keys = map(lambda x: ndb.Key(config['recos']['kind'], x),
        scores.keys())
    entities = [e for e in ndb.get_multi(keys) if e]
    if not entities:
        result = {'results': [], 'statistics':
            {'elapsed_time': time.time() - t0}}
        return jsonify(result)
    results = utils.process_recommendations(entities, scores,
        int(request.args.get('n', 10))) 
    results['statistics'] = {}
    results['statistics']['elapsed_time'] = time.time() - t0 
    return jsonify(results) 

def get_by_auth_token(cls, user_id, token, subject='auth'):
      """Returns a user object based on a user ID and token.

      :param user_id:
          The user_id of the requesting user.
      :param token:
          The token string to be verified.
      :returns:
          A tuple ``(User, timestamp)``, with a user object and
          the token timestamp, or ``(None, None)`` if both were not found.
      """
      token_key = cls.token_model.get_key(user_id, subject, token)
      user_key = ndb.Key(cls, user_id)
      # Use get_multi() to save a RPC call.
      valid_token, user = ndb.get_multi([token_key, user_key])
      if valid_token and user:
          timestamp = int(time.mktime(valid_token.created.timetuple()))
          return user, timestamp

      return None, None

    # Since the original create_user method calls user.put() (where the exception occurs), only *after*
    # calling cls.unique_model.create_multi(k for k, v in uniques), this means we'll have to delete
    # those created uniques (other they'll just stay as garbage data in the DB, while not allowing
    # the user to re-register with the same username/email/etc. 

def get_stats(handler, resolution, now, num_items, as_dict):
  """Wrapper calls that returns items for the specified resolution.

  Arguments:
  - handler: Instance of StatisticsFramework.
  - resolution: One of 'days', 'hours' or 'minutes'
  - now: datetime.datetime or None.
  - num_items: Maximum number of items to return ending at 'now'.
  - as_dict: When True, preprocess the entities to convert them to_dict(). If
        False, returns the raw objects that needs to be handled manually.
  """
  mapping = {
    'days': _get_days_keys,
    'hours': _get_hours_keys,
    'minutes': _get_minutes_keys,
  }
  keys = mapping[resolution](handler, now, num_items)
  if as_dict:
    return [
      i.get_result() for i in _get_snapshot_as_dict_future(keys)
      if i.get_result()
    ]
  # Automatically skip missing entities.
  return [i for i in ndb.get_multi(keys) if i] 

def get_stats(handler, resolution, now, num_items, as_dict):
  """Wrapper calls that returns items for the specified resolution.

  Arguments:
  - handler: Instance of StatisticsFramework.
  - resolution: One of 'days', 'hours' or 'minutes'
  - now: datetime.datetime or None.
  - num_items: Maximum number of items to return ending at 'now'.
  - as_dict: When True, preprocess the entities to convert them to_dict(). If
        False, returns the raw objects that needs to be handled manually.
  """
  mapping = {
    'days': _get_days_keys,
    'hours': _get_hours_keys,
    'minutes': _get_minutes_keys,
  }
  keys = mapping[resolution](handler, now, num_items)
  if as_dict:
    return [
      i.get_result() for i in _get_snapshot_as_dict_future(keys)
      if i.get_result()
    ]
  # Automatically skip missing entities.
  return [i for i in ndb.get_multi(keys) if i] 

def GetCounterCount(name):
  """Sums a cumulative value from all the shard counts for the given name.

  Args:
    name: The name of the counter.

  Returns:
    Integer; the cumulative count of all sharded counters for the given
    counter name.
  """
  total = memcache.get(key=name)
  if total is not None:
    return total

  total = 0
  all_keys = CounterShardConfig.AllKeys(name)
  for counter in ndb.get_multi(all_keys):
    if counter is not None:
      total += counter.count
  if memcache.add(key=name, value=total,
                  time=_COUNTER_MEMCACHE_EXPIRATION_S):
    return total
  raise recall_errors.MessageRecallCounterError(
      'Unexpected problem adding to memcache: %s.' % name) 

def get_states(self, request):
    """Returns task state for a specific set of tasks.
    """
    logging.debug('%s', request)
    result_keys = [_to_keys(task_id)[1] for task_id in request.task_id]

    # Hot path. Fetch everything we can from memcache.
    entities = ndb.get_multi(
        result_keys, use_cache=True, use_memcache=True, use_datastore=False)
    states = [t.state if t else task_result.State.PENDING for t in entities]
    # Now fetch both the ones in non-stable state or not in memcache.
    missing_keys = [
      result_keys[i] for i, state in enumerate(states)
      if state in task_result.State.STATES_RUNNING
    ]
    if missing_keys:
      more = ndb.get_multi(
          missing_keys, use_cache=False, use_memcache=False, use_datastore=True)
      # This relies on missing_keys being in the same order as states (for
      # common elements).
      for i, s in enumerate(states):
        if s in task_result.State.STATES_RUNNING:
          states[i] = more.pop(0).state

    return swarming_rpcs.TaskStates(
        states=[swarming_rpcs.TaskState(state) for state in states]) 

def requests(self, request):
    """Returns tasks requests based on the filters.

    This endpoint is slightly slower than 'list'. Use 'list' or 'count' when
    possible.
    """
    logging.debug('%s', request)
    if request.include_performance_stats:
      raise endpoints.BadRequestException(
          'Can\'t set include_performance_stats for tasks/list')
    now = utils.utcnow()
    try:
      # Get the TaskResultSummary keys, then fetch the corresponding
      # TaskRequest entities.
      keys, cursor = datastore_utils.fetch_page(
          self._query_from_request(request),
          request.limit, request.cursor, keys_only=True)
      items = ndb.get_multi(
          task_pack.result_summary_key_to_request_key(k) for k in keys)
    except ValueError as e:
      raise endpoints.BadRequestException(
          'Inappropriate filter for tasks/requests: %s' % e)
    except datastore_errors.NeedIndexError as e:
      logging.error('%s', e)
      raise endpoints.BadRequestException(
          'Requires new index, ask admin to create one.')
    except datastore_errors.BadArgumentError as e:
      logging.error('%s', e)
      raise endpoints.BadRequestException(
          'This combination is unsupported, sorry.')
    return swarming_rpcs.TaskRequests(
        cursor=cursor,
        items=[message_conversion.task_request_to_rpc(i) for i in items],
        now=now) 

def get_hints(pool, oses, names):
  """Returns the hints for each named caches.

  Returns:
    list of hints in bytes for each named cache, or -1 when there's no hint
    available.
  """
  assert isinstance(oses, list), repr(oses)
  assert isinstance(names, list), repr(names)
  os = _reduce_oses(oses)
  keys = [_named_cache_key(pool, os, name) for name in names]
  entities = ndb.get_multi(keys)
  hints = [e.hint if e else -1 for e in entities]
  ancestor = ndb.Key(NamedCacheRoot, pool)
  for i, hint in enumerate(hints):
    if hint > -1:
      continue
    # Look for named cache in other OSes in the same pool.
    q = NamedCache.query(NamedCache.name == names[i], ancestor=ancestor)
    other_oses = q.fetch()
    if not other_oses:
      # TODO(maruel): We could define default hints in the pool.
      continue
    # Found something! Take the largest value.
    hints[i] = max(e.hint for e in other_oses)

  return hints 

def get_missing_groups(groups):
  """Given a list of group names, returns a list of groups that do not exist."""
  # We need to iterate over |groups| twice. It won't work if |groups|
  # is a generator. So convert to list first.
  groups = list(groups)
  entities = ndb.get_multi(group_key(name) for name in groups)
  return [name for name, ent in zip(groups, entities) if not ent] 

def get_missing_groups(groups):
  """Given a list of group names, returns a list of groups that do not exist."""
  # We need to iterate over |groups| twice. It won't work if |groups|
  # is a generator. So convert to list first.
  groups = list(groups)
  entities = ndb.get_multi(group_key(name) for name in groups)
  return [name for name, ent in zip(groups, entities) if not ent] 

def get_missing_groups(groups):
  """Given a list of group names, returns a list of groups that do not exist."""
  # We need to iterate over |groups| twice. It won't work if |groups|
  # is a generator. So convert to list first.
  groups = list(groups)
  entities = ndb.get_multi(group_key(name) for name in groups)
  return [name for name, ent in zip(groups, entities) if not ent] 

def get_missing_groups(groups):
  """Given a list of group names, returns a list of groups that do not exist."""
  # We need to iterate over |groups| twice. It won't work if |groups|
  # is a generator. So convert to list first.
  groups = list(groups)
  entities = ndb.get_multi(group_key(name) for name in groups)
  return [name for name, ent in zip(groups, entities) if not ent] 

def main():
    user = users.get_current_user()
    if not user:
        return flask.redirect(users.create_login_url(flask.request.path))

    if flask.request.method == 'POST':
        util.csrf_protect()
        tab_ids = flask.request.values.getlist('tab_id')
        keys = [ndb.Key(Session, tab_id) for tab_id in tab_ids]
        if 'delete' in flask.request.values:
            if not all(s and s.user_id == user.user_id() for s in ndb.get_multi(keys)):
                return 'Not authorized to delete some sessions', 403
            ndb.delete_multi(keys)
        elif 'share' in flask.request.values:
            for key in keys:
                session = key.get()
                if session and session.user_id == user.user_id():
                    session.shared = True
                    session.put()
        else:
            return 'Incorrect POST name', 400
    date = flask.request.values.get('date', datetime.now().strftime('%Y-%m-%d'))
    cur_day = datetime.strptime(date, '%Y-%m-%d')
    next_day = cur_day + timedelta(days=1)
    sessions = (Session.query(Session.user_id == user.user_id(),
                Session.start_ts >= cur_day, Session.start_ts < next_day)
            .order(-Session.start_ts))
    num_shared = Session.query(Session.user_id == user.user_id(), Session.shared == True).count()
    return flask.render_template('main.html',
                                 user=user,
                                 date=date,
                                 year=datetime.now().year,
                                 logout_url=users.create_logout_url('/'),
                                 sessions=sessions,
                                 num_shared=num_shared) 

def GetProjects(user):
  projects = ndb.get_multi(user.projects)
  if None in projects:
    # users.projects references projects which do not exist
    missing_projects = [key for (key, prj) in zip(user.projects, projects)
                        if prj is None]
    if common.IsDevMode():
      shared.e('Missing project(s): {0}'.format(missing_projects))
    else:
      shared.w('Missing project(s): {0}'.format(missing_projects))
      projects = [p for p in projects if p is not None]
  return projects 

def _GetAssociatedHosts(self, user):

    # Build Keys for the associated Hosts, along with their corresponding
    # Exemptions.
    host_keys = model_utils.GetHostKeysForUser(user)

    # Grab all the Hosts.
    hosts = ndb.get_multi(host_keys)
    hosts = [host for host in hosts if host is not None]

    # Get a mapping of Host Keys to Exemptions.
    exm_dict = model_utils.GetExemptionsForHosts(host_keys)

    # If Santa hosts have never synced rules or Bit9 hosts never reported an
    # event, push them to the end of the list.
    epoch = datetime.datetime.utcfromtimestamp(0)

    def ByFreshness(host):
      if isinstance(host, host_models.Bit9Host):
        return host.last_event_dt or epoch
      elif isinstance(host, host_models.SantaHost):
        return host.rule_sync_dt or epoch

    hosts = sorted(hosts, key=ByFreshness, reverse=True)

    # Convert the Host entities to dicts for the frontend, and stuff each one
    # with its corresponding Exemption (if one exists).
    host_dicts = []
    for host in hosts:
      host_dict = host.to_dict()
      if host.key in exm_dict:
        host_dict['exemption'] = exm_dict[host.key]
      host_dicts.append(host_dict)

    return host_dicts 

def GetExemptionsForUser(email_addr, state=None):
  user = user_models.User.GetById(email_addr)
  exm_keys = [
      exemption_models.Exemption.CreateKey(host_id)
      for host_id in GetHostIdsForUser(user)]
  exms = [exm for exm in ndb.get_multi(exm_keys) if exm]
  if state:
    exms = [exm for exm in exms if exm.state == state]
  return exms 

def _CommitChangeSet(change_key):
  """Attempts to commit and delete a given RuleChangeSet."""

  change = change_key.get()
  if change is None:
    logging.info('Change no longer exists. (already committed?)')
    return

  logging.info(
      'Committing a %s change set of %s rules for blockable %s',
      change.change_type, len(change.rule_keys), change.blockable_key.id())

  blockable = change.blockable_key.get()
  rules = ndb.get_multi(change.rule_keys)

  if change.change_type == constants.RULE_POLICY.WHITELIST:
    change_func = _Whitelist
  elif change.change_type == constants.RULE_POLICY.BLACKLIST:
    change_func = _Blacklist
  elif change.change_type == constants.RULE_POLICY.REMOVE:
    change_func = _Remove
  elif change.change_type in constants.RULE_POLICY.SET_INSTALLER:
    change_func = _ChangeInstallerState
  else:
    raise NotImplementedError

  # Attempt to perform the change. If something fails, just let the Exception
  # escape and kill the task. A retry will be attempted soon enough via cron.
  change_func(blockable, rules)

  # Clean up if the change went through.
  for rule in rules:
    rule.is_committed = True
  ndb.put_multi(rules)
  change.key.delete() 

def _SpiderBite(host_keys):
  hosts = ndb.get_multi(host_keys)
  for host in hosts:
    host.client_mode = _CLIENT_MODE.LOCKDOWN

  ndb.put_multi(hosts)
  logging.info(
      'Client mode changed to LOCKDOWN for %d host(s)', len(hosts)) 

def get_important_news():
    data = memcache.get(key="important_news")
    if data is None:
        q = NewsEntry.query(NewsEntry.important==True).order(-NewsEntry.published)
        data = ndb.get_multi(q.fetch(limit=app.config["PER_PAGE"],keys_only=True))
        memcache.add(key="important_news", value=data, time=app.config["INDEX_CACHE_TIME"])
    return data 

def merge_user_dbs(user_db, deprecated_keys):
  # TODO: Merge possible user data before handling deprecated users
  deprecated_dbs = ndb.get_multi(deprecated_keys)
  for deprecated_db in deprecated_dbs:
    deprecated_db.auth_ids = []
    deprecated_db.active = False
    deprecated_db.verified = False
    if not deprecated_db.username.startswith('_'):
      deprecated_db.username = '_%s' % deprecated_db.username
  ndb.put_multi(deprecated_dbs) 

def get(self):
    user_keys = util.param('user_keys', list)
    if user_keys:
      user_db_keys = [ndb.Key(urlsafe=k) for k in user_keys]
      user_dbs = ndb.get_multi(user_db_keys)
      return helpers.make_response(user_dbs, model.User.FIELDS)

    user_dbs, cursors = model.User.get_dbs(prev_cursor=True)
    return helpers.make_response(user_dbs, model.User.FIELDS, cursors) 

def get(self):
    repo_keys = util.param('repo_keys', list)
    if repo_keys:
      repo_db_keys = [ndb.Key(urlsafe=k) for k in repo_keys]
      repo_dbs = ndb.get_multi(repo_db_keys)
      return helpers.make_response(repo_dbs, model.repo.FIELDS)

    repo_dbs, repo_cursor = model.Repo.get_dbs()
    return helpers.make_response(repo_dbs, model.Repo.FIELDS, repo_cursor) 

def get(self):
    account_keys = util.param('account_keys', list)
    if account_keys:
      account_db_keys = [ndb.Key(urlsafe=k) for k in account_keys]
      account_dbs = ndb.get_multi(account_db_keys)
      return helpers.make_response(account_dbs, model.account.FIELDS)

    account_dbs, account_cursor = model.Account.get_dbs()
    return helpers.make_response(account_dbs, model.Account.FIELDS, account_cursor) 

def forProject(cls, project_name):
        """Returns the alerts created for supplied project."""
        alert_keys = []
        alert_keys = Alert.query(
            Alert.project == project_name,
            ancestor=Alert.entity_group).fetch(keys_only=True)
        alerts = ndb.get_multi(alert_keys)
        return alerts