MongoEngine datastore completed. Also refactored model classes to use a mixin

example/app.py

@@ -6,23 +6,17 @@ sys.path.pop(0)
 sys.path.insert(0, os.getcwd())
 
 from flask import Flask, render_template
-from flask.ext.sqlalchemy import SQLAlchemy
-from flask.ext.security import (Security, LoginForm, user_datastore, 
-                      login_required, roles_required, roles_accepted)
-from flask.ext.security.datastore.sqlalchemy import SQLAlchemyUserDatastore
 
-def create_app(auth_config=None):
-    
-    app = Flask(__name__)
-    app.debug = True
-    app.config['SECRET_KEY'] = 'secret'
-    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:'
-    app.config['AUTH'] = auth_config or {}
-    
-    db = SQLAlchemy(app)
-    
-    Security(app, SQLAlchemyUserDatastore(db))
-    
+from flask.ext.mongoengine import MongoEngine
+from flask.ext.sqlalchemy import SQLAlchemy
+
+from flask.ext.security import (Security, LoginForm, user_datastore, 
+                                login_required, roles_required, roles_accepted)
+
+from flask.ext.security.datastore.sqlalchemy import SQLAlchemyUserDatastore
+from flask.ext.security.datastore.mongoengine import MongoEngineUserDatastore
+
+def add_endpoints(app):
     @app.route('/')
     def index():
         return render_template('index.html', content='Home Page')
@@ -59,27 +53,66 @@ def create_app(auth_config=None):
     def admin_or_editor():
         return render_template('index.html', content='Admin or Editor Page')
     
+def create_users():
+    user_datastore.create_user(username='matt', email='matt@lp.com', 
+                               password='password',
+                               roles=['admin'])
+    
+    user_datastore.create_user(username='joe', email='joe@lp.com', 
+                               password='password',
+                               roles=['editor'])
+    
+    user_datastore.create_user(username='jill', email='jill@lp.com', 
+                               password='password',
+                               roles=['author'])
+    
+    user_datastore.create_user(username='tiya', email='tiya@lp.com', 
+                               password='password', active=False)
+
+def create_sqlalchemy_app(auth_config=None):
+    app = Flask(__name__)
+    app.debug = True
+    app.config['SECRET_KEY'] = 'secret'
+    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:'
+    app.config['AUTH'] = auth_config or {}
+    
+    db = SQLAlchemy(app)
+    
+    Security(app, SQLAlchemyUserDatastore(db))
+    
+    add_endpoints(app)
+    
     @app.before_first_request
     def before_first_request():
         db.create_all()
+        create_users()
         
-        user_datastore.create_user(username='matt', email='matt@lp.com', 
-                                   password='password',
-                                   roles=['admin'])
-        
-        user_datastore.create_user(username='joe', email='joe@lp.com', 
-                                   password='password',
-                                   roles=['editor'])
-        
-        user_datastore.create_user(username='jill', email='jill@lp.com', 
-                                   password='password',
-                                   roles=['author'])
-        
-        user_datastore.create_user(username='tiya', email='tiya@lp.com', 
-                                   password='password', active=False)
+    return app
+
+def create_mongoengine_app(auth_config=None):
+    app = Flask(__name__)
+    app.debug = True
+    app.config['SECRET_KEY'] = 'secret'
+    app.config['AUTH'] = auth_config or {}
+    app.config['MONGODB_DB'] = 'flask_security_example'
+    app.config['MONGODB_HOST'] = 'localhost'
+    app.config['MONGODB_PORT'] = 27017
+    
+    db = MongoEngine(app)
+    
+    Security(app, MongoEngineUserDatastore(db))
+    
+    add_endpoints(app)
+    
+    @app.before_first_request
+    def before_first_request():
+        from flask.ext.security import User, Role
+        User.drop_collection()
+        Role.drop_collection()
+        create_users()
         
     return app
 
 if __name__ == '__main__':
-    app = create_app()
+    app = create_sqlalchemy_app()
     app.run()

flask_security/__init__.py

@@ -14,13 +14,14 @@ from __future__ import absolute_import
 import sys
 
 from datetime import datetime
+from types import StringType
 
 from flask import (current_app, Blueprint, flash, redirect, request, 
     session, _request_ctx_stack, url_for, abort, g)
 
-from flask.ext.login import (AnonymousUser as AnonymousUserBase, UserMixin, 
-    LoginManager, login_required, login_user, logout_user, 
-    current_user, user_logged_in, user_logged_out)
+from flask.ext.login import (AnonymousUser as AnonymousUserBase, 
+    UserMixin as BaseUserMixin, LoginManager, login_required, login_user, 
+    logout_user, current_user, user_logged_in, user_logged_out)
 
 from flask.ext.principal import (Identity, Principal, RoleNeed, UserNeed,
     Permission, AnonymousIdentity, identity_changed, identity_loaded)
@@ -140,6 +141,28 @@ def roles_accepted(*args):
         return decorated_view
     return wrapper
 
+class RoleMixin(object):
+    def __eq__(self, other):
+        return self.name == other.name
+    
+    def __ne__(self, other):
+        return self.name != other.name
+        
+    def __str__(self):
+        return '<Role name=%s, description=%s>' % (self.name, self.description)
+
+class UserMixin(BaseUserMixin):
+    def is_active(self):
+        return self.active
+    
+    def has_role(self, role):
+        if type(role) is StringType:
+            role = Role(name=role)
+        return role in self.roles
+    
+    def __str__(self):
+        ctx = (str(self.id), self.username, self.email)
+        return '<User id=%s, username=%s, email=%s>' % ctx
 
 class AnonymousUser(AnonymousUserBase):
     def __init__(self):

flask_security/datastore/mongoengine.py

@@ -1 +1,67 @@
-# TODO: Add mongoengine datastore
+from flask.ext import security
+from flask.ext.security import UserMixin, RoleMixin
+from flask.ext.security.datastore import UserDatastore
+    
+class MongoEngineUserDatastore(UserDatastore):
+    """MongoEngine datastore"""
+    
+    def __init__(self, db):
+        self.db = db
+        
+        class Role(db.Document, RoleMixin):
+            name = db.StringField(required=True, max_length=80)
+            description = db.StringField(max_length=255)
+              
+        class User(db.Document, UserMixin):
+            username = db.StringField(max_length=255)
+            email = db.StringField(max_length=255)
+            password = db.StringField(max_length=120)
+            active = db.BooleanField(default=True)
+            roles= db.ListField(db.ReferenceField(Role), default=[])
+            created_at = db.DateTimeField()
+            modified_at = db.DateTimeField()
+            
+        security.User = User
+        security.Role = Role
+        
+    def with_id(self, id):
+        try: return security.User.objects.get(id=id)
+        except: raise security.UserIdNotFoundError()
+    
+    def find(self, user_identifier):
+        user = security.User.objects(username=user_identifier).first()
+        if user: return user
+        user = security.User.objects(email=user_identifier).first()
+        if user: return user
+        raise security.UserNotFoundError()
+    
+    def create_role(self, **kwargs):
+        if not kwargs.has_key('name'):
+            raise TypeError("create_role() did not receive "
+                            "keyword argument 'name'")
+        
+        name = kwargs.get('name')
+        description = kwargs.get('description', None)
+        
+        role = security.Role.objects(name=name).first()
+        
+        if role is None:
+            role = security.Role(name=name, description=description)
+            role.save()
+            
+        return role
+    
+    def create_user(self, **kwargs):
+        kwargs = self._prepare_create_args(kwargs)
+        
+        roles = kwargs.get('roles', [])
+        user_roles = []
+        for role in roles:
+            user_roles.append(self.create_role(name=role))
+        
+        kwargs['roles'] = user_roles
+        
+        user = security.User(**kwargs)
+        user.save()
+        
+        return user

flask_security/datastore/sqlalchemy.py

@@ -1,6 +1,5 @@
-from types import StringType
 from flask.ext import security
-from flask.ext.login import UserMixin
+from flask.ext.security import UserMixin, RoleMixin
 from flask.ext.security.datastore import UserDatastore
     
 class SQLAlchemyUserDatastore(UserDatastore):
@@ -13,7 +12,7 @@ class SQLAlchemyUserDatastore(UserDatastore):
             db.Column('user_id', db.Integer(), db.ForeignKey('role.id')),
             db.Column('role_id', db.Integer(), db.ForeignKey('user.id')))
         
-        class Role(db.Model):
+        class Role(db.Model, RoleMixin):
             id = db.Column(db.Integer(), primary_key=True)
             name = db.Column(db.String(80), unique=True)
             description = db.Column(db.String(255))
@@ -21,15 +20,6 @@ class SQLAlchemyUserDatastore(UserDatastore):
             def __init__(self, name=None, description=None):
                 self.name = name
                 self.description = description
-            
-            def __eq__(self, other):
-                return self.name == other.name
-            
-            def __ne__(self, other):
-                return self.name != other.name
-                
-            def __str__(self):
-                return '<Role name=%s, description=%s>' % (self.name, self.description)
               
         class User(db.Model, UserMixin):
             id = db.Column(db.Integer, primary_key=True)
@@ -53,17 +43,6 @@ class SQLAlchemyUserDatastore(UserDatastore):
                 self.roles = roles or []
                 self.created_at = created_at
                 self.modified_at = modified_at
-                
-            def is_active(self):
-                return self.active
-            
-            def has_role(self, role):
-                if type(role) is StringType:
-                    role = security.Role(name=role)
-                return role in self.roles
-            
-            def __str__(self):
-                return '<User id=%(id)s, email=%(email)s>' % self.__dict__
             
         security.User = User
         security.Role = Role

tests/functional_tests.py

@@ -1,19 +1,22 @@
 import unittest
 from example import app
 
-class SecurityTests(unittest.TestCase):
+class SecurityTest(unittest.TestCase):
     
     AUTH_CONFIG = None
     
     def setUp(self):
-        super(SecurityTests, self).setUp()
+        super(SecurityTest, self).setUp()
         
-        self.app = app.create_app(self.AUTH_CONFIG or None)
+        self.app = self._create_app(self.AUTH_CONFIG or None)
         self.app.debug = False
         self.app.config['TESTING'] = True
         
         self.client = self.app.test_client()
         
+    def _create_app(self, auth_config):
+        return app.create_sqlalchemy_app(auth_config)
+    
     def _get(self, route, content_type=None, follow_redirects=None):
         return self.client.get(route, follow_redirects=follow_redirects,
                 content_type=content_type or 'text/html')
@@ -22,6 +25,7 @@ class SecurityTests(unittest.TestCase):
         return self.client.post(route, data=data, 
                 follow_redirects=follow_redirects,
                 content_type=content_type or 'text/html')
+        
     
     def authenticate(self, username, password, endpoint=None):
         data = dict(username=username, password=password)
@@ -31,7 +35,7 @@ class SecurityTests(unittest.TestCase):
     def logout(self, endpoint=None):
         return self._get(endpoint or '/logout', follow_redirects=True)
 
-class DefaultSecurityTests(SecurityTests):
+class DefaultSecurityTests(SecurityTest):
     
     def test_login_view(self):
         r = self._get('/login')
@@ -97,7 +101,7 @@ class DefaultSecurityTests(SecurityTests):
         self.assertIn('Login Page', r.data)
 
 
-class ConfiguredSecurityTests(SecurityTests):    
+class ConfiguredSecurityTests(SecurityTest):    
     
     AUTH_CONFIG = {
         'password_hash': 'bcrypt',
@@ -121,3 +125,9 @@ class ConfiguredSecurityTests(SecurityTests):
         self.authenticate("matt", "password", endpoint="/custom_auth")
         r = self.logout(endpoint="/custom_logout")
         assert 'Post Logout' in r.data
+
+        
+class MongoEngineSecurityTests(DefaultSecurityTests):
+    
+    def _create_app(self, auth_config):
+        return app.create_mongoengine_app(auth_config)