From f928db298d15ef29f0b2536b320032d73d5293b7 Mon Sep 17 00:00:00 2001 From: Matt Wright Date: Fri, 24 Aug 2012 11:38:25 -0400 Subject: [PATCH] Refactor datastore implementation --- flask_security/confirmable.py | 4 +- flask_security/core.py | 6 +- flask_security/datastore.py | 234 +++++++++++++++------------------ flask_security/recoverable.py | 2 +- flask_security/registerable.py | 2 +- flask_security/script.py | 2 +- flask_security/utils.py | 2 +- flask_security/views.py | 2 +- tests/functional_tests.py | 2 +- tests/test_app/__init__.py | 5 +- tests/unit_tests.py | 59 +++++++-- 11 files changed, 167 insertions(+), 153 deletions(-) diff --git a/flask_security/confirmable.py b/flask_security/confirmable.py index 32ac958..d53c0a3 100644 --- a/flask_security/confirmable.py +++ b/flask_security/confirmable.py @@ -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()) diff --git a/flask_security/core.py b/flask_security/core.py index ec60e92..1197d82 100644 --- a/flask_security/core.py +++ b/flask_security/core.py @@ -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): diff --git a/flask_security/datastore.py b/flask_security/datastore.py index 3d46ca2..975c0f0 100644 --- a/flask_security/datastore.py +++ b/flask_security/datastore.py @@ -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() diff --git a/flask_security/recoverable.py b/flask_security/recoverable.py index b9030f7..5138680 100644 --- a/flask_security/recoverable.py +++ b/flask_security/recoverable.py @@ -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()) diff --git a/flask_security/registerable.py b/flask_security/registerable.py index fa842d2..2e29af1 100644 --- a/flask_security/registerable.py +++ b/flask_security/registerable.py @@ -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) diff --git a/flask_security/script.py b/flask_security/script.py index 9e19e1f..7c798f2 100644 --- a/flask_security/script.py +++ b/flask_security/script.py @@ -32,7 +32,7 @@ def pprint(obj): def commit(fn): def wrapper(*args, **kwargs): fn(*args, **kwargs) - _datastore._commit() + _datastore.commit() return wrapper diff --git a/flask_security/utils.py b/flask_security/utils.py index 9a8fa6d..1751132 100644 --- a/flask_security/utils.py +++ b/flask_security/utils.py @@ -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)) diff --git a/flask_security/views.py b/flask_security/views.py index 376ff09..dca3f61 100644 --- a/flask_security/views.py +++ b/flask_security/views.py @@ -51,7 +51,7 @@ def _render_json(form): def _commit(response=None): - _datastore._commit() + _datastore.commit() return response diff --git a/tests/functional_tests.py b/tests/functional_tests.py index 9415fea..d0e99d7 100644 --- a/tests/functional_tests.py +++ b/tests/functional_tests.py @@ -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) diff --git a/tests/test_app/__init__.py b/tests/test_app/__init__.py index 93e686f..4b7b557 100644 --- a/tests/test_app/__init__.py +++ b/tests/test_app/__init__.py @@ -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() diff --git a/tests/unit_tests.py b/tests/unit_tests.py index 44860a4..d08db38 100644 --- a/tests/unit_tests.py +++ b/tests/unit_tests.py @@ -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))