include a link to PyMongo docs for batch_size

mongoengine/fields.py

@@ -1249,7 +1249,7 @@ class GenericReferenceField(BaseField):
         if document is None:
             return None
 
-        if isinstance(document, (dict, SON, ObjectId, DBRef)):
+        if isinstance(document, (dict, SON)):
             return document
 
         id_field_name = document.__class__._meta['id_field']

mongoengine/queryset/base.py

@@ -82,6 +82,7 @@ class BaseQuerySet(object):
         self._limit = None
         self._skip = None
         self._hint = -1  # Using -1 as None is a valid value for hint
+        self._batch_size = None
         self.only_fields = []
         self._max_time_ms = None
 
@@ -781,6 +782,19 @@ class BaseQuerySet(object):
         queryset._hint = index
         return queryset
 
+    def batch_size(self, size):
+        """Limit the number of documents returned in a single batch (each
+        batch requires a round trip to the server).
+
+        See http://api.mongodb.com/python/current/api/pymongo/cursor.html#pymongo.cursor.Cursor.batch_size
+        for details.
+
+        :param size: desired size of each batch.
+        """
+        queryset = self.clone()
+        queryset._batch_size = size
+        return queryset
+
     def distinct(self, field):
         """Return a list of distinct values for a given field.
 
@@ -1467,6 +1481,9 @@ class BaseQuerySet(object):
             if self._hint != -1:
                 self._cursor_obj.hint(self._hint)
 
+            if self._batch_size is not None:
+                self._cursor_obj.batch_size(self._batch_size)
+
         return self._cursor_obj
 
     def __deepcopy__(self, memo):

mongoengine/queryset/transform.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
 
-from bson import ObjectId, SON
-from bson.dbref import DBRef
+from bson import SON
 import pymongo
 
 from mongoengine.base.fields import UPDATE_OPERATORS
@@ -27,7 +26,6 @@ MATCH_OPERATORS = (COMPARISON_OPERATORS + GEO_OPERATORS +
                    STRING_OPERATORS + CUSTOM_OPERATORS)
 
 
-# TODO make this less complex
 def query(_doc_cls=None, **kwargs):
     """Transform a query from Django-style format to Mongo format.
     """
@@ -64,7 +62,6 @@ def query(_doc_cls=None, **kwargs):
             parts = []
 
             CachedReferenceField = _import_class('CachedReferenceField')
-            GenericReferenceField = _import_class('GenericReferenceField')
 
             cleaned_fields = []
             for field in fields:
@@ -104,16 +101,6 @@ def query(_doc_cls=None, **kwargs):
                 # 'in', 'nin' and 'all' require a list of values
                 value = [field.prepare_query_value(op, v) for v in value]
 
-            # If we're querying a GenericReferenceField, we need to alter the
-            # key depending on the value:
-            # * If the value is a DBRef, the key should be "field_name._ref".
-            # * If the value is an ObjectId, the key should be "field_name._ref.$id".
-            if isinstance(field, GenericReferenceField):
-                if isinstance(value, DBRef):
-                    parts[-1] += '._ref'
-                elif isinstance(value, ObjectId):
-                    parts[-1] += '._ref.$id'
-
         # if op and op not in COMPARISON_OPERATORS:
         if op:
             if op in GEO_OPERATORS:
@@ -141,13 +128,11 @@ def query(_doc_cls=None, **kwargs):
 
         for i, part in indices:
             parts.insert(i, part)
-
         key = '.'.join(parts)
-
         if op is None or key not in mongo_query:
             mongo_query[key] = value
         elif key in mongo_query:
-            if isinstance(mongo_query[key], dict):
+            if key in mongo_query and isinstance(mongo_query[key], dict):
                 mongo_query[key].update(value)
                 # $max/minDistance needs to come last - convert to SON
                 value_dict = mongo_query[key]

setup.cfg

@@ -9,5 +9,5 @@ tests = tests
 [flake8]
 ignore=E501,F401,F403,F405,I201
 exclude=build,dist,docs,venv,.tox,.eggs,tests
-max-complexity=45
+max-complexity=42
 application-import-names=mongoengine,tests

tests/fields/fields.py

@@ -2810,38 +2810,6 @@ class FieldTest(unittest.TestCase):
         Post.drop_collection()
         User.drop_collection()
 
-    def test_generic_reference_filter_by_dbref(self):
-        """Ensure we can search for a specific generic reference by
-        providing its ObjectId.
-        """
-        class Doc(Document):
-            ref = GenericReferenceField()
-
-        Doc.drop_collection()
-
-        doc1 = Doc.objects.create()
-        doc2 = Doc.objects.create(ref=doc1)
-
-        doc = Doc.objects.get(ref=DBRef('doc', doc1.pk))
-        self.assertEqual(doc, doc2)
-
-    def test_generic_reference_filter_by_objectid(self):
-        """Ensure we can search for a specific generic reference by
-        providing its DBRef.
-        """
-        class Doc(Document):
-            ref = GenericReferenceField()
-
-        Doc.drop_collection()
-
-        doc1 = Doc.objects.create()
-        doc2 = Doc.objects.create(ref=doc1)
-
-        self.assertTrue(isinstance(doc1.pk, ObjectId))
-
-        doc = Doc.objects.get(ref=doc1.pk)
-        self.assertEqual(doc, doc2)
-
     def test_binary_fields(self):
         """Ensure that binary fields can be stored and retrieved.
         """

tests/queryset/queryset.py

@@ -337,6 +337,34 @@ class QuerySetTest(unittest.TestCase):
         query = query.filter(boolfield=True)
         self.assertEqual(query.count(), 1)
 
+    def test_batch_size(self):
+        """Ensure that batch_size works."""
+        class A(Document):
+            s = StringField()
+
+        A.drop_collection()
+
+        for i in range(100):
+            A.objects.create(s=str(i))
+
+        # test iterating over the result set
+        cnt = 0
+        for a in A.objects.batch_size(10):
+            cnt += 1
+        self.assertEqual(cnt, 100)
+
+        # test chaining
+        qs = A.objects.all()
+        qs = qs.limit(10).batch_size(20).skip(91)
+        cnt = 0
+        for a in qs:
+            cnt += 1
+        self.assertEqual(cnt, 9)
+
+        # test invalid batch size
+        qs = A.objects.batch_size(-1)
+        self.assertRaises(ValueError, lambda: list(qs))
+
     def test_update_write_concern(self):
         """Test that passing write_concern works"""
         self.Person.drop_collection()