github.com/munnerz/test-infra@v0.0.0-20190108210205-ce3d181dc989/gubernator/third_party/cloudstorage/rest_api.py (about)

     1  # Copyright 2012 Google Inc. All Rights Reserved.
     2  #
     3  # Licensed under the Apache License, Version 2.0 (the "License");
     4  # you may not use this file except in compliance with the License.
     5  # You may obtain a copy of the License at
     6  #
     7  #    http://www.apache.org/licenses/LICENSE-2.0
     8  #
     9  # Unless required by applicable law or agreed to in writing,
    10  # software distributed under the License is distributed on an
    11  # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
    12  # either express or implied. See the License for the specific
    13  # language governing permissions and limitations under the License.
    14  
    15  """Base and helper classes for Google RESTful APIs."""
    16  
    17  
    18  
    19  
    20  
    21  __all__ = ['add_sync_methods']
    22  
    23  import logging
    24  import os
    25  import random
    26  import time
    27  
    28  from . import api_utils
    29  
    30  try:
    31    from google.appengine.api import app_identity
    32    from google.appengine.ext import ndb
    33  except ImportError:
    34    from google.appengine.api import app_identity
    35    from google.appengine.ext import ndb
    36  
    37  
    38  
    39  def _make_sync_method(name):
    40    """Helper to synthesize a synchronous method from an async method name.
    41  
    42    Used by the @add_sync_methods class decorator below.
    43  
    44    Args:
    45      name: The name of the synchronous method.
    46  
    47    Returns:
    48      A method (with first argument 'self') that retrieves and calls
    49      self.<name>, passing its own arguments, expects it to return a
    50      Future, and then waits for and returns that Future's result.
    51    """
    52  
    53    def sync_wrapper(self, *args, **kwds):
    54      method = getattr(self, name)
    55      future = method(*args, **kwds)
    56      return future.get_result()
    57  
    58    return sync_wrapper
    59  
    60  
    61  def add_sync_methods(cls):
    62    """Class decorator to add synchronous methods corresponding to async methods.
    63  
    64    This modifies the class in place, adding additional methods to it.
    65    If a synchronous method of a given name already exists it is not
    66    replaced.
    67  
    68    Args:
    69      cls: A class.
    70  
    71    Returns:
    72      The same class, modified in place.
    73    """
    74    for name in cls.__dict__.keys():
    75      if name.endswith('_async'):
    76        sync_name = name[:-6]
    77        if not hasattr(cls, sync_name):
    78          setattr(cls, sync_name, _make_sync_method(name))
    79    return cls
    80  
    81  
    82  class _AE_TokenStorage_(ndb.Model):
    83    """Entity to store app_identity tokens in memcache."""
    84  
    85    token = ndb.StringProperty()
    86    expires = ndb.FloatProperty()
    87  
    88  
    89  @ndb.tasklet
    90  def _make_token_async(scopes, service_account_id):
    91    """Get a fresh authentication token.
    92  
    93    Args:
    94      scopes: A list of scopes.
    95      service_account_id: Internal-use only.
    96  
    97    Raises:
    98      An ndb.Return with a tuple (token, expiration_time) where expiration_time is
    99      seconds since the epoch.
   100    """
   101    rpc = app_identity.create_rpc()
   102    app_identity.make_get_access_token_call(rpc, scopes, service_account_id)
   103    token, expires_at = yield rpc
   104    raise ndb.Return((token, expires_at))
   105  
   106  
   107  class _RestApi(object):
   108    """Base class for REST-based API wrapper classes.
   109  
   110    This class manages authentication tokens and request retries.  All
   111    APIs are available as synchronous and async methods; synchronous
   112    methods are synthesized from async ones by the add_sync_methods()
   113    function in this module.
   114  
   115    WARNING: Do NOT directly use this api. It's an implementation detail
   116    and is subject to change at any release.
   117    """
   118  
   119    def __init__(self, scopes, service_account_id=None, token_maker=None,
   120                 retry_params=None):
   121      """Constructor.
   122  
   123      Args:
   124        scopes: A scope or a list of scopes.
   125        service_account_id: Internal use only.
   126        token_maker: An asynchronous function of the form
   127          (scopes, service_account_id) -> (token, expires).
   128        retry_params: An instance of api_utils.RetryParams. If None, the
   129          default for current thread will be used.
   130      """
   131  
   132      if isinstance(scopes, basestring):
   133        scopes = [scopes]
   134      self.scopes = scopes
   135      self.service_account_id = service_account_id
   136      self.make_token_async = token_maker or _make_token_async
   137      if not retry_params:
   138        retry_params = api_utils._get_default_retry_params()
   139      self.retry_params = retry_params
   140      self.user_agent = {'User-Agent': retry_params._user_agent}
   141      self.expiration_headroom = random.randint(60, 240)
   142  
   143    def __getstate__(self):
   144      """Store state as part of serialization/pickling."""
   145      return {'scopes': self.scopes,
   146              'id': self.service_account_id,
   147              'a_maker': (None if self.make_token_async == _make_token_async
   148                          else self.make_token_async),
   149              'retry_params': self.retry_params,
   150              'expiration_headroom': self.expiration_headroom}
   151  
   152    def __setstate__(self, state):
   153      """Restore state as part of deserialization/unpickling."""
   154      self.__init__(state['scopes'],
   155                    service_account_id=state['id'],
   156                    token_maker=state['a_maker'],
   157                    retry_params=state['retry_params'])
   158      self.expiration_headroom = state['expiration_headroom']
   159  
   160    @ndb.tasklet
   161    def do_request_async(self, url, method='GET', headers=None, payload=None,
   162                         deadline=None, callback=None):
   163      """Issue one HTTP request.
   164  
   165      It performs async retries using tasklets.
   166  
   167      Args:
   168        url: the url to fetch.
   169        method: the method in which to fetch.
   170        headers: the http headers.
   171        payload: the data to submit in the fetch.
   172        deadline: the deadline in which to make the call.
   173        callback: the call to make once completed.
   174  
   175      Yields:
   176        The async fetch of the url.
   177      """
   178      retry_wrapper = api_utils._RetryWrapper(
   179          self.retry_params,
   180          retriable_exceptions=api_utils._RETRIABLE_EXCEPTIONS,
   181          should_retry=api_utils._should_retry)
   182      resp = yield retry_wrapper.run(
   183          self.urlfetch_async,
   184          url=url,
   185          method=method,
   186          headers=headers,
   187          payload=payload,
   188          deadline=deadline,
   189          callback=callback,
   190          follow_redirects=False)
   191      raise ndb.Return((resp.status_code, resp.headers, resp.content))
   192  
   193    @ndb.tasklet
   194    def get_token_async(self, refresh=False):
   195      """Get an authentication token.
   196  
   197      The token is cached in memcache, keyed by the scopes argument.
   198      Uses a random token expiration headroom value generated in the constructor
   199      to eliminate a burst of GET_ACCESS_TOKEN API requests.
   200  
   201      Args:
   202        refresh: If True, ignore a cached token; default False.
   203  
   204      Yields:
   205        An authentication token. This token is guaranteed to be non-expired.
   206      """
   207      key = '%s,%s' % (self.service_account_id, ','.join(self.scopes))
   208      ts = yield _AE_TokenStorage_.get_by_id_async(
   209          key, use_cache=True, use_memcache=True,
   210          use_datastore=self.retry_params.save_access_token)
   211      if refresh or ts is None or ts.expires < (
   212          time.time() + self.expiration_headroom):
   213        token, expires_at = yield self.make_token_async(
   214            self.scopes, self.service_account_id)
   215        timeout = int(expires_at - time.time())
   216        ts = _AE_TokenStorage_(id=key, token=token, expires=expires_at)
   217        if timeout > 0:
   218          yield ts.put_async(memcache_timeout=timeout,
   219                             use_datastore=self.retry_params.save_access_token,
   220                             use_cache=True, use_memcache=True)
   221      raise ndb.Return(ts.token)
   222  
   223    @ndb.tasklet
   224    def urlfetch_async(self, url, method='GET', headers=None,
   225                       payload=None, deadline=None, callback=None,
   226                       follow_redirects=False):
   227      """Make an async urlfetch() call.
   228  
   229      This is an async wrapper around urlfetch(). It adds an authentication
   230      header.
   231  
   232      Args:
   233        url: the url to fetch.
   234        method: the method in which to fetch.
   235        headers: the http headers.
   236        payload: the data to submit in the fetch.
   237        deadline: the deadline in which to make the call.
   238        callback: the call to make once completed.
   239        follow_redirects: whether or not to follow redirects.
   240  
   241      Yields:
   242        This returns a Future despite not being decorated with @ndb.tasklet!
   243      """
   244      headers = {} if headers is None else dict(headers)
   245      headers.update(self.user_agent)
   246      try:
   247        self.token = yield self.get_token_async()
   248      except app_identity.InternalError, e:
   249        if os.environ.get('DATACENTER', '').endswith('sandman'):
   250          self.token = None
   251          logging.warning('Could not fetch an authentication token in sandman '
   252                       'based Appengine devel setup; proceeding without one.')
   253        else:
   254          raise e
   255      if self.token:
   256        headers['authorization'] = 'OAuth ' + self.token
   257  
   258      deadline = deadline or self.retry_params.urlfetch_timeout
   259  
   260      ctx = ndb.get_context()
   261      resp = yield ctx.urlfetch(
   262          url, payload=payload, method=method,
   263          headers=headers, follow_redirects=follow_redirects,
   264          deadline=deadline, callback=callback)
   265      raise ndb.Return(resp)
   266  
   267  
   268  _RestApi = add_sync_methods(_RestApi)