Fixes #17327 -- Add --database option to createsuperuser and change password management commands

git-svn-id: http://code.djangoproject.com/svn/django/trunk@17665 bcc190cf-cafb-0310-a4f2-bffc1f526a37

5 files changed, 137 insertions, 8 deletions

diff --git a/django/contrib/auth/management/commands/changepassword.py b/django/contrib/auth/management/commands/changepassword.py
index 56448f1424..e2815ed117 100644
--- a/django/contrib/auth/management/commands/changepassword.py
+++ b/django/contrib/auth/management/commands/changepassword.py
@@ -1,8 +1,17 @@
+import getpass
+from optparse import make_option
+
 from django.core.management.base import BaseCommand, CommandError
 from django.contrib.auth.models import User
-import getpass
+from django.db import DEFAULT_DB_ALIAS
+
 
 class Command(BaseCommand):
+    option_list = BaseCommand.option_list + (
+        make_option('--database', action='store', dest='database',
+            default=DEFAULT_DB_ALIAS, help='Nominates a database to query for the user. '
+                    'Defaults to the "default" database.'),
+    )
     help = "Change a user's password for django.contrib.auth."
 
     requires_model_validation = False
@@ -23,11 +32,11 @@ class Command(BaseCommand):
             username = getpass.getuser()
 
         try:
-            u = User.objects.get(username=username)
+            u = User.objects.using(options.get('database')).get(username=username)
         except User.DoesNotExist:
             raise CommandError("user '%s' does not exist" % username)
 
-        print "Changing password for user '%s'" % u.username
+        self.stdout.write("Changing password for user '%s'\n" % u.username)
 
         MAX_TRIES = 3
         count = 0
@@ -36,7 +45,7 @@ class Command(BaseCommand):
             p1 = self._get_pass()
             p2 = self._get_pass("Password (again): ")
             if p1 != p2:
-                print "Passwords do not match. Please try again."
+                self.stdout.write("Passwords do not match. Please try again.\n")
                 count = count + 1
 
         if count == MAX_TRIES:
diff --git a/django/contrib/auth/management/commands/createsuperuser.py b/django/contrib/auth/management/commands/createsuperuser.py
index b3a3479c28..89a76b3279 100644
--- a/django/contrib/auth/management/commands/createsuperuser.py
+++ b/django/contrib/auth/management/commands/createsuperuser.py
@@ -6,10 +6,12 @@ import getpass
 import re
 import sys
 from optparse import make_option
+
 from django.contrib.auth.models import User
 from django.contrib.auth.management import get_default_username
 from django.core import exceptions
 from django.core.management.base import BaseCommand, CommandError
+from django.db import DEFAULT_DB_ALIAS
 from django.utils.translation import ugettext as _
 
 RE_VALID_USERNAME = re.compile('[\w.@+-]+$')
@@ -19,10 +21,12 @@ EMAIL_RE = re.compile(
     r'|^"([\001-\010\013\014\016-\037!#-\[\]-\177]|\\[\001-\011\013\014\016-\177])*"' # quoted-string
     r')@(?:[A-Z0-9-]+\.)+[A-Z]{2,6}$', re.IGNORECASE)  # domain
 
+
 def is_valid_email(value):
     if not EMAIL_RE.search(value):
         raise exceptions.ValidationError(_('Enter a valid e-mail address.'))
 
+
 class Command(BaseCommand):
     option_list = BaseCommand.option_list + (
         make_option('--username', dest='username', default=None,
@@ -34,6 +38,9 @@ class Command(BaseCommand):
                   'You must use --username and --email with --noinput, and '
                   'superusers created with --noinput will not be able to log '
                   'in until they\'re given a valid password.')),
+        make_option('--database', action='store', dest='database',
+            default=DEFAULT_DB_ALIAS, help='Nominates a database to save the user to. '
+                    'Defaults to the "default" database.'),
     )
     help = 'Used to create a superuser.'
 
@@ -42,6 +49,7 @@ class Command(BaseCommand):
         email = options.get('email', None)
         interactive = options.get('interactive')
         verbosity = int(options.get('verbosity', 1))
+        database = options.get('database')
 
         # Do quick and dirty validation if --noinput
         if not interactive:
@@ -77,7 +85,7 @@ class Command(BaseCommand):
                         username = None
                         continue
                     try:
-                        User.objects.get(username=username)
+                        User.objects.using(database).get(username=username)
                     except User.DoesNotExist:
                         break
                     else:
@@ -114,7 +122,7 @@ class Command(BaseCommand):
                 sys.stderr.write("\nOperation cancelled.\n")
                 sys.exit(1)
 
-        User.objects.create_superuser(username, email, password)
+        User.objects.db_manager(database).create_superuser(username, email, password)
         if verbosity >= 1:
           self.stdout.write("Superuser created successfully.\n")
 
diff --git a/django/contrib/auth/tests/__init__.py b/django/contrib/auth/tests/__init__.py
index 0c650d7237..cd314f6050 100644
--- a/django/contrib/auth/tests/__init__.py
+++ b/django/contrib/auth/tests/__init__.py
@@ -9,7 +9,12 @@ from django.contrib.auth.tests.forms import (UserCreationFormTest,
     UserChangeFormTest, PasswordResetFormTest)
 from django.contrib.auth.tests.remote_user import (RemoteUserTest,
     RemoteUserNoCreateTest, RemoteUserCustomTest)
-from django.contrib.auth.tests.management import GetDefaultUsernameTestCase
+from django.contrib.auth.tests.management import (
+    GetDefaultUsernameTestCase,
+    ChangepasswordManagementCommandTestCase,
+    MultiDBChangepasswordManagementCommandTestCase,
+    MultiDBCreatesuperuserTestCase,
+)
 from django.contrib.auth.tests.models import (ProfileTestCase, NaturalKeysTestCase,
     LoadDataWithoutNaturalKeysTestCase, LoadDataWithNaturalKeysTestCase,
     UserManagerTestCase)
diff --git a/django/contrib/auth/tests/management.py b/django/contrib/auth/tests/management.py
index 0af68736ff..2287ab91ad 100644
--- a/django/contrib/auth/tests/management.py
+++ b/django/contrib/auth/tests/management.py
@@ -1,5 +1,9 @@
-from django.test import TestCase
+from StringIO import StringIO
+
 from django.contrib.auth import models, management
+from django.contrib.auth.management.commands import changepassword
+from django.core.management import call_command
+from django.test import TestCase
 
 
 class GetDefaultUsernameTestCase(TestCase):
@@ -25,3 +29,94 @@ class GetDefaultUsernameTestCase(TestCase):
         # 'Julia' with accented 'u':
         management.get_system_username = lambda: u'J\xfalia'
         self.assertEqual(management.get_default_username(), 'julia')
+
+
+class ChangepasswordManagementCommandTestCase(TestCase):
+
+    def setUp(self):
+        self.user = models.User.objects.create_user(username='joe', password='qwerty')
+        self.stdout = StringIO()
+        self.stderr = StringIO()
+
+    def tearDown(self):
+        self.stdout.close()
+        self.stderr.close()
+
+    def test_that_changepassword_command_changes_joes_password(self):
+        " Executing the changepassword management command should change joe's password "
+        self.assertTrue(self.user.check_password('qwerty'))
+        command = changepassword.Command()
+        command._get_pass = lambda *args: 'not qwerty'
+
+        command.execute("joe", stdout=self.stdout)
+        command_output = self.stdout.getvalue().strip()
+
+        self.assertEquals(command_output, "Changing password for user 'joe'\nPassword changed successfully for user 'joe'")
+        self.assertTrue(models.User.objects.get(username="joe").check_password("not qwerty"))
+
+    def test_that_max_tries_exits_1(self):
+        """
+        A CommandError should be thrown by handle() if the user enters in
+        mismatched passwords three times. This should be caught by execute() and
+        converted to a SystemExit
+        """
+        command = changepassword.Command()
+        command._get_pass = lambda *args: args or 'foo'
+
+        self.assertRaises(
+            SystemExit,
+            command.execute,
+            "joe",
+            stdout=self.stdout,
+            stderr=self.stderr
+        )
+
+
+class MultiDBChangepasswordManagementCommandTestCase(TestCase):
+    multi_db = True
+
+    def setUp(self):
+        self.user = models.User.objects.db_manager('other').create_user(username='joe', password='qwerty')
+        self.stdout = StringIO()
+
+    def tearDown(self):
+        self.stdout.close()
+
+    def test_that_changepassword_command_with_database_option_uses_given_db(self):
+        """
+        Executing the changepassword management command with a database option
+        should operate on the specified DB
+        """
+        self.assertTrue(self.user.check_password('qwerty'))
+        command = changepassword.Command()
+        command._get_pass = lambda *args: 'not qwerty'
+
+        command.execute("joe", database='other', stdout=self.stdout)
+        command_output = self.stdout.getvalue().strip()
+
+        self.assertEquals(command_output, "Changing password for user 'joe'\nPassword changed successfully for user 'joe'")
+        self.assertTrue(models.User.objects.using('other').get(username="joe").check_password("not qwerty"))
+
+
+class MultiDBCreatesuperuserTestCase(TestCase):
+    multi_db = True
+
+    def test_createsuperuser_command_with_database_option(self):
+        " createsuperuser command should operate on specified DB"
+        new_io = StringIO()
+
+        call_command("createsuperuser",
+            interactive=False,
+            username="joe",
+            email="joe@somewhere.org",
+            database='other',
+            stdout=new_io
+        )
+        command_output = new_io.getvalue().strip()
+
+        self.assertEqual(command_output, 'Superuser created successfully.')
+
+        u = models.User.objects.using('other').get(username="joe")
+        self.assertEqual(u.email, 'joe@somewhere.org')
+
+        new_io.close()
diff --git a/docs/ref/django-admin.txt b/docs/ref/django-admin.txt
index 508ec4183c..5af56fa248 100644
--- a/docs/ref/django-admin.txt
+++ b/docs/ref/django-admin.txt
@@ -1197,6 +1197,12 @@ the user given as parameter. If they both match, the new password will be
 changed immediately. If you do not supply a user, the command will attempt to
 change the password whose username matches the current user.
 
+.. versionadded:: 1.4
+.. django-admin-option:: --database
+
+The ``--database`` option can be used to specify the database to query for the
+user. If it is not supplied the ``default`` database will be used.
+
 Example usage::
 
     django-admin.py changepassword ringo
@@ -1227,6 +1233,12 @@ using the ``--username`` and ``--email`` arguments on the command
 line. If either of those is not supplied, ``createsuperuser`` will prompt for
 it when running interactively.
 
+.. versionadded:: 1.4
+.. django-admin-option:: --database
+
+The ``--database`` option can be used to specify the database into which the 
+superuser object will be saved.
+
 ``django.contrib.gis``
 ----------------------