INTR-341 Fix search breaking during deployment

src/core/tasks.py

@@ -30,9 +30,6 @@ def ingest_uk_staff_locations(self):
 @celery_app.task(bind=True)
 @cache_lock(cache_key="update_search_index")
 def update_search_index(self):
-    # Run update_index --schema-only
-    call_command("update_index", schema_only=True)
-
     # Run update_index
     call_command("update_index")
 

src/extended_search/apps.py

@@ -8,15 +8,10 @@ class ExtendedSearchConfig(AppConfig):
     def ready(self):
         import extended_search.signals  # noqa
         from extended_search import query_builder, settings
-        from extended_search.index import get_indexed_models
 
         settings.settings_singleton.initialise_field_dict()
         settings.settings_singleton.initialise_env_dict()
         if django_settings.APP_ENV not in ["test", "build"]:
             settings.settings_singleton.initialise_db_dict()
         settings.extended_search_settings = settings.settings_singleton.to_dict()
         query_builder.extended_search_settings = settings.extended_search_settings
-
-        for model_class in get_indexed_models():
-            if hasattr(model_class, "indexed_fields") and model_class.indexed_fields:
-                query_builder.CustomQueryBuilder.build_search_query(model_class, True)

src/extended_search/backends/backend.py

@@ -1,17 +1,21 @@
 from typing import Optional, Union
 
+from wagtail.search.backends import get_search_backend
 from wagtail.search.backends.elasticsearch7 import (
     Elasticsearch7Mapping,
     Elasticsearch7SearchBackend,
     Elasticsearch7SearchQueryCompiler,
+    ElasticsearchAtomicIndexRebuilder,
     Field,
 )
 from wagtail.search.index import SearchField
+from wagtail.search.management.commands.update_index import group_models_by_index
 from wagtail.search.query import MATCH_NONE, Fuzzy, MatchAll, Not, Phrase, PlainText
 
 from extended_search import settings as search_settings
-from extended_search.index import RelatedFields
+from extended_search.index import RelatedFields, get_indexed_models
 from extended_search.query import Filtered, FunctionScore, Nested, OnlyFields
+from extended_search.query_builder import build_queries
 
 
 class FilteredSearchMapping(Elasticsearch7Mapping):
@@ -390,9 +394,29 @@ class CustomSearchQueryCompiler(
     mapping_class = CustomSearchMapping
 
 
+class CustomAtomicIndexRebuilder(ElasticsearchAtomicIndexRebuilder):
+    def start(self):
+        index = super().start()
+
+        models_grouped_by_index = group_models_by_index(
+            get_search_backend(), get_indexed_models()
+        )
+        models_for_current_index = []
+
+        for index_models in models_grouped_by_index.keys():
+            if index.name.startswith(index_models.name):
+                models_for_current_index = models_grouped_by_index[index_models]
+
+        if models_for_current_index:
+            build_queries(models=models_for_current_index)
+
+        return index
+
+
 class CustomSearchBackend(Elasticsearch7SearchBackend):
     query_compiler_class = CustomSearchQueryCompiler
     mapping_class = CustomSearchMapping
+    atomic_rebuilder_class = CustomAtomicIndexRebuilder
 
 
 SearchBackend = CustomSearchBackend

src/extended_search/query_builder.py

@@ -1,13 +1,15 @@
 import inspect
 import logging
-from typing import Optional, Type
+from typing import TYPE_CHECKING, Optional, Type
 
 from django.conf import settings
 from django.core.cache import cache
 from django.db import models
 from wagtail.search import index
+from wagtail.search.backends import get_search_backend
 from wagtail.search.query import Boost, Fuzzy, Phrase, PlainText, SearchQuery
 
+from extended_search import query_builder
 from extended_search import settings as search_settings
 from extended_search.index import (
     BaseField,
@@ -23,6 +25,9 @@
 from extended_search.types import AnalysisType, SearchQueryType
 
 
+if TYPE_CHECKING:
+    from extended_search.backends.backend import CustomSearchBackend
+
 logger = logging.getLogger(__name__)
 
 
@@ -332,7 +337,9 @@ def build_search_query(cls, model_class, ignore_cache=False):
         type, and all are joined together at the end.
         """
         if settings.SEARCH_ENABLE_QUERY_CACHE:
-            cache_key = model_class.__name__
+            search_backend: "CustomSearchBackend" = get_search_backend()
+            model_index = search_backend.get_index_for_model(model_class)
+            cache_key = f"{model_index.name}__{model_class.__name__}"
             if not ignore_cache:
                 built_query = cache.get(cache_key, None)
                 if built_query:
@@ -410,3 +417,12 @@ def get_extended_models_with_unique_indexed_fields(
             ):
                 extended_model_classes.append(indexed_model)
         return extended_model_classes
+
+
+def build_queries(models: list[models.Model] | None = None):
+    if not models:
+        models = get_indexed_models()
+
+    for model_class in models:
+        if hasattr(model_class, "indexed_fields") and model_class.indexed_fields:
+            query_builder.CustomQueryBuilder.build_search_query(model_class, True)