From f6fa886973c71415932500e4ca16519001429d31 Mon Sep 17 00:00:00 2001
From: Andrii Lahuta <13280256+andriilahuta@users.noreply.github.com>
Date: Thu, 18 May 2023 02:52:03 +0300
Subject: [PATCH] Filter an existing django queryset (#168)

* Add queryset filtering

* Remove _get_queryset

* Add get_queryset

---------

Co-authored-by: Andrey Laguta <13280256+laevilgenius@users.noreply.github.com>
---
 django_elasticsearch_dsl/search.py | 34 +++++++++++++++++++++++-------
 tests/test_integration.py          | 18 ++++++++++++++++
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/django_elasticsearch_dsl/search.py b/django_elasticsearch_dsl/search.py
index 024e43b1..bc66d82e 100644
--- a/django_elasticsearch_dsl/search.py
+++ b/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)
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 0dce997c..fad7b753 100644
--- a/tests/test_integration.py
+++ b/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(