mirror of
https://github.com/wassname/flask-security.git
synced 2026-07-16 01:20:15 +08:00
Refactor datastore implementation
This commit is contained in:
Matt Wright
@@ -58,7 +58,7 @@ def generate_confirmation_token(user):
|
||||
|
||||
def requires_confirmation(user):
|
||||
"""Returns `True` if the user requires confirmation."""
|
||||
return user.confirmed_at == None and _security.confirmable
|
||||
return _security.confirmable and user.confirmed_at == None
|
||||
|
||||
|
||||
def confirm_email_token_status(token):
|
||||
@@ -78,5 +78,5 @@ def confirm_user(user):
|
||||
:param user: The user to confirm
|
||||
"""
|
||||
user.confirmed_at = datetime.utcnow()
|
||||
_datastore._save_model(user)
|
||||
_datastore.put(user)
|
||||
user_confirmed.send(user, app=app._get_current_object())
|
||||
|
||||
@@ -175,10 +175,12 @@ def _context_processor():
|
||||
class RoleMixin(object):
|
||||
"""Mixin for `Role` model definitions"""
|
||||
def __eq__(self, other):
|
||||
return self.name == other or self.name == getattr(other, 'name', None)
|
||||
return (self.name == other or \
|
||||
self.name == getattr(other, 'name', None))
|
||||
|
||||
def __ne__(self, other):
|
||||
return self.name != other and self.name != getattr(other, 'name', None)
|
||||
return (self.name != other and
|
||||
self.name != getattr(other, 'name', None))
|
||||
|
||||
|
||||
class UserMixin(BaseUserMixin):
|
||||
|
||||
+107
-127
@@ -9,65 +9,52 @@
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
class Datastore(object):
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
|
||||
def commit(self):
|
||||
pass
|
||||
|
||||
def put(self, model):
|
||||
raise NotImplementedError
|
||||
|
||||
def delete(self, model):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class SQLAlchemyDatastore(Datastore):
|
||||
def commit(self):
|
||||
self.db.session.commit()
|
||||
|
||||
def put(self, model):
|
||||
self.db.session.add(model)
|
||||
return model
|
||||
|
||||
def delete(self, model):
|
||||
self.db.session.delete(model)
|
||||
|
||||
|
||||
class MongoEngineDatastore(Datastore):
|
||||
def put(self, model):
|
||||
model.save()
|
||||
return model
|
||||
|
||||
def delete(self, model):
|
||||
model.delete()
|
||||
|
||||
|
||||
class UserDatastore(object):
|
||||
"""Abstracted user datastore. Always extend this class and implement the
|
||||
:attr:`_save_model`, :attr:`_delete_model`, :attr:`_do_find_user`, and
|
||||
:attr:`_do_find_role` methods.
|
||||
"""Abstracted user datastore.
|
||||
|
||||
:param db: An instance of a configured databse manager from a Flask
|
||||
extension such as Flask-SQLAlchemy or Flask-MongoEngine
|
||||
:param user_model: A user model class definition
|
||||
:param role_model: A role model class definition
|
||||
"""
|
||||
|
||||
def __init__(self, db, user_model, role_model):
|
||||
self.db = db
|
||||
def __init__(self, user_model, role_model):
|
||||
self.user_model = user_model
|
||||
self.role_model = role_model
|
||||
|
||||
def _commit(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def _save_model(self, model, **kwargs):
|
||||
raise NotImplementedError(
|
||||
"User datastore does not implement _save_model method")
|
||||
|
||||
def _delete_model(self, model):
|
||||
raise NotImplementedError(
|
||||
"User datastore does not implement _delete_model method")
|
||||
|
||||
def _do_find_user(self, **kwargs):
|
||||
raise NotImplementedError(
|
||||
"User datastore does not implement _do_find_user method")
|
||||
|
||||
def _do_find_role(self, **kwargs):
|
||||
raise NotImplementedError(
|
||||
"User datastore does not implement _do_find_role method")
|
||||
|
||||
def _do_add_role(self, user, role):
|
||||
user, role = self._prepare_role_modify_args(user, role)
|
||||
if role not in user.roles:
|
||||
user.roles.append(role)
|
||||
return user
|
||||
|
||||
def _do_remove_role(self, user, role):
|
||||
user, role = self._prepare_role_modify_args(user, role)
|
||||
if role in user.roles:
|
||||
user.roles.remove(role)
|
||||
return user
|
||||
|
||||
def _do_toggle_active(self, user, active):
|
||||
user = self.find_user(email=user.email)
|
||||
if active != user.active:
|
||||
user.active = active
|
||||
return user
|
||||
|
||||
def _do_deactive_user(self, user):
|
||||
return self._do_toggle_active(user, False)
|
||||
|
||||
def _do_active_user(self, user):
|
||||
return self._do_toggle_active(user, True)
|
||||
|
||||
def _prepare_role_modify_args(self, user, role):
|
||||
role = role.name if isinstance(role, self.role_model) else role
|
||||
return self.find_user(email=user.email), self.find_role(role)
|
||||
@@ -75,124 +62,117 @@ class UserDatastore(object):
|
||||
def _prepare_create_user_args(self, **kwargs):
|
||||
kwargs.setdefault('active', True)
|
||||
roles = kwargs.get('roles', [])
|
||||
|
||||
for i, role in enumerate(roles):
|
||||
rn = role.name if isinstance(role, self.role_model) else role
|
||||
# see if the role exists
|
||||
roles[i] = self.find_role(rn)
|
||||
|
||||
kwargs['roles'] = roles
|
||||
|
||||
return kwargs
|
||||
|
||||
def find_user(self, **kwargs):
|
||||
"""Returns a user based on the specified identifier.
|
||||
"""Returns a user matching the provided paramters."""
|
||||
raise NotImplementedError
|
||||
|
||||
:param user: User identifier, usually email address
|
||||
def find_role(self, **kwargs):
|
||||
"""Returns a role matching the provided paramters."""
|
||||
raise NotImplementedError
|
||||
|
||||
def add_role_to_user(self, user, role):
|
||||
"""Adds a role tp a user
|
||||
|
||||
:param user: The user to manipulate
|
||||
:param role: The role to add to the user
|
||||
"""
|
||||
return self._do_find_user(**kwargs)
|
||||
rv = False
|
||||
user, role = self._prepare_role_modify_args(user, role)
|
||||
if role not in user.roles:
|
||||
rv = True
|
||||
user.roles.append(role)
|
||||
return rv
|
||||
|
||||
def find_role(self, role):
|
||||
"""Returns a role based on its name.
|
||||
def remove_role_from_user(self, user, role):
|
||||
"""Removes a role from a user
|
||||
|
||||
:param role: Role name
|
||||
:param user: The user to manipulate
|
||||
:param role: The role to remove from the user
|
||||
"""
|
||||
return self._do_find_role(role)
|
||||
rv = False
|
||||
user, role = self._prepare_role_modify_args(user, role)
|
||||
if role in user.roles:
|
||||
rv = True
|
||||
user.roles.remove(role)
|
||||
return rv
|
||||
|
||||
def toggle_active(self, user):
|
||||
"""Toggles a user's active status. Always returns True."""
|
||||
user.active = not user.active
|
||||
return True
|
||||
|
||||
def deactivate_user(self, user):
|
||||
"""Deactivates a specified user. Returns `True` if a change was made.
|
||||
|
||||
:param user: The user to deactivate
|
||||
"""
|
||||
if user.active:
|
||||
user.active = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def activate_user(self, user):
|
||||
"""Activates a specified user. Returns `True` if a change was made.
|
||||
|
||||
:param user: The user to activate
|
||||
"""
|
||||
if not user.active:
|
||||
user.active = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def create_role(self, **kwargs):
|
||||
"""Creates and returns a new role.
|
||||
"""Creates and returns a new role from the given parameters."""
|
||||
|
||||
:param name: Role name
|
||||
"""
|
||||
role = self.role_model(**kwargs)
|
||||
return self._save_model(role)
|
||||
return self.put(role)
|
||||
|
||||
def create_user(self, **kwargs):
|
||||
"""Creates and returns a new user.
|
||||
"""Creates and returns a new user from the given parameters."""
|
||||
|
||||
:param email: Email address
|
||||
:param password: Unencrypted password
|
||||
:param active: The optional active state
|
||||
"""
|
||||
user = self.user_model(**self._prepare_create_user_args(**kwargs))
|
||||
return self._save_model(user)
|
||||
return self.put(user)
|
||||
|
||||
def delete_user(self, user):
|
||||
"""Delete the specified user
|
||||
|
||||
:param user: The user to delete_user
|
||||
:param user: The user to delete
|
||||
"""
|
||||
self._delete_model(user)
|
||||
|
||||
def add_role_to_user(self, user, role):
|
||||
"""Adds a role to a user if the user does not have it already. Returns
|
||||
the modified user.
|
||||
|
||||
:param user: A User instance or a user identifier
|
||||
:param role: A Role instance or a role name
|
||||
"""
|
||||
return self._save_model(self._do_add_role(user, role))
|
||||
|
||||
def remove_role_from_user(self, user, role):
|
||||
"""Removes a role from a user if the user has the role. Returns the
|
||||
modified user.
|
||||
|
||||
:param user: A User instance or a user identifier
|
||||
:param role: A Role instance or a role name
|
||||
"""
|
||||
return self._save_model(self._do_remove_role(user, role))
|
||||
|
||||
def deactivate_user(self, user):
|
||||
"""Deactivates a user and returns the modified user.
|
||||
|
||||
:param user: A User instance or a user identifier
|
||||
"""
|
||||
return self._save_model(self._do_deactive_user(user))
|
||||
|
||||
def activate_user(self, user):
|
||||
"""Activates a user and returns the modified user.
|
||||
|
||||
:param user: A User instance or a user identifier
|
||||
"""
|
||||
return self._save_model(self._do_active_user(user))
|
||||
self.delete(user)
|
||||
|
||||
|
||||
class SQLAlchemyUserDatastore(UserDatastore):
|
||||
class SQLAlchemyUserDatastore(SQLAlchemyDatastore, UserDatastore):
|
||||
"""A SQLAlchemy datastore implementation for Flask-Security that assumes the
|
||||
use of the Flask-SQLAlchemy extension.
|
||||
"""
|
||||
def __init__(self, db, user_model, role_model):
|
||||
SQLAlchemyDatastore.__init__(self, db)
|
||||
UserDatastore.__init__(self, user_model, role_model)
|
||||
|
||||
def _commit(self, *args, **kwargs):
|
||||
self.db.session.commit()
|
||||
|
||||
def _save_model(self, model):
|
||||
self.db.session.add(model)
|
||||
return model
|
||||
|
||||
def _delete_model(self, model):
|
||||
self.db.session.delete(model)
|
||||
|
||||
def _do_find_user(self, **kwargs):
|
||||
def find_user(self, **kwargs):
|
||||
return self.user_model.query.filter_by(**kwargs).first()
|
||||
|
||||
def _do_find_role(self, role):
|
||||
def find_role(self, role):
|
||||
return self.role_model.query.filter_by(name=role).first()
|
||||
|
||||
|
||||
class MongoEngineUserDatastore(UserDatastore):
|
||||
class MongoEngineUserDatastore(MongoEngineDatastore, UserDatastore):
|
||||
"""A MongoEngine datastore implementation for Flask-Security that assumes
|
||||
the use of the Flask-MongoEngine extension.
|
||||
"""
|
||||
def __init__(self, db, user_model, role_model):
|
||||
MongoEngineDatastore.__init__(self, db)
|
||||
UserDatastore.__init__(self, user_model, role_model)
|
||||
|
||||
def _save_model(self, model):
|
||||
model.save()
|
||||
return model
|
||||
|
||||
def _delete_model(self, model):
|
||||
model.delete()
|
||||
|
||||
def _do_find_user(self, **kwargs):
|
||||
def find_user(self, **kwargs):
|
||||
return self.user_model.objects(**kwargs).first()
|
||||
|
||||
def _do_find_role(self, role):
|
||||
def find_role(self, role):
|
||||
return self.role_model.objects(name=role).first()
|
||||
|
||||
@@ -75,6 +75,6 @@ def update_password(user, password):
|
||||
:param password: The unencrypted new password
|
||||
"""
|
||||
user.password = encrypt_password(password)
|
||||
_datastore._save_model(user)
|
||||
_datastore.put(user)
|
||||
send_password_reset_notice(user)
|
||||
password_reset.send(user, app=app._get_current_object())
|
||||
|
||||
@@ -26,7 +26,7 @@ def register_user(**kwargs):
|
||||
confirmation_link, token = None, None
|
||||
kwargs['password'] = encrypt_password(kwargs['password'])
|
||||
user = _datastore.create_user(**kwargs)
|
||||
_datastore._commit()
|
||||
_datastore.commit()
|
||||
|
||||
if _security.confirmable:
|
||||
confirmation_link, token = generate_confirmation_link(user)
|
||||
|
||||
@@ -32,7 +32,7 @@ def pprint(obj):
|
||||
def commit(fn):
|
||||
def wrapper(*args, **kwargs):
|
||||
fn(*args, **kwargs)
|
||||
_datastore._commit()
|
||||
_datastore.commit()
|
||||
return wrapper
|
||||
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ def login_user(user, remember=True):
|
||||
|
||||
user.login_count = user.login_count + 1 if user.login_count else 1
|
||||
|
||||
_datastore._save_model(user)
|
||||
_datastore.put(user)
|
||||
identity_changed.send(current_app._get_current_object(),
|
||||
identity=Identity(user.id))
|
||||
|
||||
|
||||
@@ -51,7 +51,7 @@ def _render_json(form):
|
||||
|
||||
|
||||
def _commit(response=None):
|
||||
_datastore._commit()
|
||||
_datastore.commit()
|
||||
return response
|
||||
|
||||
|
||||
|
||||
@@ -182,7 +182,7 @@ class DefaultSecurityTests(SecurityTest):
|
||||
with self.app.test_request_context('/'):
|
||||
user = self.app.security.datastore.find_user(email='matt@lp.com')
|
||||
self.app.security.datastore.delete_user(user)
|
||||
self.app.security.datastore._commit()
|
||||
self.app.security.datastore.commit()
|
||||
|
||||
r = self._get('/')
|
||||
self.assertNotIn('Hello matt@lp.com', r.data)
|
||||
|
||||
@@ -111,8 +111,7 @@ def create_app(config):
|
||||
def create_roles():
|
||||
for role in ('admin', 'editor', 'author'):
|
||||
ds.create_role(name=role)
|
||||
ds._commit()
|
||||
|
||||
ds.commit()
|
||||
|
||||
def create_users():
|
||||
for u in (('matt@lp.com', 'password', ['admin'], True),
|
||||
@@ -122,7 +121,7 @@ def create_users():
|
||||
('tiya@lp.com', 'password', [], False)):
|
||||
ds.create_user(email=u[0], password=encrypt_password(u[1]),
|
||||
roles=u[2], active=u[3])
|
||||
ds._commit()
|
||||
ds.commit()
|
||||
|
||||
def populate_data():
|
||||
create_roles()
|
||||
|
||||
+46
-13
@@ -3,18 +3,16 @@
|
||||
import unittest
|
||||
|
||||
from flask_security import RoleMixin, UserMixin, AnonymousUser
|
||||
from flask_security.datastore import UserDatastore
|
||||
from flask_security.datastore import Datastore, UserDatastore
|
||||
|
||||
|
||||
class Role(RoleMixin):
|
||||
def __init__(self, name, description=None):
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
self.description = description
|
||||
|
||||
|
||||
class User(UserMixin):
|
||||
def __init__(self, username, email, roles):
|
||||
self.username = username
|
||||
def __init__(self, email, roles):
|
||||
self.email = email
|
||||
self.roles = roles
|
||||
|
||||
@@ -22,7 +20,7 @@ admin = Role('admin')
|
||||
admin2 = Role('admin')
|
||||
editor = Role('editor')
|
||||
|
||||
user = User('matt', 'matt@lp.com', [admin, editor])
|
||||
user = User('matt@lp.com', [admin, editor])
|
||||
|
||||
|
||||
class SecurityEntityTests(unittest.TestCase):
|
||||
@@ -45,11 +43,46 @@ class SecurityEntityTests(unittest.TestCase):
|
||||
self.assertFalse(au.has_role('admin'))
|
||||
|
||||
|
||||
class UserDatastoreTests(unittest.TestCase):
|
||||
class DatastoreTests(unittest.TestCase):
|
||||
|
||||
def test_unimplemented(self):
|
||||
ds = UserDatastore(None, None, None)
|
||||
self.assertRaises(NotImplementedError, ds._save_model, None)
|
||||
self.assertRaises(NotImplementedError, ds._delete_model, None)
|
||||
self.assertRaises(NotImplementedError, ds._do_find_user)
|
||||
self.assertRaises(NotImplementedError, ds._do_find_role)
|
||||
def setUp(self):
|
||||
super(DatastoreTests, self).setUp()
|
||||
self.ds = UserDatastore(None, None)
|
||||
|
||||
def test_unimplemented_datastore_methods(self):
|
||||
ds = Datastore(None)
|
||||
self.assertRaises(NotImplementedError, ds.put, None)
|
||||
self.assertRaises(NotImplementedError, ds.delete, None)
|
||||
|
||||
def test_unimplemented_user_datastore_methods(self):
|
||||
self.assertRaises(NotImplementedError, self.ds.find_user)
|
||||
self.assertRaises(NotImplementedError, self.ds.find_role)
|
||||
|
||||
def test_toggle_active(self):
|
||||
user.active = True
|
||||
rv = self.ds.toggle_active(user)
|
||||
self.assertTrue(rv)
|
||||
self.assertFalse(user.active)
|
||||
rv = self.ds.toggle_active(user)
|
||||
self.assertTrue(rv)
|
||||
self.assertTrue(user.active)
|
||||
|
||||
def test_deactivate_user(self):
|
||||
user.active = True
|
||||
rv = self.ds.deactivate_user(user)
|
||||
self.assertTrue(rv)
|
||||
self.assertFalse(user.active)
|
||||
|
||||
def test_activate_user(self):
|
||||
ds = UserDatastore(None, None)
|
||||
user.active = False
|
||||
ds.activate_user(user)
|
||||
self.assertTrue(user.active)
|
||||
|
||||
def test_deactivate_returns_false_if_already_false(self):
|
||||
user.active = False
|
||||
self.assertFalse(self.ds.deactivate_user(user))
|
||||
|
||||
def test_activate_returns_false_if_already_true(self):
|
||||
user.active = True
|
||||
self.assertFalse(self.ds.activate_user(user))
|
||||
|
||||
Reference in New Issue
Block a user