fix(gms): Write back lineage search results to in-memory cache bound to feature flag

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/featureflags/FeatureFlags.java

@@ -6,4 +6,5 @@
 @Data
 public class FeatureFlags {
   private boolean showSimplifiedHomepageByDefault = false;
+  private boolean lineageSearchCacheEnabled = false;
 }

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

@@ -39,7 +39,9 @@
 public class LineageSearchService {
   private final SearchService _searchService;
   private final GraphService _graphService;
+  @Nullable
   private final Cache cache;
+  private final boolean cacheEnabled;
 
   private static final String DEGREE_FILTER = "degree";
   private static final String DEGREE_FILTER_INPUT = "degree.keyword";
@@ -71,10 +73,14 @@ public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull
       @Nonnull List<String> entities, @Nullable String input, @Nullable Integer maxHops, @Nullable Filter inputFilters,
       @Nullable SortCriterion sortCriterion, int from, int size) {
     // Cache multihop result for faster performance
-    EntityLineageResult lineageResult = cache.get(Pair.of(sourceUrn, direction), EntityLineageResult.class);
+    EntityLineageResult lineageResult = cacheEnabled ? cache.get(Pair.of(sourceUrn, direction), EntityLineageResult.class)
+        : null;
     if (lineageResult == null) {
       maxHops = maxHops != null ? maxHops : 1000;
       lineageResult = _graphService.getLineage(sourceUrn, direction, 0, MAX_RELATIONSHIPS, maxHops);
+      if (cacheEnabled) {
+        cache.put(Pair.of(sourceUrn, direction), lineageResult);
+      }
     }
 
     // Filter hopped result based on the set of entities to return and inputFilters before sending to search

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

@@ -87,10 +87,10 @@ public void setup() {
     _elasticSearchService.configure();
     _cacheManager = new ConcurrentMapCacheManager();
     _graphService = mock(GraphService.class);
-    resetService();
+    resetService(true);
   }
 
-  private void resetService() {
+  private void resetService(boolean withCache) {
     CachingEntitySearchService cachingEntitySearchService = new CachingEntitySearchService(_cacheManager, _elasticSearchService, 100, true);
     _lineageSearchService = new LineageSearchService(
         new SearchService(
@@ -102,7 +102,7 @@ private void resetService() {
                 100,
                 true),
             new SimpleRanker()),
-        _graphService, _cacheManager.getCache("test"));
+        _graphService, _cacheManager.getCache("test"), withCache);
   }
 
   @BeforeMethod
@@ -126,7 +126,7 @@ private ElasticSearchService buildEntitySearchService() {
 
   private void clearCache() {
     _cacheManager.getCacheNames().forEach(cache -> _cacheManager.getCache(cache).clear());
-    resetService();
+    resetService(true);
   }
 
   @AfterClass

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

@@ -1,6 +1,7 @@
 package com.linkedin.gms.factory.search;
 
 import com.linkedin.gms.factory.common.GraphServiceFactory;
+import com.linkedin.gms.factory.config.ConfigurationProvider;
 import com.linkedin.gms.factory.spring.YamlPropertySourceFactory;
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.search.LineageSearchService;
@@ -21,22 +22,13 @@
 @PropertySource(value = "classpath:/application.yml", factory = YamlPropertySourceFactory.class)
 public class LineageSearchServiceFactory {
 
-  @Autowired
-  @Qualifier("searchService")
-  private SearchService searchService;
-
-  @Autowired
-  @Qualifier("graphService")
-  private GraphService graphService;
-
-  @Autowired
-  private CacheManager cacheManager;
-
   @Bean(name = "relationshipSearchService")
   @Primary
   @Nonnull
-  protected LineageSearchService getInstance() {
+  protected LineageSearchService getInstance(CacheManager cacheManager, GraphService graphService,
+       SearchService searchService, ConfigurationProvider configurationProvider) {
+    boolean cacheEnabled = configurationProvider.getFeatureFlags().isLineageSearchCacheEnabled();
     return new LineageSearchService(searchService, graphService,
-        cacheManager.getCache("relationshipSearchService"));
+        cacheEnabled ? cacheManager.getCache("relationshipSearchService") : null, cacheEnabled);
   }
 }

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

@@ -203,6 +203,7 @@ bootstrap:
 
 featureFlags:
   showSimplifiedHomepageByDefault: ${SHOW_SIMPLIFIED_HOMEPAGE_BY_DEFAULT:false} # shows a simplified homepage with just datasets, charts and dashboards by default to users. this can be configured in user settings
+  lineageSearchCacheEnabled: ${LINEAGE_SEARCH_CACHE_ENABLED:false} # Enables in-memory cache for searchAcrossLineage query, disabled by default to prevent unexpected update delays
 
 entityChangeEvents:
   enabled: ${ENABLE_ENTITY_CHANGE_EVENTS_HOOK:true}