Merge pull request pandas-dev#22 from manahl/ceate_user_improvements

Enhance arctic_create_user

Author: jamesblackburn

arctic/scripts/arctic_create_user.py

@@ -1,59 +1,61 @@
-import optparse
-import pymongo
-import uuid
+import argparse
 import base64
+from pymongo import MongoClient
+import uuid
+import logging
 
-from ..auth import get_auth, authenticate
 from ..hooks import get_mongodb_uri
-from .utils import setup_logging
+from .utils import do_db_auth
+from arctic.arctic import Arctic
+
+logger = logging.getLogger(__name__)
 
 
 def main():
-    usage = """usage: %prog [options] username ...
+    usage = """arctic_create_user --host research [--db mongoose_user] [--write] user
 
-    Create the user's personal Arctic database, and adds them, read-only
-    to the central admin database.
+    Creates the user's personal Arctic mongo database
+    Or add a user to an existing Mongo Database.
     """
-    setup_logging()
-    parser = optparse.OptionParser(usage=usage)
-    parser.add_option("--host", default='localhost', help="Hostname, or clustername. Default: localhost")
-    parser.add_option("--password", dest="password", default=None, help="Password. Default: random")
-    parser.add_option("--admin-write", dest="admin", action='store_false', default=True,
-                      help="Give write access to the admin DB. Default: False")
-    parser.add_option("--dryrun", "-n", dest="dryrun", action="store_true", help="Don't really do anything", default=False)
-    parser.add_option("--verbose", "-v", dest="verbose", action="store_true", help="Print some commentary", default=False)
-    parser.add_option("--nodb", dest="nodb", help="Don't create a 'personal' database", action="store_true", default=False)
-
-    (opts, args) = parser.parse_args()
 
-    c = pymongo.MongoClient(get_mongodb_uri(opts.host))
-    credentials = get_auth(opts.host, 'admin', 'admin')
-    if not credentials:
-        raise ValueError("You have no admin credentials for instance '%s'" % (opts.host))
+    parser = argparse.ArgumentParser(usage=usage)
+    parser.add_argument("--host", default='localhost', help="Hostname, or clustername. Default: localhost")
+    parser.add_argument("--db", default=None, help="Database to add user on. Default: mongoose_<user>")
+    parser.add_argument("--password", default=None, help="Password. Default: random")
+    parser.add_argument("--write", action='store_true', default=False, help="Used for granting write access to someone else's DB")
+    parser.add_argument("users", nargs='+', help="Users to add.")
 
-    if not authenticate(c.admin, credentials.user, credentials.password):
-        raise ValueError("Failed to authenticate to '%s' as '%s'" % (opts.host, credentials.user))
+    args = parser.parse_args()
 
-    for user in args:
+    c = MongoClient(get_mongodb_uri(args.host))
 
-        p = opts.password
+    if not do_db_auth(args.host, c, args.db if args.db else 'admin'):
+        logger.error("Failed to authenticate to '%s'. Check your admin password!" % (args.host))
+        return
 
+    for user in args.users:
+        write_access = args.write
+        p = args.password
         if p is None:
             p = base64.b64encode(uuid.uuid4().bytes).replace('/', '')[:12]
+        db = args.db
+        if not db:
+            # Users always have write access to their database
+            write_access = True
+            db = Arctic.DB_PREFIX + '_' + user
+
+        # Add the user to the database
+        c[db].add_user(user, p, read_only=not write_access)
+
+        logger.info("Granted: {user} [{permission}] to {db}".format(user=user,
+                                                                    permission='WRITE' if write_access else 'READ',
+                                                                    db=db))
+        logger.info("User creds: {db}/{user}/{password}".format(user=user,
+                                                                host=args.host,
+                                                                db=db,
+                                                                password=p,
+                                                                ))
 
-        if not opts.dryrun:
-            if opts.verbose:
-                print "Adding user %s to DB %s" % (user, opts.host)
-            if not opts.nodb:
-                if opts.verbose:
-                    print "Adding database arctic_%s to DB %s" % (user, opts.host)
-                c['arctic_' + user].add_user(user, p)
-            c.admin.add_user(user, p, read_only=opts.admin)
-        else:
-            print "DRYRUN: add user %s readonly %s nodb %s" % (user, opts.admin, opts.nodb)
-
-        if not opts.password:
-            print "%-16s %s" % (user, p)
 
 if __name__ == '__main__':
     main()

tests/integration/scripts/test_create_user.py

@@ -0,0 +1,71 @@
+from mock import patch, sentinel, call, ANY
+
+from arctic.scripts.arctic_create_user import main
+
+from ...util import run_as_main
+
+
+def test_main_minimal():
+    with patch('arctic.scripts.arctic_create_user.logger', autospec=True) as logger, \
+         patch('arctic.scripts.arctic_create_user.MongoClient', autospec=True) as MC, \
+         patch('arctic.scripts.arctic_create_user.get_mongodb_uri', autospec=True) as get_mongodb_uri, \
+         patch('arctic.scripts.arctic_create_user.do_db_auth', autospec=True) as do_db_auth:
+        run_as_main(main, '--host', 'some_host',
+                          '--password', 'asdf',
+                          'user')
+    get_mongodb_uri.assert_called_once_with('some_host')
+    MC.assert_called_once_with(get_mongodb_uri.return_value)
+    assert do_db_auth.call_args_list == [call('some_host',
+                                              MC.return_value,
+                                              'admin')]
+    assert MC.return_value.__getitem__.call_args_list == [call('arctic_user')]
+    db = MC.return_value.__getitem__.return_value
+    assert [call('user', ANY, read_only=False)] == db.add_user.call_args_list
+    assert logger.info.call_args_list == [call('Granted: user [WRITE] to arctic_user'),
+                                          call('User creds: arctic_user/user/asdf')]
+
+
+def test_main_with_db():
+    with patch('arctic.scripts.arctic_create_user.MongoClient', autospec=True) as MC, \
+         patch('arctic.scripts.arctic_create_user.get_mongodb_uri', autospec=True) as get_mongodb_uri, \
+         patch('arctic.scripts.arctic_create_user.do_db_auth', autospec=True) as do_db_auth:
+        run_as_main(main, '--host', 'some_host',
+                    '--db', 'some_db',
+                    'jblackburn')
+    get_mongodb_uri.assert_called_once_with('some_host')
+    MC.assert_called_once_with(get_mongodb_uri.return_value)
+    assert do_db_auth.call_args_list == [call('some_host',
+                                              MC.return_value,
+                                              'some_db')]
+    assert MC.return_value.__getitem__.call_args_list == [call('some_db')]
+    db = MC.return_value.__getitem__.return_value
+    assert [call('jblackburn', ANY, read_only=True)] == db.add_user.call_args_list
+
+
+def test_main_with_db_write():
+    with patch('arctic.scripts.arctic_create_user.MongoClient', autospec=True) as MC, \
+         patch('arctic.scripts.arctic_create_user.get_mongodb_uri', autospec=True) as get_mongodb_uri, \
+         patch('arctic.scripts.arctic_create_user.do_db_auth', autospec=True) as do_db_auth:
+        run_as_main(main, '--host', 'some_host',
+                    '--db', 'some_db',
+                    '--write',
+                    'jblackburn')
+    get_mongodb_uri.assert_called_once_with('some_host')
+    MC.assert_called_once_with(get_mongodb_uri.return_value)
+    assert do_db_auth.call_args_list == [call('some_host',
+                                              MC.return_value,
+                                              'some_db')]
+    assert MC.return_value.__getitem__.call_args_list == [call('some_db')]
+    db = MC.return_value.__getitem__.return_value
+    assert [call('jblackburn', ANY, read_only=False)] == db.add_user.call_args_list
+
+
+def test_no_auth():
+    with patch('arctic.scripts.arctic_create_user.logger', autospec=True) as logger, \
+         patch('arctic.scripts.arctic_create_user.MongoClient', autospec=True) as MC, \
+         patch('arctic.scripts.arctic_create_user.get_mongodb_uri', autospec=True) as get_mongodb_uri, \
+         patch('arctic.scripts.arctic_create_user.do_db_auth', autospec=True,
+               return_value=False) as do_db_auth:
+        run_as_main(main, '--host', 'some_host',
+                          'jblackburn')
+    assert logger.error.call_args_list == [call("Failed to authenticate to 'some_host'. Check your admin password!")]