Implement Phrase() query in Elasticsearch backend

wagtail · May 20, 2020 · a501dc2 · a501dc2
1 parent b320ac7
commit a501dc2
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 10 deletions.
diff --git a/wagtail/search/backends/elasticsearch2.py b/wagtail/search/backends/elasticsearch2.py
@@ -14,7 +14,7 @@
     BaseSearchBackend, BaseSearchQueryCompiler, BaseSearchResults, FilterFieldError)
 from wagtail.search.index import (
     AutocompleteField, FilterField, Indexed, RelatedFields, SearchField, class_is_indexed)
-from wagtail.search.query import And, Boost, MatchAll, Not, Or, PlainText
+from wagtail.search.query import And, Boost, MatchAll, Not, Or, Phrase, PlainText
 from wagtail.utils.utils import deep_update
 
 
@@ -416,6 +416,22 @@ def _compile_plaintext_query(self, query, fields, boost=1.0):
                 'multi_match': match_query
             }
 
+    def _compile_phrase_query(self, query, fields):
+        if len(fields) == 1:
+            return {
+                'match_phrase': {
+                    fields[0]: query.query_string
+                }
+            }
+        else:
+            return {
+                'multi_match': {
+                    'query': query.query_string,
+                    'fields': fields,
+                    'type': 'phrase',
+                }
+            }
+
     def _compile_query(self, query, field, boost=1.0):
         if isinstance(query, MatchAll):
             match_all_query = {}
@@ -455,6 +471,8 @@ def _compile_query(self, query, field, boost=1.0):
         elif isinstance(query, PlainText):
             return self._compile_plaintext_query(query, [field], boost)
 
+        elif isinstance(query, Phrase):
+            return self._compile_phrase_query(query, [field])
 
         elif isinstance(query, Boost):
             return self._compile_query(query.subquery, field, boost * query.boost)
@@ -489,6 +507,9 @@ def get_inner_query(self):
         elif isinstance(self.query, PlainText):
             return self._compile_plaintext_query(self.query, fields)
 
+        elif isinstance(self.query, Phrase):
+            return self._compile_phrase_query(self.query, fields)
+
         else:
             if len(fields) == 1:
                 return self._compile_query(self.query, fields[0])

diff --git a/wagtail/search/tests/elasticsearch_common_tests.py b/wagtail/search/tests/elasticsearch_common_tests.py
@@ -191,8 +191,3 @@ def test_prefix_multiple_words(self):
     @unittest.expectedFailure
     def test_incomplete_plain_text(self):
         super().test_incomplete_plain_text()
-
-    # Elasticsearch backend doesn't support Phrase() query class
-    @unittest.expectedFailure
-    def test_phrase(self):
-        super().test_phrase()
diff --git a/wagtail/search/tests/test_elasticsearch2_backend.py b/wagtail/search/tests/test_elasticsearch2_backend.py
@@ -8,7 +8,7 @@
 from elasticsearch.serializer import JSONSerializer
 
 from wagtail.search.backends.elasticsearch2 import Elasticsearch2SearchBackend, get_model_root
-from wagtail.search.query import MATCH_ALL
+from wagtail.search.query import MATCH_ALL, Phrase
 from wagtail.tests.search import models
 
 from .elasticsearch_common_tests import ElasticsearchCommonSearchBackendTests
@@ -316,6 +316,22 @@ def test_custom_ordering_multiple(self):
         expected_result = [{'publication_date_filter': 'asc'}, {'number_of_pages_filter': 'asc'}]
         self.assertDictEqual(query_compiler.get_sort(), expected_result)
 
+    def test_phrase_query(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"))
+
+        # Check it
+        expected_result = {'multi_match': {'fields': ['_all', '_partials'], 'query': "Hello world", 'type': 'phrase'}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
+    def test_phrase_query_single_field(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"), fields=['title'])
+
+        # Check it
+        expected_result = {'match_phrase': {'title': "Hello world"}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
 
 class TestElasticsearch2SearchResults(TestCase):
     fixtures = ['search']

diff --git a/wagtail/search/tests/test_elasticsearch5_backend.py b/wagtail/search/tests/test_elasticsearch5_backend.py
@@ -8,7 +8,7 @@
 from elasticsearch.serializer import JSONSerializer
 
 from wagtail.search.backends.elasticsearch5 import Elasticsearch5SearchBackend
-from wagtail.search.query import MATCH_ALL
+from wagtail.search.query import MATCH_ALL, Phrase
 from wagtail.tests.search import models
 
 from .elasticsearch_common_tests import ElasticsearchCommonSearchBackendTests
@@ -316,6 +316,22 @@ def test_custom_ordering_multiple(self):
         expected_result = [{'publication_date_filter': 'asc'}, {'number_of_pages_filter': 'asc'}]
         self.assertDictEqual(query_compiler.get_sort(), expected_result)
 
+    def test_phrase_query(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"))
+
+        # Check it
+        expected_result = {'multi_match': {'fields': ['_all', '_partials'], 'query': "Hello world", 'type': 'phrase'}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
+    def test_phrase_query_single_field(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"), fields=['title'])
+
+        # Check it
+        expected_result = {'match_phrase': {'title': "Hello world"}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
 
 class TestElasticsearch5SearchResults(TestCase):
     fixtures = ['search']

diff --git a/wagtail/search/tests/test_elasticsearch6_backend.py b/wagtail/search/tests/test_elasticsearch6_backend.py
@@ -8,7 +8,7 @@
 from elasticsearch.serializer import JSONSerializer
 
 from wagtail.search.backends.elasticsearch6 import Elasticsearch6SearchBackend
-from wagtail.search.query import MATCH_ALL
+from wagtail.search.query import MATCH_ALL, Phrase
 from wagtail.tests.search import models
 
 from .elasticsearch_common_tests import ElasticsearchCommonSearchBackendTests
@@ -316,6 +316,22 @@ def test_custom_ordering_multiple(self):
         expected_result = [{'publication_date_filter': 'asc'}, {'number_of_pages_filter': 'asc'}]
         self.assertDictEqual(query.get_sort(), expected_result)
 
+    def test_phrase_query(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"))
+
+        # Check it
+        expected_result = {'multi_match': {'fields': ['_all_text', '_edgengrams'], 'query': "Hello world", 'type': 'phrase'}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
+    def test_phrase_query_single_field(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"), fields=['title'])
+
+        # Check it
+        expected_result = {'match_phrase': {'title': "Hello world"}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
 
 class TestElasticsearch6SearchResults(TestCase):
     fixtures = ['search']

diff --git a/wagtail/search/tests/test_elasticsearch7_backend.py b/wagtail/search/tests/test_elasticsearch7_backend.py
@@ -8,7 +8,7 @@
 from elasticsearch.serializer import JSONSerializer
 
 from wagtail.search.backends.elasticsearch7 import Elasticsearch7SearchBackend
-from wagtail.search.query import MATCH_ALL
+from wagtail.search.query import MATCH_ALL, Phrase
 from wagtail.tests.search import models
 
 from .elasticsearch_common_tests import ElasticsearchCommonSearchBackendTests
@@ -316,6 +316,22 @@ def test_custom_ordering_multiple(self):
         expected_result = [{'publication_date_filter': 'asc'}, {'number_of_pages_filter': 'asc'}]
         self.assertDictEqual(query.get_sort(), expected_result)
 
+    def test_phrase_query(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"))
+
+        # Check it
+        expected_result = {'multi_match': {'fields': ['_all_text', '_edgengrams'], 'query': "Hello world", 'type': 'phrase'}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
+    def test_phrase_query_single_field(self):
+        # Create a query
+        query_compiler = self.query_compiler_class(models.Book.objects.all(), Phrase("Hello world"), fields=['title'])
+
+        # Check it
+        expected_result = {'match_phrase': {'title': "Hello world"}}
+        self.assertDictEqual(query_compiler.get_inner_query(), expected_result)
+
 
 class TestElasticsearch7SearchResults(TestCase):
     fixtures = ['search']