Ticket #14662: multidb_management_patch.2.diff

django/contrib/contenttypes/management.py

@@ -1 +1 @@
 from django.contrib.contenttypes.models import ContentType
 from django.db.models import get_apps, get_models, signals
 from django.utils.encoding import smart_unicode
+from django.db import DEFAULT_DB_ALIAS
 
 def update_contenttypes(app, created_models, verbosity=2, **kwargs):
     """
     Creates content types for models in the given app, removing any model
     entries that no longer have a matching model class.
     """
-    ContentType.objects.clear_cache()
-    content_types = list(ContentType.objects.filter(app_label=app.__name__.split('.')[-2]))
+    database = kwargs.get('db', DEFAULT_DB_ALIAS)
+    ContentType.objects.db_manager(database).clear_cache()
+    content_types = list(ContentType.objects.using(database).filter(app_label=app.__name__.split('.')[-2]))
     app_models = get_models(app)
     if not app_models:
         return
     for klass in app_models:
         opts = klass._meta
         try:
-            ct = ContentType.objects.get(app_label=opts.app_label,
-                                         model=opts.object_name.lower())
+            ct = ContentType.objects.using(database).get(app_label=opts.app_label,
+                                                         model=opts.object_name.lower())
             content_types.remove(ct)
         except ContentType.DoesNotExist:
             ct = ContentType(name=smart_unicode(opts.verbose_name_raw),
                 app_label=opts.app_label, model=opts.object_name.lower())
-            ct.save()
+            ct.save(using=database)
             if verbosity >= 2:
                 print "Adding content type '%s | %s'" % (ct.app_label, ct.model)
     # The presence of any remaining content types means the supplied app has an
@@ -45 +47 @@
             for ct in content_types:
                 if verbosity >= 2:
                     print "Deleting stale content type '%s | %s'" % (ct.app_label, ct.model)
-                ct.delete()
+                ct.delete(using=database)
         else:
             if verbosity >= 2:
                 print "Stale content types remain."

django/contrib/auth/management/commands/createsuperuser.py

@@ -11 +11 @@
 from django.core import exceptions
 from django.core.management.base import BaseCommand, CommandError
 from django.utils.translation import ugettext as _
+from django.db import DEFAULT_DB_ALIAS
 
 RE_VALID_USERNAME = re.compile('[\w.@+-]+$')
 
@@ -34 +35 @@
                  '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 for superuser creation. Defaults to the "default" database.'),
     )
     help = 'Used to create a superuser.'
 
@@ -42 +46 @@
         email = options.get('email', None)
         interactive = options.get('interactive')
         verbosity = int(options.get('verbosity', 1))
+        database = options.get('database', DEFAULT_DB_ALIAS)
 
         # Do quick and dirty validation if --noinput
         if not interactive:
@@ -70 +75 @@
         # it as an option.
         if default_username:
             try:
-                User.objects.get(username=default_username)
+                User.objects.using(database).get(username=default_username)
             except User.DoesNotExist:
                 pass
             else:
@@ -95 +100 @@
                         username = None
                         continue
                     try:
-                        User.objects.get(username=username)
+                        User.objects.using(database).get(username=username)
                     except User.DoesNotExist:
                         break
                     else:
@@ -132 +137 @@
                 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")
 

django/contrib/auth/management/__init__.py

@@ -4 +4 @@
 
 from django.contrib.auth import models as auth_app
 from django.db.models import get_models, signals
+from django.db import DEFAULT_DB_ALIAS
 
-
 def _get_permission_codename(action, opts):
     return u'%s_%s' % (action, opts.object_name.lower())
 
@@ -21 +21 @@
 
     app_models = get_models(app)
 
+    database = kwargs.get('db', DEFAULT_DB_ALIAS)
     # This will hold the permissions we're looking for as
     # (content_type, (codename, name))
     searched_perms = set()
     # The codenames and ctypes that should exist.
     ctypes = set()
     for klass in app_models:
-        ctype = ContentType.objects.get_for_model(klass)
+        ctype = ContentType.objects.db_manager(database).get_for_model(klass)
         ctypes.add(ctype)
         for perm in _get_all_permissions(klass._meta):
             searched_perms.add((ctype, perm))
@@ -35 +36 @@
     # Find all the Permissions that have a context_type for a model we're
     # looking for.  We don't need to check for codenames since we already have
     # a list of the ones we're going to create.
-    all_perms = set(auth_app.Permission.objects.filter(
+    all_perms = set(auth_app.Permission.objects.using(database).filter(
         content_type__in=ctypes,
     ).values_list(
         "content_type", "codename"
@@ -45 +46 @@
         # If the permissions exists, move on.
         if (ctype.pk, codename) in all_perms:
             continue
-        p = auth_app.Permission.objects.create(
+        p = auth_app.Permission.objects.using(database).create(
             codename=codename,
             name=name,
             content_type=ctype
@@ -56 +57 @@
 
 def create_superuser(app, created_models, verbosity, **kwargs):
     from django.core.management import call_command
+    database = kwargs.get('db', DEFAULT_DB_ALIAS)
 
     if auth_app.User in created_models and kwargs.get('interactive', True):
         msg = ("\nYou just installed Django's auth system, which means you "
@@ -67 +69 @@
                 confirm = raw_input('Please enter either "yes" or "no": ')
                 continue
             if confirm == 'yes':
-                call_command("createsuperuser", interactive=True)
+                call_command("createsuperuser", interactive=True, database=database)
             break
 
 signals.post_syncdb.connect(create_permissions,