Update script commands a bit

flask_security/script.py

@@ -1,18 +1,25 @@
+# -*- coding: utf-8 -*-
+"""
+    flask.ext.security.script
+    ~~~~~~~~~~~~~~~~~~~~~~~~~
 
+    Flask-Security script module
+
+    :copyright: (c) 2012 by Matt Wright.
+    :license: MIT, see LICENSE for more details.
+"""
 try:
     import simplejson as json
 except ImportError:
     import json
 
-import inspect
-import os
 import re
 
 from flask import current_app
-from flask.ext.script import Command, Option, prompt_bool
+from flask.ext.script import Command, Option
 from werkzeug.local import LocalProxy
 
-from flask_security import views, utils
+from .utils import encrypt_password
 
 
 _datastore = LocalProxy(lambda: current_app.extensions['security'].datastore)
@@ -22,6 +29,13 @@ def pprint(obj):
     print json.dumps(obj, sort_keys=True, indent=4)
 
 
+def commit(fn):
+    def wrapper(*args, **kwargs):
+        fn(*args, **kwargs)
+        _datastore._commit()
+    return wrapper
+
+
 class CreateUserCommand(Command):
     """Create a user"""
 
@@ -32,6 +46,7 @@ class CreateUserCommand(Command):
         Option('-r', '--roles',    dest='roles',    default=''),
     )
 
+    @commit
     def run(self, **kwargs):
         # sanitize active input
         ai = re.sub(r'\s', '', str(kwargs['active']))
@@ -40,7 +55,7 @@ class CreateUserCommand(Command):
         # sanitize role input a bit
         ri = re.sub(r'\s', '', kwargs['roles'])
         kwargs['roles'] = [] if ri == '' else ri.split(',')
-        kwargs['password'] = utils.encrypt_password(kwargs['password'])
+        kwargs['password'] = encrypt_password(kwargs['password'])
 
         _datastore.create_user(**kwargs)
 
@@ -57,6 +72,7 @@ class CreateRoleCommand(Command):
         Option('-d', '--desc', dest='description', default=None),
     )
 
+    @commit
     def run(self, **kwargs):
         _datastore.create_role(**kwargs)
         print 'Role "%(name)s" created successfully.' % kwargs
@@ -72,6 +88,7 @@ class _RoleCommand(Command):
 class AddRoleCommand(_RoleCommand):
     """Add a role to a user"""
 
+    @commit
     def run(self, user_identifier, role_name):
         _datastore.add_role_to_user(user_identifier, role_name)
         print "Role '%s' added to user '%s' successfully" % (role_name, user_identifier)
@@ -80,6 +97,7 @@ class AddRoleCommand(_RoleCommand):
 class RemoveRoleCommand(_RoleCommand):
     """Add a role to a user"""
 
+    @commit
     def run(self, user_identifier, role_name):
         _datastore.remove_role_from_user(user_identifier, role_name)
         print "Role '%s' removed from user '%s' successfully" % (role_name, user_identifier)
@@ -94,6 +112,7 @@ class _ToggleActiveCommand(Command):
 class DeactivateUserCommand(_ToggleActiveCommand):
     """Deactive a user"""
 
+    @commit
     def run(self, user_identifier):
         _datastore.deactivate_user(user_identifier)
         print "User '%s' has been deactivated" % user_identifier
@@ -102,6 +121,7 @@ class DeactivateUserCommand(_ToggleActiveCommand):
 class ActivateUserCommand(_ToggleActiveCommand):
     """Deactive a user"""
 
+    @commit
     def run(self, user_identifier):
         _datastore.activate_user(user_identifier)
         print "User '%s' has been activated" % user_identifier