Fix: OpenSearch queries are huge (azul-private#316)

Author: nadove-ucsc

src/azul/service/query_service.py

@@ -15,9 +15,11 @@
 from typing import (
     Any,
     Generic,
+    Iterator,
     Self,
     TypeVar,
 )
+import unittest.mock
 
 import attr
 from furl import (
@@ -235,23 +237,30 @@ def prepare_query(self, skip_field_paths: tuple[FieldPath] = ()) -> Query:
         Converts the given filters into an Elasticsearch DSL Query object.
         """
         filter_list = []
+        plugin = self.plugin
+        source_id_field_name = plugin.special_fields.source_id.name
+        source_id_field_path = plugin.field_mapping[source_id_field_name]
         for field_path, filter in self.prepared_filters.items():
             if field_path not in skip_field_paths:
+                values: Sequence[PrimitiveJSON] | ToJsonTemplate
                 operator, values = one(filter.items())
+                original_values = values
                 # Note that `is_not` is only used internally (for filtering by
                 # inaccessible sources)
                 if operator in ('is', 'is_not'):
                     field_type = self.service.field_type(self.catalog, field_path)
+                    if field_path == source_id_field_path:
+                        values = ToJsonTemplate(param_name=source_id_field_name, value=values)
                     if isinstance(field_type, Nested):
                         term_queries = []
-                        for nested_field, nested_value in one(values).items():
+                        for nested_field, nested_value in one(original_values).items():
                             nested_body = {dotted(field_path, nested_field, 'keyword'): nested_value}
                             term_queries.append(Q('term', **nested_body))
                         query = Q('nested', path=dotted(field_path), query=Q('bool', must=term_queries))
                     else:
                         query = Q('terms', **{dotted(field_path, 'keyword'): values})
                         translated_none = field_type.to_index(None)
-                        if translated_none in values:
+                        if translated_none in original_values:
                             # Note that at this point None values in filters have already
                             # been translated e.g. {'is': ['~null']} and if the filter has a
                             # None our query needs to find fields with None values as well
@@ -668,6 +677,29 @@ def to_source(self) -> AnyJSON:
         raise NotImplementedError
 
 
+@attr.s(frozen=True, auto_attribs=True, kw_only=True)
+class ToJsonTemplate(Template):
+
+    def to_source(self):
+        return RawStr('{{#toJson}}' + self.param_name + '{{/toJson}}')
+
+
+class RawStr(str):
+    """
+    Instances of this class will not be surrounded by quotes when encoded as
+    JSON using a :class:`TemplateSearchJSONEncoder`.
+    """
+
+
+_original = json.encoder.py_encode_basestring_ascii
+
+
+def _encode_basestring_ascii(s: str) -> str:
+    result = _original(s)
+    assert result[0] == result[-1] == '"', result
+    return result[1:-1] if isinstance(s, RawStr) else result
+
+
 class TemplateSearchJSONEncoder(json.JSONEncoder):
 
     def __init__(self, **kwargs):
@@ -686,6 +718,11 @@ def default(self, obj):
         else:
             return super().default(obj)
 
+    def iterencode(self, o: AnyJSON, _one_shot: bool = False) -> Iterator[str]:
+        with unittest.mock.patch('json.encoder.encode_basestring_ascii',
+                                 wraps=_encode_basestring_ascii):
+            return super().iterencode(o, _one_shot=_one_shot)
+
 
 class TemplateSearch(Search):
 

test/service/test_request_builder.py

@@ -100,7 +100,7 @@ def facets(self) -> Sequence[str]:
         'constant_score': {
             'filter': {
                 'terms': {
-                    'sources.id.keyword': []
+                    'sources.id.keyword': '{{#toJson}}sourceId{{/toJson}}'
                 }
             }
         }
@@ -336,8 +336,12 @@ def _test_create_request_with_service(self,
         filters = Filters(explicit=sample_filter, source_ids=set())
         request = self._prepare_request(filters, post_filter, service)
         expected_output = {
-            'source': json.dumps(expected_output, sort_keys=True),
-            'params': {}
+            'source': (
+                json.dumps(expected_output, sort_keys=True)
+                .replace('"{{#toJson}}', '{{#toJson}}')
+                .replace('{{/toJson}}"', '{{/toJson}}')
+            ),
+            'params': {'sourceId': []}
         }
         actual_request_body = request.to_dict()
         self.assertEqual(expected_output, actual_request_body)