ES service uses template_search instead of search

Author: nadove-ucsc

src/azul/service/elasticsearch_service.py

@@ -29,6 +29,9 @@
     Q,
     Search,
 )
+from opensearchpy.connection.connections import (
+    get_connection,
+)
 from opensearchpy.helpers.aggs import (
     Agg,
     Terms,
@@ -664,6 +667,28 @@ def page_link(*, previous):
                                   order=pagination.order)
 
 
+class TemplateSearch(Search):
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._template_params = {}
+
+    def execute(self, ignore_cache: bool = False) -> Any:
+        if ignore_cache or not hasattr(self, '_response'):
+            opensearch = get_connection(self._using)
+            body = {
+                'source': json.dumps(self.to_dict()),
+                'params': self._template_params,
+            }
+            self._response = self._response_class(
+                self,
+                opensearch.search_template(
+                    index=self._index, body=body, **self._params
+                ),
+            )
+        return self._response
+
+
 class ElasticsearchService(DocumentService):
 
     @cached_property
@@ -701,12 +726,12 @@ def create_request(self,
                        catalog: CatalogName,
                        entity_type: str,
                        doc_type: DocumentType = DocumentType.aggregate
-                       ) -> Search:
+                       ) -> TemplateSearch:
         """
         Create an Elasticsearch request against the index containing documents
         of the given entity and document types, in the given catalog.
         """
-        return Search(using=self._es_client,
-                      index=str(IndexName.create(catalog=catalog,
-                                                 qualifier=entity_type,
-                                                 doc_type=doc_type)))
+        return TemplateSearch(using=self._es_client,
+                              index=str(IndexName.create(catalog=catalog,
+                                                         qualifier=entity_type,
+                                                         doc_type=doc_type)))

test/es_test_case.py

@@ -72,7 +72,17 @@ def setUpClass(cls):
                     # indices in a single request. Speeds up deletion of indices
                     # between tests.
                     #
-                    'action.destructive_requires_name': False
+                    'action.destructive_requires_name': False,
+
+                    # The service uses template queries when reading from the
+                    # index which, during tests, can far exceed the default
+                    # script compilation rate of 75/5m. Rendering a template
+                    # query using mustache is a very cheap operation compared to
+                    # other contexts (e.g. generating bytecode for a painless
+                    # script), so performance shouldn't be significantly
+                    # affected.
+                    #
+                    'script.context.template.max_compilations_rate': 'unlimited'
                 }
             })
         except BaseException:  # no coverage