Merge pull request #34493 from dimagi/gh/elasticsearch/preference-kwarg

Add preference as query parameter in ES adapter
dimagi · Apr 23, 2024 · 78faa5f · 78faa5f
2 parents 7e6e095 + 18e75b4
commit 78faa5f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 13 deletions.
diff --git a/corehq/apps/es/client.py b/corehq/apps/es/client.py
@@ -21,6 +21,7 @@
     TransportError,
     bulk,
 )
+from corehq.toggles import ES_QUERY_PREFERENCE
 from corehq.util.global_request import get_request_domain
 from corehq.util.metrics import (
     limit_domains,
@@ -644,12 +645,19 @@ def _search(self, query, **kw):
         """Perform a "low-level" search and return the raw result. This is
         split into a separate method for ease of testing the result format.
         """
+        domain = get_request_domain()
+        if ES_QUERY_PREFERENCE.enabled(domain):
+            # Use domain as key to route to a consistent set of shards.kwargs
+            # See https://www.elastic.co/guide/en/elasticsearch/reference/5.6/search-request-preference.html
+            if 'preference' not in kw:
+                kw['preference'] = domain
+
         with metrics_histogram_timer(
                 'commcare.elasticsearch.search.timing',
                 timing_buckets=(1, 10),
                 tags={
                     'index': self.canonical_name,
-                    'domain': limit_domains(get_request_domain()),
+                    'domain': limit_domains(domain),
                 },
         ):
             return self._es.search(self.index_name, self.type, query, **kw)

diff --git a/corehq/apps/es/es_query.py b/corehq/apps/es/es_query.py
@@ -92,9 +92,7 @@
 
 from memoized import memoized
 
-from corehq.toggles import ES_QUERY_PREFERENCE
 from corehq.util.files import TransientTempfile
-from corehq.util.global_request import get_request_domain
 
 from . import aggregations, filters, queries
 from .const import SCROLL_SIZE, SIZE_LIMIT
@@ -152,16 +150,6 @@ def __init__(self, index=None, for_export=False):
                 }
             }
         }
-        self._set_preference()
-
-    def _set_preference(self):
-        """
-        If the specified domain has ES_QUERY_PREFERENCE enabled, use domain as key to route to a consistent set
-        of shards. See https://www.elastic.co/guide/en/elasticsearch/reference/5.6/search-request-preference.html
-        """
-        domain = get_request_domain()
-        if ES_QUERY_PREFERENCE.enabled(domain):
-            self.es_query['preference'] = domain
 
     def clone(self):
         adapter = self.adapter