Move conditionally defined ndb helpers into their own module.

Towards #433.

2 files changed, 183 insertions, 140 deletions

diff --git a/oauth2client/contrib/_appengine_ndb.py b/oauth2client/contrib/_appengine_ndb.py
new file mode 100644
index 0000000..44c0dac
--- /dev/null
+++ b/oauth2client/contrib/_appengine_ndb.py
@@ -0,0 +1,163 @@
+# Copyright 2016 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Google App Engine utilities helper.
+
+Classes that directly require App Engine's ndb library. Provided
+as a separate module in case of failure to import ndb while
+other App Engine libraries are present.
+"""
+
+import logging
+
+from google.appengine.ext import ndb
+
+from oauth2client import client
+
+
+NDB_KEY = ndb.Key
+"""Key constant used by :mod:`oauth2client.contrib.appengine`."""
+
+NDB_MODEL = ndb.Model
+"""Model constant used by :mod:`oauth2client.contrib.appengine`."""
+
+_LOGGER = logging.getLogger(__name__)
+
+
+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'
+
+
+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, client.Flow):
+            raise TypeError('Property %s must be convertible to a flow '
+                            'instance; received: %s.' % (self._name,
+                                                         value))
+
+
+class CredentialsNDBProperty(ndb.BlobProperty):
+    """App Engine NDB datastore Property for Credentials.
+
+    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, client.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 = client.Credentials.new_from_json(value)
+        except ValueError:
+            credentials = None
+        return credentials
+
+
+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'
diff --git a/oauth2client/contrib/appengine.py b/oauth2client/contrib/appengine.py
index 2f05254..a8b7163 100644
--- a/oauth2client/contrib/appengine.py
+++ b/oauth2client/contrib/appengine.py
@@ -46,12 +46,11 @@ from oauth2client.client import OAuth2WebServerFlow
 from oauth2client.client import Storage
 from oauth2client.contrib import xsrfutil
 
-# TODO(dhermes): Resolve import issue.
 # This is a temporary fix for a Google internal issue.
 try:
-    from google.appengine.ext import ndb
-except ImportError:
-    ndb = None
+    from oauth2client.contrib import _appengine_ndb
+except ImportError:  # pragma: NO COVER
+    _appengine_ndb = None
 
 
 __author__ = 'jcgregorio@google.com (Joe Gregorio)'
@@ -62,6 +61,21 @@ OAUTH2CLIENT_NAMESPACE = 'oauth2client#ns'
 
 XSRF_MEMCACHE_ID = 'xsrf_secret_key'
 
+if _appengine_ndb is None:
+    CredentialsNDBModel = None
+    CredentialsNDBProperty = None
+    FlowNDBProperty = None
+    _NDB_KEY = None
+    _NDB_MODEL = None
+    SiteXsrfSecretKeyNDB = None
+else:
+    CredentialsNDBModel = _appengine_ndb.CredentialsNDBModel
+    CredentialsNDBProperty = _appengine_ndb.CredentialsNDBProperty
+    FlowNDBProperty = _appengine_ndb.FlowNDBProperty
+    _NDB_KEY = _appengine_ndb.NDB_KEY
+    _NDB_MODEL = _appengine_ndb.NDB_MODEL
+    SiteXsrfSecretKeyNDB = _appengine_ndb.SiteXsrfSecretKeyNDB
+
 
 def _safe_html(s):
     """Escape text to make it safe to display.
@@ -91,24 +105,6 @@ class SiteXsrfSecretKey(db.Model):
     """
     secret = db.StringProperty()
 
-if ndb is not None:
-    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."""
@@ -273,35 +269,6 @@ class FlowProperty(db.Property):
         return not value
 
 
-if ndb is not None:
-    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.
 
@@ -346,73 +313,6 @@ class CredentialsProperty(db.Property):
         return value
 
 
-if ndb is not None:
-    # 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.
-
-        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.
 
@@ -460,7 +360,7 @@ class StorageByKeyName(Storage):
         # need worry about new-style classes since ndb and db models are
         # new-style
         if isinstance(self._model, type):
-            if ndb is not None and issubclass(self._model, ndb.Model):
+            if _NDB_MODEL is not None and issubclass(self._model, _NDB_MODEL):
                 return True
             elif issubclass(self._model, db.Model):
                 return False
@@ -489,7 +389,7 @@ class StorageByKeyName(Storage):
         not the given key is in the datastore.
         """
         if self._is_ndb():
-            ndb.Key(self._model, self._key_name).delete()
+            _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)
@@ -548,26 +448,6 @@ class CredentialsModel(db.Model):
     credentials = CredentialsProperty()
 
 
-if ndb is not None:
-    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.