feat(search): Add flag to enable caching on search service

metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java

@@ -25,13 +25,15 @@ public class SearchService {
   private final AllEntitiesSearchAggregatorCache _allEntitiesSearchAggregatorCache;
 
   public SearchService(EntityRegistry entityRegistry, EntitySearchService entitySearchService,
-      SearchRanker searchRanker, CacheManager cacheManager, int batchSize) {
+      SearchRanker searchRanker, CacheManager cacheManager, int batchSize, boolean enableCache) {
     _entitySearchService = entitySearchService;
     _searchRanker = searchRanker;
     _aggregator =
-        new AllEntitiesSearchAggregator(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize);
-    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize);
-    _allEntitiesSearchAggregatorCache = new AllEntitiesSearchAggregatorCache(cacheManager, _aggregator, batchSize);
+        new AllEntitiesSearchAggregator(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize,
+            enableCache);
+    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize, enableCache);
+    _allEntitiesSearchAggregatorCache =
+        new AllEntitiesSearchAggregatorCache(cacheManager, _aggregator, batchSize, enableCache);
   }
 
   /**

metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java

@@ -47,13 +47,13 @@ public class AllEntitiesSearchAggregator {
   private final EntitySearchServiceCache _entitySearchServiceCache;
 
   public AllEntitiesSearchAggregator(EntityRegistry entityRegistry, EntitySearchService entitySearchService,
-      SearchRanker searchRanker, CacheManager cacheManager, int batchSize) {
+      SearchRanker searchRanker, CacheManager cacheManager, int batchSize, boolean enableCache) {
     _entityRegistry = entityRegistry;
     _entitySearchService = entitySearchService;
     _searchRanker = searchRanker;
     _cacheManager = cacheManager;
     _nonEmptyEntitiesCache = new NonEmptyEntitiesCache(entityRegistry, entitySearchService, cacheManager);
-    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize);
+    _entitySearchServiceCache = new EntitySearchServiceCache(cacheManager, entitySearchService, batchSize, enableCache);
   }
 
   @Nonnull

metadata-io/src/main/java/com/linkedin/metadata/search/cache/AllEntitiesSearchAggregatorCache.java

@@ -19,12 +19,13 @@ public class AllEntitiesSearchAggregatorCache {
   private final CacheManager cacheManager;
   private final AllEntitiesSearchAggregator aggregator;
   private final int batchSize;
+  private final boolean enableCache;
 
   public CacheableSearcher<?> getSearcher(List<String> entities, @Nonnull String input, @Nullable Filter postFilters,
       @Nullable SortCriterion sortCriterion, @Nullable SearchFlags searchFlags) {
     return new CacheableSearcher<>(cacheManager.getCache(ALL_ENTITIES_SEARCH_AGGREGATOR_CACHE_NAME), batchSize,
         querySize -> aggregator.search(entities, input, postFilters, sortCriterion, querySize.getFrom(),
             querySize.getSize(), searchFlags),
-        querySize -> Quintet.with(entities, input, postFilters, sortCriterion, querySize), searchFlags);
+        querySize -> Quintet.with(entities, input, postFilters, sortCriterion, querySize), searchFlags, enableCache);
   }
 }

metadata-io/src/main/java/com/linkedin/metadata/search/cache/CacheableSearcher.java

@@ -28,6 +28,7 @@ public class CacheableSearcher<K> {
   private final Function<QueryPagination, K> cacheKeyGenerator;
   @Nullable
   private final SearchFlags searchFlags;
+  private final boolean enableCache;
 
   @Value
   public static class QueryPagination {
@@ -95,6 +96,6 @@ private SearchResult getBatch(int batchId) {
   }
 
   private boolean enableCache() {
-    return searchFlags == null || !searchFlags.isSkipCache();
+    return enableCache && (searchFlags == null || !searchFlags.isSkipCache());
   }
 }

metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntitySearchServiceCache.java

@@ -18,12 +18,13 @@ public class EntitySearchServiceCache {
   private final CacheManager cacheManager;
   private final EntitySearchService entitySearchService;
   private final int batchSize;
+  private final boolean enableCache;
 
   public CacheableSearcher<?> getSearcher(@Nonnull String entityName, @Nonnull String input,
       @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, @Nullable SearchFlags searchFlags) {
     return new CacheableSearcher<>(cacheManager.getCache(ENTITY_SEARCH_SERVICE_CACHE_NAME), batchSize,
         querySize -> entitySearchService.search(entityName, input, postFilters, sortCriterion, querySize.getFrom(),
             querySize.getSize()), querySize -> Quintet.with(entityName, input, postFilters, sortCriterion, querySize),
-        searchFlags);
+        searchFlags, enableCache);
   }
 }

metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java

@@ -83,7 +83,7 @@ public void setup() {
     _cacheManager = new ConcurrentMapCacheManager();
     _graphService = mock(GraphService.class);
     _lineageSearchService = new LineageSearchService(
-        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100),
+        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100, true),
         _graphService, _cacheManager.getCache("test"));
   }
 

metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTest.java

@@ -60,7 +60,8 @@ public void setup() {
     _elasticSearchService = buildEntitySearchService();
     _elasticSearchService.configure();
     _cacheManager = new ConcurrentMapCacheManager();
-    _searchService = new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100);
+    _searchService =
+        new SearchService(_entityRegistry, _elasticSearchService, new SimpleRanker(), _cacheManager, 100, true);
   }
 
   @BeforeMethod

metadata-io/src/test/java/com/linkedin/metadata/search/cache/CacheableSearcherTest.java

@@ -26,7 +26,7 @@ public class CacheableSearcherTest {
   public void testCacheableSearcherWhenEmpty() {
     CacheableSearcher<Integer> emptySearcher =
         new CacheableSearcher<>(cacheManager.getCache("emptySearcher"), 10, this::getEmptySearchResult,
-            CacheableSearcher.QueryPagination::getFrom, null);
+            CacheableSearcher.QueryPagination::getFrom, null, true);
     assertTrue(emptySearcher.getSearchResults(0, 0).getEntities().isEmpty());
     assertTrue(emptySearcher.getSearchResults(0, 10).getEntities().isEmpty());
     assertTrue(emptySearcher.getSearchResults(5, 10).getEntities().isEmpty());
@@ -36,7 +36,7 @@ public void testCacheableSearcherWhenEmpty() {
   public void testCacheableSearcherWithFixedNumResults() {
     CacheableSearcher<Integer> fixedBatchSearcher =
         new CacheableSearcher<>(cacheManager.getCache("fixedBatchSearcher"), 10, qs -> getSearchResult(qs, 10),
-            CacheableSearcher.QueryPagination::getFrom, null);
+            CacheableSearcher.QueryPagination::getFrom, null, true);
 
     SearchResult result = fixedBatchSearcher.getSearchResults(0, 0);
     assertTrue(result.getEntities().isEmpty());
@@ -59,7 +59,8 @@ public void testCacheableSearcherWithFixedNumResults() {
   public void testCacheableSearcherWithVariableNumResults() {
     CacheableSearcher<Integer> variableBatchSearcher =
         new CacheableSearcher<>(cacheManager.getCache("variableBatchSearcher"), 10,
-            qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, null);
+            qs -> getSearchResult(qs, qs.getFrom() + qs.getSize()), CacheableSearcher.QueryPagination::getFrom, null,
+            true);
 
     SearchResult result = variableBatchSearcher.getSearchResults(0, 0);
     assertTrue(result.getEntities().isEmpty());

metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java

@@ -38,10 +38,13 @@ public class SearchServiceFactory {
   @Value("${searchService.resultBatchSize}")
   private Integer batchSize;
 
+  @Value("${searchService.enableCache}")
+  private Boolean enableCache;
+
   @Bean(name = "searchService")
   @Primary
   @Nonnull
   protected SearchService getInstance() {
-    return new SearchService(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize);
+    return new SearchService(entityRegistry, entitySearchService, searchRanker, cacheManager, batchSize, enableCache);
   }
 }

metadata-service/factories/src/main/resources/application.yml

@@ -58,6 +58,7 @@ graphService:
 
 searchService:
   resultBatchSize: ${SEARCH_SERVICE_BATCH_SIZE:100}
+  enableCache: ${SEARCH_SERVICE_ENABLE_CACHE:false}
 
 configEntityRegistry:
   # TODO: Change to read from resources on classpath.