web/lib/django/core/management/commands/loaddata.py
changeset 29 cc9b7e14412b
parent 0 0d40e90630ef
--- a/web/lib/django/core/management/commands/loaddata.py	Wed May 19 17:43:59 2010 +0200
+++ b/web/lib/django/core/management/commands/loaddata.py	Tue May 25 02:43:45 2010 +0200
@@ -4,13 +4,13 @@
 import zipfile
 from optparse import make_option
 
+from django.conf import settings
+from django.core import serializers
 from django.core.management.base import BaseCommand
 from django.core.management.color import no_style
-
-try:
-    set
-except NameError:
-    from sets import Set as set   # Python 2.3 fallback
+from django.db import connections, router, transaction, DEFAULT_DB_ALIAS
+from django.db.models import get_apps
+from django.utils.itercompat import product
 
 try:
     import bz2
@@ -22,12 +22,16 @@
     help = 'Installs the named fixture(s) in the database.'
     args = "fixture [fixture ...]"
 
+    option_list = BaseCommand.option_list + (
+        make_option('--database', action='store', dest='database',
+            default=DEFAULT_DB_ALIAS, help='Nominates a specific database to load '
+                'fixtures into. Defaults to the "default" database.'),
+    )
+
     def handle(self, *fixture_labels, **options):
-        from django.db.models import get_apps
-        from django.core import serializers
-        from django.db import connection, transaction
-        from django.conf import settings
+        using = options.get('database', DEFAULT_DB_ALIAS)
 
+        connection = connections[using]
         self.style = no_style()
 
         verbosity = int(options.get('verbosity', 1))
@@ -56,9 +60,9 @@
         # Start transaction management. All fixtures are installed in a
         # single transaction to ensure that all references are resolved.
         if commit:
-            transaction.commit_unless_managed()
-            transaction.enter_transaction_management()
-            transaction.managed(True)
+            transaction.commit_unless_managed(using=using)
+            transaction.enter_transaction_management(using=using)
+            transaction.managed(True, using=using)
 
         class SingleZipReader(zipfile.ZipFile):
             def __init__(self, *args, **kwargs):
@@ -76,7 +80,17 @@
         if has_bz2:
             compression_types['bz2'] = bz2.BZ2File
 
-        app_fixtures = [os.path.join(os.path.dirname(app.__file__), 'fixtures') for app in get_apps()]
+        app_module_paths = []
+        for app in get_apps():
+            if hasattr(app, '__path__'):
+                # It's a 'models/' subpackage
+                for path in app.__path__:
+                    app_module_paths.append(path)
+            else:
+                # It's a models.py module
+                app_module_paths.append(app.__file__)
+
+        app_fixtures = [os.path.join(os.path.dirname(path), 'fixtures') for path in app_module_paths]
         for fixture_label in fixture_labels:
             parts = fixture_label.split('.')
 
@@ -103,8 +117,8 @@
                 sys.stderr.write(
                     self.style.ERROR("Problem installing fixture '%s': %s is not a known serialization format." %
                         (fixture_name, format)))
-                transaction.rollback()
-                transaction.leave_transaction_management()
+                transaction.rollback(using=using)
+                transaction.leave_transaction_management(using=using)
                 return
 
             if os.path.isabs(fixture_name):
@@ -117,73 +131,75 @@
                     print "Checking %s for fixtures..." % humanize(fixture_dir)
 
                 label_found = False
-                for format in formats:
-                    for compression_format in compression_formats:
-                        if compression_format: 
-                            file_name = '.'.join([fixture_name, format, 
-                                                  compression_format])
-                        else: 
-                            file_name = '.'.join([fixture_name, format])
-                    
-                        if verbosity > 1:
-                            print "Trying %s for %s fixture '%s'..." % \
-                                (humanize(fixture_dir), file_name, fixture_name)
-                        full_path = os.path.join(fixture_dir, file_name)
-                        open_method = compression_types[compression_format]                                
-                        try: 
-                            fixture = open_method(full_path, 'r')
-                            if label_found:
-                                fixture.close()
-                                print self.style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
-                                    (fixture_name, humanize(fixture_dir)))
-                                transaction.rollback()
-                                transaction.leave_transaction_management()
-                                return
-                            else:
-                                fixture_count += 1
-                                objects_in_fixture = 0
-                                if verbosity > 0:
-                                    print "Installing %s fixture '%s' from %s." % \
-                                        (format, fixture_name, humanize(fixture_dir))
-                                try:
-                                    objects = serializers.deserialize(format, fixture)
-                                    for obj in objects:
+                for combo in product([using, None], formats, compression_formats):
+                    database, format, compression_format = combo
+                    file_name = '.'.join(
+                        p for p in [
+                            fixture_name, database, format, compression_format
+                        ]
+                        if p
+                    )
+
+                    if verbosity > 1:
+                        print "Trying %s for %s fixture '%s'..." % \
+                            (humanize(fixture_dir), file_name, fixture_name)
+                    full_path = os.path.join(fixture_dir, file_name)
+                    open_method = compression_types[compression_format]
+                    try:
+                        fixture = open_method(full_path, 'r')
+                        if label_found:
+                            fixture.close()
+                            print self.style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
+                                (fixture_name, humanize(fixture_dir)))
+                            transaction.rollback(using=using)
+                            transaction.leave_transaction_management(using=using)
+                            return
+                        else:
+                            fixture_count += 1
+                            objects_in_fixture = 0
+                            if verbosity > 0:
+                                print "Installing %s fixture '%s' from %s." % \
+                                    (format, fixture_name, humanize(fixture_dir))
+                            try:
+                                objects = serializers.deserialize(format, fixture, using=using)
+                                for obj in objects:
+                                    if router.allow_syncdb(using, obj.object.__class__):
                                         objects_in_fixture += 1
                                         models.add(obj.object.__class__)
-                                        obj.save()
-                                    object_count += objects_in_fixture
-                                    label_found = True
-                                except (SystemExit, KeyboardInterrupt):
-                                    raise
-                                except Exception:
-                                    import traceback
-                                    fixture.close()
-                                    transaction.rollback()
-                                    transaction.leave_transaction_management()
-                                    if show_traceback:
-                                        traceback.print_exc()
-                                    else:
-                                        sys.stderr.write(
-                                            self.style.ERROR("Problem installing fixture '%s': %s\n" %
-                                                 (full_path, ''.join(traceback.format_exception(sys.exc_type, 
-                                                     sys.exc_value, sys.exc_traceback))))) 
-                                    return
+                                        obj.save(using=using)
+                                object_count += objects_in_fixture
+                                label_found = True
+                            except (SystemExit, KeyboardInterrupt):
+                                raise
+                            except Exception:
+                                import traceback
                                 fixture.close()
-
-                                # If the fixture we loaded contains 0 objects, assume that an
-                                # error was encountered during fixture loading.
-                                if objects_in_fixture == 0:
+                                transaction.rollback(using=using)
+                                transaction.leave_transaction_management(using=using)
+                                if show_traceback:
+                                    traceback.print_exc()
+                                else:
                                     sys.stderr.write(
-                                        self.style.ERROR("No fixture data found for '%s'. (File format may be invalid.)" %
-                                            (fixture_name)))
-                                    transaction.rollback()
-                                    transaction.leave_transaction_management()
-                                    return
+                                        self.style.ERROR("Problem installing fixture '%s': %s\n" %
+                                             (full_path, ''.join(traceback.format_exception(sys.exc_type,
+                                                 sys.exc_value, sys.exc_traceback)))))
+                                return
+                            fixture.close()
 
-                        except Exception, e:
-                            if verbosity > 1:
-                                print "No %s fixture '%s' in %s." % \
-                                    (format, fixture_name, humanize(fixture_dir))
+                            # If the fixture we loaded contains 0 objects, assume that an
+                            # error was encountered during fixture loading.
+                            if objects_in_fixture == 0:
+                                sys.stderr.write(
+                                    self.style.ERROR("No fixture data found for '%s'. (File format may be invalid.)" %
+                                        (fixture_name)))
+                                transaction.rollback(using=using)
+                                transaction.leave_transaction_management(using=using)
+                                return
+
+                    except Exception, e:
+                        if verbosity > 1:
+                            print "No %s fixture '%s' in %s." % \
+                                (format, fixture_name, humanize(fixture_dir))
 
         # If we found even one object in a fixture, we need to reset the
         # database sequences.
@@ -196,11 +212,11 @@
                     cursor.execute(line)
 
         if commit:
-            transaction.commit()
-            transaction.leave_transaction_management()
+            transaction.commit(using=using)
+            transaction.leave_transaction_management(using=using)
 
         if object_count == 0:
-            if verbosity > 1:
+            if verbosity > 0:
                 print "No fixtures found."
         else:
             if verbosity > 0: