add caching to all search apis

Author: Deepthi-Chand

api/views/search_collaborative.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import structlog
+from django.core.cache import cache
 from elasticsearch_dsl import A
 from elasticsearch_dsl import Q as ESQ
 from elasticsearch_dsl import Search
@@ -14,6 +15,8 @@
 from api.views.paginated_elastic_view import PaginatedElasticSearchAPIView
 from search.documents import CollaborativeDocument
 
+METADATA_CACHE_TTL = 60 * 30  # 30 minutes
+
 logger = structlog.get_logger(__name__)
 
 
@@ -147,6 +150,12 @@ def __init__(self, **kwargs: Any) -> None:
         attributes={"component": "search_collaborative"},
     )
     def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
+        cached: Optional[Tuple[List[str], Dict[str, str]]] = cache.get(
+            "collaborative_search_metadata_config"
+        )
+        if cached:
+            return cached
+
         searchable_fields = [
             "title",
             "summary",
@@ -173,7 +182,9 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         for metadata in filterable_metadata:
             aggregations[f"metadata.{metadata.label}"] = "terms"  # type: ignore
 
-        return searchable_fields, aggregations
+        result = (searchable_fields, aggregations)
+        cache.set("collaborative_search_metadata_config", result, timeout=METADATA_CACHE_TTL)
+        return result
 
     @trace_method(name="add_aggregations", attributes={"component": "search_collaborative"})
     def add_aggregations(self, search: Search) -> Search:
@@ -189,8 +200,17 @@ def add_aggregations(self, search: Search) -> Search:
                 )
 
         if aggregate_fields:
-            metadata_qs = Metadata.objects.filter(filterable=True)
-            filterable_metadata = [str(meta.label) for meta in metadata_qs]  # type: ignore
+            filterable_metadata: List[str] = (
+                cache.get("collaborative_filterable_metadata_labels") or []
+            )
+            if not filterable_metadata:
+                metadata_qs = Metadata.objects.filter(filterable=True)
+                filterable_metadata = [str(meta.label) for meta in metadata_qs]  # type: ignore
+                cache.set(
+                    "collaborative_filterable_metadata_labels",
+                    filterable_metadata,
+                    timeout=METADATA_CACHE_TTL,
+                )
 
             metadata_bucket = search.aggs.bucket("metadata", "nested", path="metadata")
             composite_agg = A(
@@ -277,8 +297,15 @@ def generate_q_expression(self, query: str) -> Optional[Union[ESQuery, List[ESQu
 
     @trace_method(name="add_filters", attributes={"component": "search_collaborative"})
     def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
-        non_filter_metadata = Metadata.objects.filter(filterable=False).all()
-        excluded_labels: List[str] = [e.label for e in non_filter_metadata]  # type: ignore
+        excluded_labels: List[str] = cache.get("collaborative_non_filter_metadata_labels") or []
+        if not excluded_labels:
+            non_filter_metadata = Metadata.objects.filter(filterable=False).all()
+            excluded_labels = [e.label for e in non_filter_metadata]  # type: ignore
+            cache.set(
+                "collaborative_non_filter_metadata_labels",
+                excluded_labels,
+                timeout=METADATA_CACHE_TTL,
+            )
 
         for filter in filters:
             if filter in excluded_labels:

api/views/search_dataset.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union, cast
 
 import structlog
+from django.core.cache import cache
 from elasticsearch_dsl import A
 from elasticsearch_dsl import Q as ESQ
 from elasticsearch_dsl import Search
@@ -14,6 +15,8 @@
 from api.views.paginated_elastic_view import PaginatedElasticSearchAPIView
 from search.documents import DatasetDocument
 
+METADATA_CACHE_TTL = 60 * 30  # 30 minutes
+
 logger = structlog.get_logger(__name__)
 
 
@@ -147,6 +150,12 @@ def __init__(self, **kwargs: Any) -> None:
     )
     def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         """Get searchable fields and aggregations for the search."""
+        cached: Optional[Tuple[List[str], Dict[str, str]]] = cache.get(
+            "dataset_search_metadata_config"
+        )
+        if cached:
+            return cached
+
         enabled_metadata = Metadata.objects.filter(enabled=True).all()
         searchable_fields: List[str] = []
         searchable_fields.extend(
@@ -170,7 +179,9 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         for metadata in enabled_metadata:  # type: Metadata
             if metadata.filterable:
                 aggregations[f"metadata.{metadata.label}"] = "terms"
-        return searchable_fields, aggregations
+        result = (searchable_fields, aggregations)
+        cache.set("dataset_search_metadata_config", result, timeout=METADATA_CACHE_TTL)
+        return result
 
     @trace_method(name="add_aggregations", attributes={"component": "search_dataset"})
     def add_aggregations(self, search: Search) -> Search:
@@ -248,8 +259,13 @@ def generate_q_expression(self, query: str) -> Optional[Union[ESQuery, List[ESQu
     @trace_method(name="add_filters", attributes={"component": "search_dataset"})
     def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
         """Add filters to the search query."""
-        non_filter_metadata = Metadata.objects.filter(filterable=False).all()
-        excluded_labels: List[str] = [e.label for e in non_filter_metadata]  # type: ignore
+        excluded_labels: List[str] = cache.get("dataset_non_filter_metadata_labels") or []
+        if not excluded_labels:
+            non_filter_metadata = Metadata.objects.filter(filterable=False).all()
+            excluded_labels = [e.label for e in non_filter_metadata]  # type: ignore
+            cache.set(
+                "dataset_non_filter_metadata_labels", excluded_labels, timeout=METADATA_CACHE_TTL
+            )
 
         for filter in filters:
             if filter in excluded_labels:

api/views/search_unified.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Tuple
 
 import structlog
+from django.core.cache import cache
 from elasticsearch_dsl import Q as ESQ
 from elasticsearch_dsl import Search
 from rest_framework import serializers
@@ -13,6 +14,7 @@
 from api.models import Dataset, Geography, Metadata, UseCase
 from api.models.AIModel import AIModel
 from api.models.Collaborative import Collaborative
+from api.signals.dataset_signals import SEARCH_CACHE_VERSION_KEY
 from api.utils.telemetry_utils import trace_method
 from DataSpace import settings
 from search.documents import (
@@ -440,10 +442,27 @@ def perform_unified_search(
 
         return results, total, aggregations
 
+    def _generate_unified_cache_key(self, request: Any) -> str:
+        """Generate a unique cache key for unified search based on request parameters."""
+        params: Dict[str, str] = {
+            "query": request.GET.get("query", ""),
+            "page": request.GET.get("page", "1"),
+            "size": request.GET.get("size", "10"),
+            "types": request.GET.get("types", "dataset,usecase,aimodel,collaborative,publisher"),
+            "filters": str(sorted(request.GET.dict().items())),
+            "version": str(cache.get(SEARCH_CACHE_VERSION_KEY, 0)),
+        }
+        return f"unified_search:{hash(frozenset(params.items()))}"
+
     @trace_method(name="get", attributes={"component": "unified_search"})
     def get(self, request: Any) -> Response:
         """Handle GET request and return unified search results."""
         try:
+            cache_key = self._generate_unified_cache_key(request)
+            cached_result = cache.get(cache_key)
+            if cached_result:
+                return Response(cached_result)
+
             query: str = request.GET.get("query", "")
             page: int = int(request.GET.get("page", 1))
             size: int = int(request.GET.get("size", 10))
@@ -476,6 +495,8 @@ def get(self, request: Any) -> Response:
                 "types_searched": types_list,
             }
 
+            cache.set(cache_key, result, timeout=3600)
+
             return Response(result)
 
         except Exception as e:

api/views/search_usecase.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union, cast
 
 import structlog
+from django.core.cache import cache
 from elasticsearch_dsl import A
 from elasticsearch_dsl import Q as ESQ
 from elasticsearch_dsl import Search
@@ -14,6 +15,8 @@
 from api.views.paginated_elastic_view import PaginatedElasticSearchAPIView
 from search.documents import UseCaseDocument
 
+METADATA_CACHE_TTL = 60 * 30  # 30 minutes
+
 logger = structlog.get_logger(__name__)
 
 
@@ -144,9 +147,7 @@ def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.searchable_fields: List[str]
         self.aggregations: Dict[str, str]
-        self.searchable_fields, self.aggregations = (
-            self.get_searchable_and_aggregations()
-        )
+        self.searchable_fields, self.aggregations = self.get_searchable_and_aggregations()
         self.logger = structlog.get_logger(__name__)
 
     @trace_method(
@@ -155,6 +156,12 @@ def __init__(self, **kwargs: Any) -> None:
     )
     def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         """Get searchable fields and aggregations for the search."""
+        cached: Optional[Tuple[List[str], Dict[str, str]]] = cache.get(
+            "usecase_search_metadata_config"
+        )
+        if cached:
+            return cached
+
         searchable_fields = [
             "title",
             "summary",
@@ -181,7 +188,9 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
         for metadata in filterable_metadata:
             aggregations[f"metadata.{metadata.label}"] = "terms"  # type: ignore
 
-        return searchable_fields, aggregations
+        result = (searchable_fields, aggregations)
+        cache.set("usecase_search_metadata_config", result, timeout=METADATA_CACHE_TTL)
+        return result
 
     @trace_method(name="add_aggregations", attributes={"component": "search_usecase"})
     def add_aggregations(self, search: Search) -> Search:
@@ -199,18 +208,21 @@ def add_aggregations(self, search: Search) -> Search:
                 )
 
         if aggregate_fields:
-            metadata_qs = Metadata.objects.filter(filterable=True)
-            filterable_metadata = [str(meta.label) for meta in metadata_qs]  # type: ignore
+            filterable_metadata: List[str] = cache.get("usecase_filterable_metadata_labels") or []
+            if not filterable_metadata:
+                metadata_qs = Metadata.objects.filter(filterable=True)
+                filterable_metadata = [str(meta.label) for meta in metadata_qs]  # type: ignore
+                cache.set(
+                    "usecase_filterable_metadata_labels",
+                    filterable_metadata,
+                    timeout=METADATA_CACHE_TTL,
+                )
 
             metadata_bucket = search.aggs.bucket("metadata", "nested", path="metadata")
             composite_agg = A(
                 "composite",
                 sources=[
-                    {
-                        "metadata_label": {
-                            "terms": {"field": "metadata.metadata_item.label"}
-                        }
-                    },
+                    {"metadata_label": {"terms": {"field": "metadata.metadata_item.label"}}},
                     {"metadata_value": {"terms": {"field": "metadata.value"}}},
                 ],
                 size=10000,
@@ -219,13 +231,7 @@ def add_aggregations(self, search: Search) -> Search:
                 "filter",
                 {  # type: ignore[arg-type]
                     "bool": {
-                        "must": [
-                            {
-                                "terms": {
-                                    "metadata.metadata_item.label": filterable_metadata
-                                }
-                            }
-                        ]
+                        "must": [{"terms": {"metadata.metadata_item.label": filterable_metadata}}]
                     }
                 },
             )
@@ -235,12 +241,8 @@ def add_aggregations(self, search: Search) -> Search:
 
         return search
 
-    @trace_method(
-        name="generate_q_expression", attributes={"component": "search_usecase"}
-    )
-    def generate_q_expression(
-        self, query: str
-    ) -> Optional[Union[ESQuery, List[ESQuery]]]:
+    @trace_method(name="generate_q_expression", attributes={"component": "search_usecase"})
+    def generate_q_expression(self, query: str) -> Optional[Union[ESQuery, List[ESQuery]]]:
         """Generate Elasticsearch Query expression."""
         if query:
             queries: List[ESQuery] = []
@@ -256,9 +258,7 @@ def generate_q_expression(
                                     ESQ("wildcard", **{field: {"value": f"*{query}*"}}),
                                     ESQ(
                                         "fuzzy",
-                                        **{
-                                            field: {"value": query, "fuzziness": "AUTO"}
-                                        },
+                                        **{field: {"value": query, "fuzziness": "AUTO"}},
                                     ),
                                 ],
                             ),
@@ -281,18 +281,14 @@ def generate_q_expression(
                                     ESQ("wildcard", **{field: {"value": f"*{query}*"}}),
                                     ESQ(
                                         "fuzzy",
-                                        **{
-                                            field: {"value": query, "fuzziness": "AUTO"}
-                                        },
+                                        **{field: {"value": query, "fuzziness": "AUTO"}},
                                     ),
                                 ],
                             ),
                         )
                     )
                 else:
-                    queries.append(
-                        ESQ("fuzzy", **{field: {"value": query, "fuzziness": "AUTO"}})
-                    )
+                    queries.append(ESQ("fuzzy", **{field: {"value": query, "fuzziness": "AUTO"}}))
         else:
             queries = [ESQ("match_all")]
 
@@ -301,8 +297,13 @@ def generate_q_expression(
     @trace_method(name="add_filters", attributes={"component": "search_usecase"})
     def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
         """Add filters to the search query."""
-        non_filter_metadata = Metadata.objects.filter(filterable=False).all()
-        excluded_labels: List[str] = [e.label for e in non_filter_metadata]  # type: ignore
+        excluded_labels: List[str] = cache.get("usecase_non_filter_metadata_labels") or []
+        if not excluded_labels:
+            non_filter_metadata = Metadata.objects.filter(filterable=False).all()
+            excluded_labels = [e.label for e in non_filter_metadata]  # type: ignore
+            cache.set(
+                "usecase_non_filter_metadata_labels", excluded_labels, timeout=METADATA_CACHE_TTL
+            )
 
         for filter in filters:
             if filter in excluded_labels:
@@ -335,9 +336,7 @@ def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
                 search = search.filter(
                     "nested",
                     path="metadata",
-                    query={
-                        "bool": {"must": {"term": {f"metadata.value": filters[filter]}}}
-                    },
+                    query={"bool": {"must": {"term": {f"metadata.value": filters[filter]}}}},
                 )
         return search