Optimize DB accesses by using an SQL JOIN when retrieving a user. (#56)

* Make `roles` in user model query optimization "optional".

* Add unit test for DB optimization

flask_security/datastore.py

@@ -234,12 +234,17 @@ def __init__(self, db, user_model, role_model):
 
     def get_user(self, identifier):
         from sqlalchemy import func as alchemyFn
+        user_model_query = self.user_model.query
+        if hasattr(self.user_model, 'roles'):
+            from sqlalchemy.orm import joinedload
+            user_model_query = user_model_query.options(joinedload('roles'))
+
         if self._is_numeric(identifier):
-            return self.user_model.query.get(identifier)
+            return user_model_query.get(identifier)
         for attr in get_identity_attributes():
             query = alchemyFn.lower(getattr(self.user_model, attr)) \
                 == alchemyFn.lower(identifier)
-            rv = self.user_model.query.filter(query).first()
+            rv = user_model_query.filter(query).first()
             if rv is not None:
                 return rv
 
@@ -251,7 +256,12 @@ def _is_numeric(self, value):
         return True
 
     def find_user(self, **kwargs):
-        return self.user_model.query.filter_by(**kwargs).first()
+        query = self.user_model.query
+        if hasattr(self.user_model, 'roles'):
+            from sqlalchemy.orm import joinedload
+            query = query.options(joinedload('roles'))
+
+        return query.filter_by(**kwargs).first()
 
     def find_role(self, role):
         return self.role_model.query.filter_by(name=role).first()

tests/test_datastore.py

@@ -7,7 +7,7 @@
 """
 
 from pytest import raises
-from utils import init_app_with_options
+from utils import init_app_with_options, get_num_queries
 
 from flask_security import RoleMixin, Security, UserMixin
 from flask_security.datastore import Datastore, UserDatastore
@@ -145,8 +145,13 @@ def test_create_user_with_roles(app, datastore):
         user = datastore.create_user(email='dude@lp.com', username='dude',
                                      password='password', roles=[role])
         datastore.commit()
+        current_nqueries = get_num_queries(datastore)
         user = datastore.find_user(email='dude@lp.com')
         assert user.has_role('admin') is True
+        end_nqueries = get_num_queries(datastore)
+        # Verify that getting user and role is just one DB query
+        assert current_nqueries is None\
+            or end_nqueries == (current_nqueries + 1)
 
 
 def test_delete_user(app, datastore):

tests/utils.py

@@ -10,6 +10,8 @@
 from flask import json
 
 from flask_security import Security
+from flask_security.datastore import SQLAlchemyUserDatastore,\
+    SQLAlchemySessionUserDatastore
 from flask_security.utils import encrypt_password
 
 _missing = object
@@ -100,3 +102,14 @@ def init_app_with_options(app, datastore, **options):
     app.config.update(**options)
     app.security = Security(app, datastore=datastore, **security_args)
     populate_data(app)
+
+
+def get_num_queries(datastore):
+    """ Return # of queries executed during test.
+    return None if datastore doesn't support this.
+    """
+    if isinstance(datastore, SQLAlchemyUserDatastore) and\
+            not isinstance(datastore, SQLAlchemySessionUserDatastore):
+        from flask_sqlalchemy import get_debug_queries
+        return len(get_debug_queries())
+    return None