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

def mock_methods(self):
    calls = []

    @ndb.tasklet
    def via_jwt(*args):
      calls.append(('jwt_based', args))
      raise ndb.Return({'access_token': 'token', 'exp_ts': 0})

    self.mock(service_account, '_mint_jwt_based_token_async', via_jwt)

    def via_gae_api(*args):
      calls.append(('gae_api', args))
      return 'token', 0
    self.mock(service_account.app_identity, 'get_access_token', via_gae_api)

    return calls 

def _get_segment(self, start, request_size):
    """Get a segment of the file from Google Storage.

    Args:
      start: start offset of the segment. Inclusive. Have to be within the
        range of the file.
      request_size: number of bytes to request. Have to be small enough
        for a single urlfetch request. May go over the logical range of the
        file.

    Yields:
      a segment [start, start + request_size) of the file.

    Raises:
      ValueError: if the file has changed while reading.
    """
    end = start + request_size - 1
    content_range = '%d-%d' % (start, end)
    headers = {'Range': 'bytes=' + content_range}
    status, resp_headers, content = yield self._api.get_object_async(
        self.name, headers=headers)
    errors.check_status(status, [200, 206], self.name, headers, resp_headers)
    self._check_etag(resp_headers.get('etag'))
    raise ndb.Return(content) 

def _get_task_request_async(task_id, request_key, viewing):
  """Returns the TaskRequest corresponding to a task ID.

  Enforces the ACL for users. Allows bots all access for the moment.

  Returns:
    TaskRequest instance.
  """
  request = yield request_key.get_async()
  if not request:
    raise endpoints.NotFoundException('%s not found.' % task_id)
  if viewing == _VIEW:
    realms.check_task_get_acl(request)
  elif viewing == _CANCEL:
    realms.check_task_cancel_acl(request)
  else:
    raise endpoints.InternalServerErrorException('_get_task_request_async()')
  raise ndb.Return(request) 

def get_diff_async(
    hostname, project, from_commit, to_commit, path, **fetch_kwargs):
  """Loads diff between two treeishes.

  Returns:
    A patch.
  """
  _validate_args(hostname, project, from_commit, path)
  _validate_treeish(to_commit)
  path = (path or '').strip('/')
  data = yield gerrit.fetch_async(
      hostname,
      '%s/+/%s..%s/%s' % _quote_all(project, from_commit, to_commit, path),
      headers={'Accept': 'text/plain'},
      **fetch_kwargs)
  raise ndb.Return(base64.b64decode(data) if data is not None else None) 

def create_invocation_async(task_run_id, realm):
  """This is wrapper for CreateInvocation API.

  Returns:
    update-token for created invocation.
  """
  hostname = app_identity.get_default_version_hostname()
  response_headers = {}

  yield _call_resultdb_recorder_api_async(
      'CreateInvocation', {
          'requestId': str(uuid.uuid4()),
          'invocationId': _get_invocation_id(task_run_id),
          'invocation': {
              'producerResource': '//%s/tasks/%s' % (hostname, task_run_id),
              'realm': realm,
          }
      },
      project_id=realm.split(':')[0],
      response_headers=response_headers)
  update_token = response_headers.get('update-token')
  assert update_token, ("response_headers should have valid update-token: %s" %
                        response_headers)
  raise ndb.Return(update_token) 

def get_tree_async(hostname, project, treeish, path=None):
  """Gets a tree object.

  Returns:
    Tree object, or None if the tree was not found.
  """
  _validate_args(hostname, project, treeish, path)
  data = yield gerrit.fetch_json_async(
      hostname, '%s/+/%s%s' % _quote_all(project, treeish, path))
  if data is None:
    raise ndb.Return(None)

  raise ndb.Return(Tree(
      id=data['id'],
      entries=[
        TreeEntry(
            id=e['id'],
            name=e['name'],
            type=e['type'],
            mode=e['mode'],
        )
        for e in data.get('entries', [])
      ])) 

def setUp(self):
    super(NetTest, self).setUp()

    @ndb.tasklet
    def get_access_token(*_args):
      raise ndb.Return(('own-token', 0))
    self.mock(auth, 'get_access_token_async', get_access_token)

    @ndb.tasklet
    def get_project_access_token(project_id, _scopes):
      raise ndb.Return(('%s-token' % project_id, 0))
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock(
        auth, 'is_internal_domain',
        lambda domain: domain == 'internal.example.com')

    self.mock(logging, 'warning', lambda *_args: None)
    self.mock(logging, 'error', lambda *_args: None) 

def delete_entry_and_gs_entry_async(key):
  """Deletes synchronously a ContentEntry and its GS file.

  It deletes the ContentEntry first, then the file in GS. The worst case is that
  the GS file is left behind and will be reaped by a lost GS task queue. The
  reverse is much worse, having a ContentEntry pointing to a deleted GS entry
  will lead to lookup failures.
  """
  bucket = config.settings().gs_bucket
  # Note that some content entries may NOT have corresponding GS files. That
  # happens for small entry stored inline in the datastore. Since this function
  # operates only on keys, it can't distinguish "large" entries stored in GS
  # from "small" ones stored inline. So instead it always tries to delete the
  # corresponding GS files, silently skipping ones that are not there.
  # Always delete ContentEntry first.
  name = key.string_id()
  yield key.delete_async()
  # This is synchronous.
  yield gcs.delete_file_async(bucket, name, ignore_missing=True)
  raise ndb.Return(None) 

def get_async(
      self, config_set, path, revision=None, dest_type=None,
      store_last_good=None):
    """Returns tuple (revision, content).

    If not found, returns (None, None).

    See api.get_async for more info.
    """
    assert config_set
    assert path

    if store_last_good:
      result = yield _get_last_good_async(config_set, path, dest_type)
      raise ndb.Return(result)

    revision, content_hash = yield self.get_config_hash_async(
        config_set, path, revision=revision)
    content = None
    if content_hash:
      content = yield self.get_config_by_hash_async(content_hash)
    config = common._convert_config(content, dest_type)
    raise ndb.Return(revision, config) 

def test_cache_with_tasklets(self):
    @utils.cache
    def f():
      ndb.sleep(0).wait()  # Yield thread.
      return 1

    @ndb.tasklet
    def g():
      yield ()  # Make g a generator.
      raise ndb.Return(f())

    def test():
      ndb.Future.wait_all([(g()), (g())])

    t = threading.Thread(target=test)
    t.daemon = True
    t.start()
    t.join(1)
    if t.is_alive():
      self.fail('deadlock') 

def test_unordered_first_concurrent_jobs(self):
    items = [10, 5, 0]

    log = []

    @ndb.tasklet
    def fn_async(x):
      log.append('%d started' % x)
      yield ndb.sleep(float(x) / 1000)
      log.append('%d finishing' % x)
      raise ndb.Return(x)

    expected = [(5, 5), (0, 0), (10, 10)]
    actual = utils.async_apply(
        items, fn_async, concurrent_jobs=2, unordered=True)
    self.assertFalse(isinstance(actual, list))
    self.assertEqual(expected, list(actual))
    self.assertEqual(log, [
        '10 started',
        '5 started',
        '5 finishing',
        '0 started',
        '0 finishing',
        '10 finishing',
    ]) 

def get_versioned_most_recent_with_root_async(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:
    raise ndb.Return(None, None)
  entity = yield ndb.Key(cls, root.current, parent=root_key).get_async()
  raise ndb.Return(root, entity) 

def _get_or_mint_token_async(
    cache_key,
    min_lifetime_sec,
    minter,
    namespace=_MEMCACHE_NS):
  """Gets an accress token from the cache or triggers mint flow."""
  # Randomize refresh time to avoid thundering herd effect when token expires.
  # Also add 5 sec extra to make sure callers will get the token that lives for
  # at least min_lifetime_sec even taking into account possible delays in
  # propagating the token up the stack. We can't give any strict guarantees
  # here though (need to be able to stop time to do that).
  token_info = yield _memcache_get(cache_key, namespace=namespace)

  min_allowed_exp = (
    utils.time_time() +
    _randint(min_lifetime_sec + 5, min_lifetime_sec + 305))

  if not token_info or token_info['exp_ts'] < min_allowed_exp:
    token_info = yield minter()
    yield _memcache_set(cache_key, token_info,
                        token_info['exp_ts'], namespace=namespace)
  raise ndb.Return(token_info['access_token'], token_info['exp_ts']) 

def get_async(self, config_set, path, dest_type=None, **kwargs):
    """Reads a (revision, config) from a file, where revision is always None.

    Kwargs are not used, but reported as warnings.
    """
    assert config_set
    assert path
    if kwargs:
      logging.warning(
          'config: parameters %r are ignored in the filesystem mode',
          kwargs.keys())
    filename = os.path.join(
        self.root,
        config_set.replace('/', os.path.sep),
        SEPARATOR,
        path.replace('/', os.path.sep))
    filename = os.path.abspath(filename)
    assert filename.startswith(os.path.abspath(self.root)), filename
    content = None
    if os.path.exists(filename):
      with open(filename, 'rb') as f:
        content = f.read()
    config = common._convert_config(content, dest_type)
    raise ndb.Return(None, config) 

def fetch_async(hostname, path, **kwargs):
  """Sends request to Gerrit, returns raw response.

  See 'net.request_async' for list of accepted kwargs.

  Returns:
    Response body on success.
    None on 404 response.

  Raises:
    net.Error on communication errors.
  """
  assert not path.startswith('/'), path
  assert 'scopes' not in kwargs, kwargs['scopes']
  try:
    url = urllib.parse.urljoin('https://' + hostname, 'a/' + path)
    result = yield net.request_async(url, scopes=[AUTH_SCOPE], **kwargs)
    raise ndb.Return(result)
  except net.NotFoundError:
    raise ndb.Return(None) 

def get_projects_async(self):
    res = yield self._api_call_async('projects', allow_not_found=False)
    raise ndb.Return(res.get('projects', [])) 

def _get_segment(self, start, request_size, check_response=True):
    """Get a segment of the file from Google Storage.

    Args:
      start: start offset of the segment. Inclusive. Have to be within the
        range of the file.
      request_size: number of bytes to request. Have to be small enough
        for a single urlfetch request. May go over the logical range of the
        file.
      check_response: True to check the validity of GCS response automatically
        before the future returns. False otherwise. See Yields section.

    Yields:
      If check_response is True, the segment [start, start + request_size)
      of the file.
      Otherwise, a tuple. The first element is the unverified file segment.
      The second element is a closure that checks response. Caller should
      first invoke the closure before consuing the file segment.

    Raises:
      ValueError: if the file has changed while reading.
    """
    end = start + request_size - 1
    content_range = '%d-%d' % (start, end)
    headers = {'Range': 'bytes=' + content_range}
    status, resp_headers, content = yield self._api.get_object_async(
        self._path, headers=headers)
    def _checker():
      errors.check_status(status, [200, 206], self._path, headers,
                          resp_headers, body=content)
      self._check_etag(resp_headers.get('etag'))
    if check_response:
      _checker()
      raise ndb.Return(content)
    raise ndb.Return(content, _checker) 

def _get_config_provider_async():  # pragma: no cover
  """Returns a config provider to load configs.

  There are two config provider implementations: remote.Provider and
  fs.Provider. Both implement get_async(), get_project_configs_async(),
  get_ref_configs_async() with same signatures.
  """
  raise ndb.Return((yield remote.get_provider_async()) or fs.get_provider()) 

def get_config_set_location_async(self, _config_set):
    """Returns URL of where configs for given config set are stored.

    Returns:
      Always None for file system.
    """
    raise ndb.Return(None) 

def tell(self):
    """Return the total number of bytes passed to write() so far.

    (There is no seek() method.)
    """
    return self._offset 

def get_config_set_location_async(self, config_set):
    """Returns URL of where configs for given config set are stored.

    Returns:
      URL or None if no such config set.
    """
    assert config_set
    res = yield self._api_call_async(
        'mapping', params={'config_set': config_set})
    if not res:
      raise ndb.Return(None)
    for entry in res.get('mappings', []):
      if entry.get('config_set') == config_set:
        raise ndb.Return(entry.get('location'))
    raise ndb.Return(None) 

def get_ref_configs_async(self, path):
    """Reads a config file in all refs of all projects.

    Returns:
      {config_set -> (revision, content)} map, where revision is always None.
    """
    assert path
    result = {}
    for pid in self.get_project_ids():
      for ref in self.get_project_refs(pid):
        config_set = 'projects/%s/%s' % (pid, ref)
        rev, content = yield self.get_async(config_set, path)
        if content is not None:
          result[config_set] = (rev, content)
    raise ndb.Return(result) 

def get_project_configs_async(self, path):
    """Reads a config file in all projects.

    Returns:
      {config_set -> (revision, content)} map, where revision is always None.
    """
    assert path
    config_sets = ['projects/%s' % pid for pid in self.get_project_ids()]
    result = {}
    for config_set in config_sets:
      rev, content = yield self.get_async(config_set, path)
      if content is not None:
        result[config_set] = (rev, content)
    raise ndb.Return(result) 

def get_projects_async(self):
    projects = [{'id': pid} for pid in sorted(self.get_project_ids())]
    # TODO(nodir): read project names from projects/<pid>:project.cfg
    raise ndb.Return(projects) 

def _lookup_cache_is_taken_async(to_run_key):
  """Queries the quick lookup cache to reduce DB operations."""
  key = _memcache_to_run_key(to_run_key)
  neg = yield ndb.get_context().memcache_get(key, namespace='task_to_run')
  raise ndb.Return(bool(neg)) 

def setUp(self):
    super(GerritFetchTestCase, self).setUp()
    self.response = mock.Mock(
        status_code=httplib.OK,
        headers={},
        content='',
    )
    self.urlfetch_mock = mock.Mock(return_value=self.response)
    @ndb.tasklet
    def mocked_urlfetch(**kwargs):
      raise ndb.Return(self.urlfetch_mock(**kwargs))
    self.mock(net, 'urlfetch_async', mocked_urlfetch)
    self.mock(auth, 'get_access_token', mock.Mock(return_value=('token', 0.0)))
    self.mock(logging, 'warning', mock.Mock()) 

def project_credentials(project_id):
  """Return credentials that use project-based authorization.

  The returned value can be used as "credentials" argument in RPC method calls.
  """
  return lambda req: req._replace(project_id=project_id) 

def mock_urlfetch(self, calls):
    calls = calls[:]

    @ndb.tasklet
    def urlfetch_mock(
        url, payload, method, headers,
        follow_redirects, deadline, validate_certificate):
      self.assertFalse(follow_redirects)
      self.assertEqual(deadline, 5)
      self.assertTrue(validate_certificate)
      if not calls:
        self.fail('Unexpected call to %s' % url)
      call = calls.pop(0).copy()
      response = call.pop('response')
      self.assertEqual({
        'url': url,
        'payload': payload,
        'method': method,
        'headers': headers,
      }, call)
      if isinstance(response, Exception):
        raise response
      if isinstance(response, MockedResponse):
        raise ndb.Return(response)
      raise ndb.Return(MockedResponse(response[0], json.dumps(response[1])))

    self.mock(service_account, '_urlfetch', urlfetch_mock)
    return calls 

def sign_claimset_async(self, claimset):
    self.claimsets.append(claimset)
    raise ndb.Return('fake_jwt') 

def get_projects_async():
  """Returns a list of registered projects (type Project)."""
  provider = yield _get_config_provider_async()
  project_dicts = yield provider.get_projects_async()
  empty = Project('', '', '', '')
  raise ndb.Return([empty._replace(**p) for p in project_dicts])