Custom OAuth2Decorator Persistence

from google.appengine.ext import ndb

from oauth2client import appengine


class CustomCredentialsModel(appengine.CredentialsNDBModel):
  user_email = ndb.StringProperty(required=True, indexed=False)


class CustomStorageByKeyName(appengine.StorageByKeyName):
  def __init__(self, model, key_name, property_name, cache=None, user=None):
    super(CustomStorageByKeyName, self).__init__(
        model, key_name, property_name, cache=cache, user=user)
    # May want to throw exception if not a users.User object
    self._user = user

  def locked_put(self, credentials):
    entity = self._model.get_or_insert(self._key_name)
    setattr(entity, self._property_name, credentials)
    entity.user_email = self._user.email()
    entity.put()
    if self._cache:
      self._cache.set(self._key_name, credentials.to_json())


decorator = appengine.OAuth2Decorator(
    'some.client.id', 'some.client.secret', 'some.scope',
    _storage_class=CustomStorageByKeyName,
    _credentials_class=CustomCredentialsModel)