Add fields to stored hits (#73)

Author: hugorodgerbrown

elasticsearch_django/models.py

@@ -630,11 +630,18 @@ def has_aggs(self) -> bool:
 
     @property
     def has_highlights(self) -> bool:
-        """Return True if the query includes aggs."""
+        """Return True if the query includes highlights."""
         if not self.query:
             raise ValueError("Missing query attribute.")
         return "highlight" in self.query
 
+    @property
+    def has_fields(self) -> bool:
+        """Return True if the query includes explicit fields."""
+        if not self.query:
+            raise ValueError("Missing query attribute.")
+        return "fields" in self.query
+
     def search_rank_annotation(self, pk_field_name: str = "pk") -> Case | None:
         """Return SQL CASE statement used to annotate results with rank."""
         if not self.hits:
@@ -761,6 +768,8 @@ def _hit(hit: dict) -> dict:
             }
             if highlight := hit.get("highlight"):
                 retval["highlight"] = highlight
+            if fields := hit.get("fields"):
+                retval["fields"] = fields
             return retval
 
         return [_hit(h) for h in self.raw_hits]

tests/test_models.py

@@ -16,6 +16,7 @@
     SearchDocumentManagerMixin,
     SearchDocumentMixin,
     SearchQuery,
+    SearchResponseParser,
 )
 
 from .models import (
@@ -412,6 +413,57 @@ def test_from_search_results(self) -> None:
         assert obj.search_score == 3.0
 
 
+@pytest.mark.django_db
+class SearchResponseParserTests:
+    def test_parse(self) -> None:
+        response = mock.MagicMock(spec=ObjectApiResponse)
+        response.body = {
+            "took": 31,
+            "timed_out": False,
+            "_shards": {"total": 1, "successful": 1, "skipped": 0, "failed": 0},
+            "hits": {
+                "total": {"value": 190, "relation": "eq"},
+                "max_score": 7.563781,
+                "hits": [
+                    {
+                        "_index": "foo",
+                        "_id": "123",
+                        "_score": 7.563781,
+                        "_source": {"country": "GB", "city": "London", "name": "Fred"},
+                        "fields": {"country": ["gb"], "city": ["london"]},
+                        "highlight": {"country": ["<em>gb</em>"]},
+                    }
+                ],
+            },
+            "aggregations": {
+                "countries": {
+                    "doc_count_error_upper_bound": 0,
+                    "sum_other_doc_count": 0,
+                    "buckets": [{"key": "gb", "doc_count": 100}],
+                }
+            },
+        }
+        parser = SearchResponseParser(response)
+        assert parser.total_hits == 190
+        assert parser.total_hits_relation == "eq"
+        assert parser.hits == [
+            {
+                "id": "123",
+                "index": "foo",
+                "score": 7.563781,
+                "fields": {"city": ["london"], "country": ["gb"]},
+                "highlight": {"country": ["<em>gb</em>"]},
+            }
+        ]
+        assert parser.aggregations == {
+            "countries": {
+                "doc_count_error_upper_bound": 0,
+                "sum_other_doc_count": 0,
+                "buckets": [{"key": "gb", "doc_count": 100}],
+            }
+        }
+
+
 @pytest.mark.django_db
 class ExecuteFunctionTests:
     raw_hits = [