Filter an existing django queryset (#168)

* Add queryset filtering

* Remove _get_queryset

* Add get_queryset

---------

Co-authored-by: Andrey Laguta <[email protected]>

Author: andriilahuta

django_elasticsearch_dsl/search.py

@@ -14,12 +14,17 @@ def _clone(self):
         s._model = self._model
         return s
 
-    def to_queryset(self, keep_order=True):
+    def filter_queryset(self, queryset, keep_search_order=True):
         """
-        This method return a django queryset from the an elasticsearch result.
-        It cost a query to the sql db.
+        Filter an existing django queryset using the elasticsearch result.
+        It costs a query to the sql db.
         """
         s = self
+        if s._model is not queryset.model:
+            raise TypeError(
+                'Unexpected queryset model '
+                '(should be: %s, got: %s)' % (s._model, queryset.model)
+            )
 
         # Do not query again if the es result is already cached
         if not hasattr(self, '_response'):
@@ -28,14 +33,27 @@ def to_queryset(self, keep_order=True):
             s = s.execute()
 
         pks = [result.meta.id for result in s]
+        queryset = queryset.filter(pk__in=pks)
 
-        qs = self._model.objects.filter(pk__in=pks)
-
-        if keep_order:
+        if keep_search_order:
             preserved_order = Case(
                 *[When(pk=pk, then=pos) for pos, pk in enumerate(pks)],
                 output_field=IntegerField()
             )
-            qs = qs.order_by(preserved_order)
+            queryset = queryset.order_by(preserved_order)
+
+        return queryset
+
+    def _get_queryset(self):
+        """
+        Return a django queryset that will be filtered by to_queryset method.
+        """
+        return self._model._default_manager.all()
 
-        return qs
+    def to_queryset(self, keep_order=True):
+        """
+        Return a django queryset from the elasticsearch result.
+        It costs a query to the sql db.
+        """
+        qs = self._get_queryset()
+        return self.filter_queryset(qs, keep_order)

tests/test_integration.py

@@ -331,6 +331,24 @@ def test_rebuild_command(self):
         result = AdDocument().search().execute()
         self.assertEqual(len(result), 3)
 
+    def test_filter_queryset(self):
+        Ad(title="Nothing that match",  car=self.car1).save()
+
+        qs = AdDocument().search().query(
+            'match', title="Ad number 2").filter_queryset(Ad.objects)
+        self.assertEqual(qs.count(), 2)
+        self.assertEqual(list(qs), [self.ad2, self.ad1])
+
+        qs = AdDocument().search().query(
+            'match', title="Ad number 2"
+        ).filter_queryset(Ad.objects.filter(url="www.ad2.com"))
+        self.assertEqual(qs.count(), 1)
+        self.assertEqual(list(qs), [self.ad2])
+
+        with self.assertRaisesMessage(TypeError, 'Unexpected queryset model'):
+            AdDocument().search().query(
+                'match', title="Ad number 2").filter_queryset(Category.objects)
+
     def test_to_queryset(self):
         Ad(title="Nothing that match",  car=self.car1).save()
         qs = AdDocument().search().query(