diff --git a/mongoengine/base.py b/mongoengine/base.py index 836817da..0cbd707d 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -255,6 +255,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): 'index_background': False, 'index_drop_dups': False, 'index_opts': {}, + 'queryset_class': QuerySet, } meta.update(base_meta) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index c199d64f..c971d387 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -999,7 +999,8 @@ class QuerySetManager(object): self._collection = db[collection] # owner is the document that contains the QuerySetManager - queryset = QuerySet(owner, self._collection) + queryset_class = owner._meta['queryset_class'] or QuerySet + queryset = queryset_class(owner, self._collection) if self._manager_func: if self._manager_func.func_code.co_argcount == 1: queryset = self._manager_func(queryset) diff --git a/tests/queryset.py b/tests/queryset.py index 3d714be2..4491be8c 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1378,6 +1378,26 @@ class QTest(unittest.TestCase): self.assertEqual(Post.objects.filter(created_user=user).count(), 1) self.assertEqual(Post.objects.filter(Q(created_user=user)).count(), 1) + def test_custom_querysets(self): + """Ensure that custom QuerySet classes may be used. + """ + class CustomQuerySet(QuerySet): + def not_empty(self): + return len(self) > 0 + + class Post(Document): + meta = {'queryset_class': CustomQuerySet} + + Post.drop_collection() + + self.assertTrue(isinstance(Post.objects, CustomQuerySet)) + self.assertFalse(Post.objects.not_empty()) + + Post().save() + self.assertTrue(Post.objects.not_empty()) + + Post.drop_collection() + if __name__ == '__main__': unittest.main()