Ticket #16330: dumpdata-primary-keys3.diff

django/core/management/commands/dumpdata.py

@@ -6 +6 @@
 
 from optparse import make_option
 
+import itertools
+import sys
+
+
 class Command(BaseCommand):
     option_list = BaseCommand.option_list + (
         make_option('--format', default='json', dest='format',
@@ -21 +25 @@
             help='Use natural keys if they are available.'),
         make_option('-a', '--all', action='store_true', dest='use_base_manager', default=False,
             help="Use Django's base manager to dump all models stored in the database, including those that would otherwise be filtered or modified by a custom manager."),
+        make_option('--pks', dest='primary_keys', action='append', default=[],
+            help="Only dump objects with given primary keys. Accepts a space seperated list of keys or '-' to read from stdin. This option is applied to all apps/models."),
     )
     help = ("Output the contents of the database as a fixture of the given "
             "format (using each model's default manager unless --all is "
@@ -34 +40 @@
         indent = options.get('indent')
         using = options.get('database')
         excludes = options.get('exclude')
+        pks = options.get('primary_keys')
         show_traceback = options.get('traceback')
         use_natural_keys = options.get('use_natural_keys')
         use_base_manager = options.get('use_base_manager')
@@ -54 +61 @@
                 except ImproperlyConfigured:
                     raise CommandError('Unknown app in excludes: %s' % exclude)
 
+        if pks == ['-']:
+            # Read from Stdin
+            primary_keys = sys.stdin.read().split()
+        else:
+            primary_keys = list(itertools.chain(*[pk.split() for pk in pks]))
+
         if len(app_labels) == 0:
             app_list = SortedDict((app, None) for app in get_apps() if app not in excluded_apps)
         else:
@@ -107 +120 @@
                         objects = model._base_manager
                     else:
                         objects = model._default_manager
-                    for obj in objects.using(using).\
-                            order_by(model._meta.pk.name).iterator():
+
+                    queryset = objects.using(using).order_by(model._meta.pk.name)
+                    if primary_keys:
+                        queryset = queryset.filter(pk__in=primary_keys)
+
+                    for obj in queryset.iterator():
                         yield obj
 
         try:
@@ -120 +137 @@
                 raise
             raise CommandError("Unable to serialize database: %s" % e)
 
+
 def sort_dependencies(app_list):
     """Sort a list of app,modellist pairs into a single list of models.
 

tests/regressiontests/fixtures_regress/tests.py

@@ -13 +13 @@
 from django.test import (TestCase, TransactionTestCase, skipIfDBFeature,
     skipUnlessDBFeature)
 from django.test.utils import override_settings
+from django.utils import simplejson
 from django.utils.six import PY3, StringIO
 
 from .models import (Animal, Stuff, Absolute, Parent, Child, Article, Widget,
@@ -316 +317 @@
         self.assertTrue(emu_json in data)
         self.assertTrue(platypus_json in data)
 
+        stdout = StringIO()
+        management.call_command(
+            'dumpdata',
+            'fixtures_regress.animal',
+            format='json',
+            stdout=stdout,
+            primary_keys=['1 10']
+        )
+        data = stdout.getvalue()
+        json = simplejson.loads(data)
+        self.assertEqual(len(json), 2)
+        keys = sorted([item['pk'] for item in json])
+        self.assertEqual(keys, [1, 10])
+
     def test_proxy_model_included(self):
         """
         Regression for #11428 - Proxy models aren't included when you dumpdata