Supporting NDB in oauth2client/appengine, and beginning *slow* transition to this datastore API.
Reviewed in https://codereview.appspot.com/6852082/.
diff --git a/oauth2client/appengine.py b/oauth2client/appengine.py
index c747528..95414c9 100644
--- a/oauth2client/appengine.py
+++ b/oauth2client/appengine.py
@@ -31,6 +31,7 @@
from google.appengine.api import memcache
from google.appengine.api import users
from google.appengine.ext import db
+from google.appengine.ext import ndb
from google.appengine.ext import webapp
from google.appengine.ext.webapp.util import login_required
from google.appengine.ext.webapp.util import run_wsgi_app
@@ -76,10 +77,29 @@
"""Storage for the sites XSRF secret key.
There will only be one instance stored of this model, the one used for the
- site. """
+ site.
+ """
secret = db.StringProperty()
+class SiteXsrfSecretKeyNDB(ndb.Model):
+ """NDB Model for storage for the sites XSRF secret key.
+
+ Since this model uses the same kind as SiteXsrfSecretKey, it can be used
+ interchangeably. This simply provides an NDB model for interacting with the
+ same data the DB model interacts with.
+
+ There should only be one instance stored of this model, the one used for the
+ site.
+ """
+ secret = ndb.StringProperty()
+
+ @classmethod
+ def _get_kind(cls):
+ """Return the kind name for this class."""
+ return 'SiteXsrfSecretKey'
+
+
def _generate_new_xsrf_secret_key():
"""Returns a random XSRF secret key.
"""
@@ -165,7 +185,7 @@
class FlowProperty(db.Property):
"""App Engine datastore Property for Flow.
- Utility property that allows easy storage and retreival of an
+ Utility property that allows easy storage and retrieval of an
oauth2client.Flow"""
# Tell what the user type is.
@@ -194,6 +214,32 @@
return not value
+class FlowNDBProperty(ndb.PickleProperty):
+ """App Engine NDB datastore Property for Flow.
+
+ Serves the same purpose as the DB FlowProperty, but for NDB models. Since
+ PickleProperty inherits from BlobProperty, the underlying representation of
+ the data in the datastore will be the same as in the DB case.
+
+ Utility property that allows easy storage and retrieval of an
+ oauth2client.Flow
+ """
+
+ def _validate(self, value):
+ """Validates a value as a proper Flow object.
+
+ Args:
+ value: A value to be set on the property.
+
+ Raises:
+ TypeError if the value is not an instance of Flow.
+ """
+ logger.info('validate: Got type %s', type(value))
+ if value is not None and not isinstance(value, Flow):
+ raise TypeError('Property %s must be convertible to a flow '
+ 'instance; received: %s.' % (self._name, value))
+
+
class CredentialsProperty(db.Property):
"""App Engine datastore Property for Credentials.
@@ -240,14 +286,73 @@
return value
-class StorageByKeyName(Storage):
- """Store and retrieve a single credential to and from
- the App Engine datastore.
+# TODO(dhermes): Turn this into a JsonProperty and overhaul the Credentials
+# and subclass mechanics to use new_from_dict, to_dict,
+# from_dict, etc.
+class CredentialsNDBProperty(ndb.BlobProperty):
+ """App Engine NDB datastore Property for Credentials.
- This Storage helper presumes the Credentials
- have been stored as a CredenialsProperty
- on a datastore model class, and that entities
- are stored by key_name.
+ Serves the same purpose as the DB CredentialsProperty, but for NDB models.
+ Since CredentialsProperty stores data as a blob and this inherits from
+ BlobProperty, the data in the datastore will be the same as in the DB case.
+
+ Utility property that allows easy storage and retrieval of Credentials and
+ subclasses.
+ """
+ def _validate(self, value):
+ """Validates a value as a proper credentials object.
+
+ Args:
+ value: A value to be set on the property.
+
+ Raises:
+ TypeError if the value is not an instance of Credentials.
+ """
+ logger.info('validate: Got type %s', type(value))
+ if value is not None and not isinstance(value, Credentials):
+ raise TypeError('Property %s must be convertible to a credentials '
+ 'instance; received: %s.' % (self._name, value))
+
+ def _to_base_type(self, value):
+ """Converts our validated value to a JSON serialized string.
+
+ Args:
+ value: A value to be set in the datastore.
+
+ Returns:
+ A JSON serialized version of the credential, else '' if value is None.
+ """
+ if value is None:
+ return ''
+ else:
+ return value.to_json()
+
+ def _from_base_type(self, value):
+ """Converts our stored JSON string back to the desired type.
+
+ Args:
+ value: A value from the datastore to be converted to the desired type.
+
+ Returns:
+ A deserialized Credentials (or subclass) object, else None if the
+ value can't be parsed.
+ """
+ if not value:
+ return None
+ try:
+ # Uses the from_json method of the implied class of value
+ credentials = Credentials.new_from_json(value)
+ except ValueError:
+ credentials = None
+ return credentials
+
+
+class StorageByKeyName(Storage):
+ """Store and retrieve a credential to and from the App Engine datastore.
+
+ This Storage helper presumes the Credentials have been stored as a
+ CredentialsProperty or CredentialsNDBProperty on a datastore model class, and
+ that entities are stored by key_name.
"""
@util.positional(4)
@@ -255,16 +360,61 @@
"""Constructor for Storage.
Args:
- model: db.Model, model class
+ model: db.Model or ndb.Model, model class
key_name: string, key name for the entity that has the credentials
property_name: string, name of the property that is a CredentialsProperty
- cache: memcache, a write-through cache to put in front of the datastore
+ or CredentialsNDBProperty.
+ cache: memcache, a write-through cache to put in front of the datastore.
+ If the model you are using is an NDB model, using a cache will be
+ redundant since the model uses an instance cache and memcache for you.
"""
self._model = model
self._key_name = key_name
self._property_name = property_name
self._cache = cache
+ def _is_ndb(self):
+ """Determine whether the model of the instance is an NDB model.
+
+ Returns:
+ Boolean indicating whether or not the model is an NDB or DB model.
+ """
+ # issubclass will fail if one of the arguments is not a class, only need
+ # worry about new-style classes since ndb and db models are new-style
+ if isinstance(self._model, type):
+ if issubclass(self._model, ndb.Model):
+ return True
+ elif issubclass(self._model, db.Model):
+ return False
+
+ raise TypeError('Model class not an NDB or DB model: %s.' % (self._model,))
+
+ def _get_entity(self):
+ """Retrieve entity from datastore.
+
+ Uses a different model method for db or ndb models.
+
+ Returns:
+ Instance of the model corresponding to the current storage object
+ and stored using the key name of the storage object.
+ """
+ if self._is_ndb():
+ return self._model.get_by_id(self._key_name)
+ else:
+ return self._model.get_by_key_name(self._key_name)
+
+ def _delete_entity(self):
+ """Delete entity from datastore.
+
+ Attempts to delete using the key_name stored on the object, whether or not
+ the given key is in the datastore.
+ """
+ if self._is_ndb():
+ ndb.Key(self._model, self._key_name).delete()
+ else:
+ entity_key = db.Key.from_path(self._model.kind(), self._key_name)
+ db.delete(entity_key)
+
def locked_get(self):
"""Retrieve Credential from datastore.
@@ -276,16 +426,16 @@
if json:
return Credentials.new_from_json(json)
- credential = None
- entity = self._model.get_by_key_name(self._key_name)
+ credentials = None
+ entity = self._get_entity()
if entity is not None:
- credential = getattr(entity, self._property_name)
- if credential and hasattr(credential, 'set_store'):
- credential.set_store(self)
+ credentials = getattr(entity, self._property_name)
+ if credentials and hasattr(credentials, 'set_store'):
+ credentials.set_store(self)
if self._cache:
- self._cache.set(self._key_name, credential.to_json())
+ self._cache.set(self._key_name, credentials.to_json())
- return credential
+ return credentials
def locked_put(self, credentials):
"""Write a Credentials to the datastore.
@@ -305,9 +455,7 @@
if self._cache:
self._cache.delete(self._key_name)
- entity = self._model.get_by_key_name(self._key_name)
- if entity is not None:
- entity.delete()
+ self._delete_entity()
class CredentialsModel(db.Model):
@@ -318,6 +466,25 @@
credentials = CredentialsProperty()
+class CredentialsNDBModel(ndb.Model):
+ """NDB Model for storage of OAuth 2.0 Credentials
+
+ Since this model uses the same kind as CredentialsModel and has a property
+ which can serialize and deserialize Credentials correctly, it can be used
+ interchangeably with a CredentialsModel to access, insert and delete the same
+ entities. This simply provides an NDB model for interacting with the
+ same data the DB model interacts with.
+
+ Storage of the model is keyed by the user.user_id().
+ """
+ credentials = CredentialsNDBProperty()
+
+ @classmethod
+ def _get_kind(cls):
+ """Return the kind name for this class."""
+ return 'CredentialsModel'
+
+
def _build_state_value(request_handler, user):
"""Composes the value for the 'state' parameter.
diff --git a/tests/test_oauth2client_appengine.py b/tests/test_oauth2client_appengine.py
index 827a31f..870b5a8 100644
--- a/tests/test_oauth2client_appengine.py
+++ b/tests/test_oauth2client_appengine.py
@@ -48,6 +48,7 @@
from google.appengine.api import users
from google.appengine.api.memcache import memcache_stub
from google.appengine.ext import db
+from google.appengine.ext import ndb
from google.appengine.ext import testbed
from google.appengine.runtime import apiproxy_errors
from oauth2client import appengine
@@ -56,6 +57,8 @@
from oauth2client.clientsecrets import InvalidClientSecretsError
from oauth2client.appengine import AppAssertionCredentials
from oauth2client.appengine import CredentialsModel
+from oauth2client.appengine import CredentialsNDBModel
+from oauth2client.appengine import FlowNDBProperty
from oauth2client.appengine import FlowProperty
from oauth2client.appengine import OAuth2Decorator
from oauth2client.appengine import StorageByKeyName
@@ -198,6 +201,7 @@
'http://www.googleapis.com/scope http://www.googleapis.com/scope2',
credentials.scope)
+
class TestFlowModel(db.Model):
flow = FlowProperty()
@@ -224,6 +228,33 @@
self.assertEqual('foo_client_id', retrieved.flow.client_id)
+class TestFlowNDBModel(ndb.Model):
+ flow = FlowNDBProperty()
+
+
+class FlowNDBPropertyTest(unittest.TestCase):
+
+ def setUp(self):
+ self.testbed = testbed.Testbed()
+ self.testbed.activate()
+ self.testbed.init_datastore_v3_stub()
+ self.testbed.init_memcache_stub()
+
+ def tearDown(self):
+ self.testbed.deactivate()
+
+ def test_flow_get_put(self):
+ instance = TestFlowNDBModel(
+ flow=flow_from_clientsecrets(datafile('client_secrets.json'), 'foo',
+ redirect_uri='oob'),
+ id='foo'
+ )
+ instance.put()
+ retrieved = TestFlowNDBModel.get_by_id('foo')
+
+ self.assertEqual('foo_client_id', retrieved.flow.client_id)
+
+
def _http_request(*args, **kwargs):
resp = httplib2.Response({'status': '200'})
content = simplejson.dumps({'access_token': 'bar'})
@@ -291,6 +322,94 @@
self.assertEqual(None, credentials)
self.assertEqual(None, memcache.get('foo'))
+ def test_get_and_put_ndb(self):
+ # Start empty
+ storage = StorageByKeyName(
+ CredentialsNDBModel, 'foo', 'credentials')
+ self.assertEqual(None, storage.get())
+
+ # Refresh storage and retrieve without using storage
+ self.credentials.set_store(storage)
+ self.credentials._refresh(_http_request)
+ credmodel = CredentialsNDBModel.get_by_id('foo')
+ self.assertEqual('bar', credmodel.credentials.access_token)
+ self.assertEqual(credmodel.credentials.to_json(),
+ self.credentials.to_json())
+
+ def test_delete_ndb(self):
+ # Start empty
+ storage = StorageByKeyName(
+ CredentialsNDBModel, 'foo', 'credentials')
+ self.assertEqual(None, storage.get())
+
+ # Add credentials to model with storage, and check equivalent w/o storage
+ storage.put(self.credentials)
+ credmodel = CredentialsNDBModel.get_by_id('foo')
+ self.assertEqual(credmodel.credentials.to_json(),
+ self.credentials.to_json())
+
+ # Delete and make sure empty
+ storage.delete()
+ self.assertEqual(None, storage.get())
+
+ def test_get_and_put_mixed_ndb_storage_db_get(self):
+ # Start empty
+ storage = StorageByKeyName(
+ CredentialsNDBModel, 'foo', 'credentials')
+ self.assertEqual(None, storage.get())
+
+ # Set NDB store and refresh to add to storage
+ self.credentials.set_store(storage)
+ self.credentials._refresh(_http_request)
+
+ # Retrieve same key from DB model to confirm mixing works
+ credmodel = CredentialsModel.get_by_key_name('foo')
+ self.assertEqual('bar', credmodel.credentials.access_token)
+ self.assertEqual(self.credentials.to_json(),
+ credmodel.credentials.to_json())
+
+ def test_get_and_put_mixed_db_storage_ndb_get(self):
+ # Start empty
+ storage = StorageByKeyName(
+ CredentialsModel, 'foo', 'credentials')
+ self.assertEqual(None, storage.get())
+
+ # Set DB store and refresh to add to storage
+ self.credentials.set_store(storage)
+ self.credentials._refresh(_http_request)
+
+ # Retrieve same key from NDB model to confirm mixing works
+ credmodel = CredentialsNDBModel.get_by_id('foo')
+ self.assertEqual('bar', credmodel.credentials.access_token)
+ self.assertEqual(self.credentials.to_json(),
+ credmodel.credentials.to_json())
+
+ def test_delete_db_ndb_mixed(self):
+ # Start empty
+ storage_ndb = StorageByKeyName(
+ CredentialsNDBModel, 'foo', 'credentials')
+ storage = StorageByKeyName(
+ CredentialsModel, 'foo', 'credentials')
+
+ # First DB, then NDB
+ self.assertEqual(None, storage.get())
+ storage.put(self.credentials)
+ self.assertNotEqual(None, storage.get())
+
+ storage_ndb.delete()
+ self.assertEqual(None, storage.get())
+
+ # First NDB, then DB
+ self.assertEqual(None, storage_ndb.get())
+ storage_ndb.put(self.credentials)
+
+ storage.delete()
+ self.assertNotEqual(None, storage_ndb.get())
+ # NDB uses memcache and an instance cache (Context)
+ ndb.get_context().clear_cache()
+ memcache.flush_all()
+ self.assertEqual(None, storage_ndb.get())
+
class MockRequest(object):
url = 'https://example.org'
@@ -596,19 +715,33 @@
# Secret shouldn't change if memcache goes away.
memcache.delete(appengine.XSRF_MEMCACHE_ID,
- namespace=appengine.OAUTH2CLIENT_NAMESPACE)
+ namespace=appengine.OAUTH2CLIENT_NAMESPACE)
secret3 = appengine.xsrf_secret_key()
self.assertEqual(secret2, secret3)
# Secret should change if both memcache and the model goes away.
memcache.delete(appengine.XSRF_MEMCACHE_ID,
- namespace=appengine.OAUTH2CLIENT_NAMESPACE)
+ namespace=appengine.OAUTH2CLIENT_NAMESPACE)
model = appengine.SiteXsrfSecretKey.get_or_insert('site')
model.delete()
secret4 = appengine.xsrf_secret_key()
self.assertNotEqual(secret3, secret4)
+ def test_ndb_insert_db_get(self):
+ secret = appengine._generate_new_xsrf_secret_key()
+ appengine.SiteXsrfSecretKeyNDB(id='site', secret=secret).put()
+
+ site_key = appengine.SiteXsrfSecretKey.get_by_key_name('site')
+ self.assertEqual(site_key.secret, secret)
+
+ def test_db_insert_ndb_get(self):
+ secret = appengine._generate_new_xsrf_secret_key()
+ appengine.SiteXsrfSecretKey(key_name='site', secret=secret).put()
+
+ site_key = appengine.SiteXsrfSecretKeyNDB.get_by_id('site')
+ self.assertEqual(site_key.secret, secret)
+
class DecoratorXsrfProtectionTests(unittest.TestCase):
"""Test _build_state_value and _parse_state_value."""