Add Invalidation for Lineage Cache during updates (#27606)

* Add Invalidation for Lineage Cache during updates

* Add test

(cherry picked from commit 557626f)

Author: mohityadav766

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/LineageRepository.java

@@ -393,6 +393,16 @@ private void addLineageToSearch(
         buildEntityLineageData(fromEntity, toEntity, lineageDetails).withToEntity(null);
     Pair<String, String> to = new ImmutablePair<>("_id", toEntity.getId().toString());
     searchClient.updateLineage(destinationIndexName, to, lineageData);
+    invalidateLineageCacheForEdge(fromEntity, toEntity);
+  }
+
+  private void invalidateLineageCacheForEdge(EntityReference from, EntityReference to) {
+    if (from != null) {
+      searchClient.invalidateLineageCache(from.getFullyQualifiedName());
+    }
+    if (to != null) {
+      searchClient.invalidateLineageCache(to.getFullyQualifiedName());
+    }
   }
 
   public static RelationshipRef buildEntityRefLineage(EntityReference entityRef) {
@@ -1187,12 +1197,28 @@ private void deleteLineageFromSearch(List<CollectionDAO.EntityRelationshipObject
     for (CollectionDAO.EntityRelationshipObject obj : relations) {
       LineageDetails lineageDetails = JsonUtils.readValue(obj.getJson(), LineageDetails.class);
       deleteLineageFromSearch(
-          new EntityReference().withId(UUID.fromString(obj.getFromId())),
-          new EntityReference().withId(UUID.fromString(obj.getToId())),
+          resolveRefForCacheInvalidation(obj.getFromEntity(), obj.getFromId()),
+          resolveRefForCacheInvalidation(obj.getToEntity(), obj.getToId()),
           lineageDetails);
     }
   }
 
+  private EntityReference resolveRefForCacheInvalidation(String entityType, String id) {
+    EntityReference ref = new EntityReference().withId(UUID.fromString(id));
+    if (nullOrEmpty(entityType)) {
+      return ref;
+    }
+    try {
+      EntityReference resolved =
+          Entity.getEntityReferenceById(entityType, UUID.fromString(id), Include.ALL);
+      return ref.withType(entityType).withFullyQualifiedName(resolved.getFullyQualifiedName());
+    } catch (Exception e) {
+      LOG.debug(
+          "Could not resolve FQN for {}:{} during lineage cache invalidation", entityType, id);
+      return ref.withType(entityType);
+    }
+  }
+
   private void deleteLineageFromSearch(
       EntityReference fromEntity, EntityReference toEntity, LineageDetails lineageDetails) {
     String uniqueValue = getDocumentUniqueId(fromEntity, toEntity);
@@ -1202,6 +1228,7 @@ private void deleteLineageFromSearch(
           new ImmutablePair<>("upstreamLineage.docUniqueId.keyword", uniqueValue),
           new ImmutablePair<>(
               REMOVE_LINEAGE_SCRIPT, Collections.singletonMap("docUniqueId", uniqueValue)));
+      invalidateLineageCacheForEdge(fromEntity, toEntity);
     } catch (Exception e) {
       SearchIndexRetryQueue.enqueue(
           fromEntity.getId() != null ? fromEntity.getId().toString() : null,

openmetadata-service/src/main/java/org/openmetadata/service/search/SearchClient.java

@@ -718,4 +718,15 @@ default void initializeLineageBuilders() {
     // Default implementation does nothing - concrete implementations can override
     // This allows backward compatibility for clients that don't need lineage features
   }
+
+  /**
+   * Evicts every cached lineage graph whose root, nodes, or edge endpoints reference
+   * the given FQN. Callers invoke this after a lineage edge involving the FQN is added
+   * or deleted so stale graphs are not served back to the UI.
+   *
+   * @param fqn Fully qualified name of the entity touched by the mutation
+   */
+  default void invalidateLineageCache(String fqn) {
+    // Default no-op; concrete clients delegate to their LineageGraphBuilder cache
+  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java

@@ -425,6 +425,14 @@ public SearchLineageResult searchPlatformLineage(
     return lineageGraphBuilder.getPlatformLineage(index, queryFilter, deleted);
   }
 
+  @Override
+  public void invalidateLineageCache(String fqn) {
+    if (lineageGraphBuilder == null) {
+      return;
+    }
+    lineageGraphBuilder.invalidateLineageCacheForFqn(fqn);
+  }
+
   @Override
   public Response searchEntityRelationship(
       String fqn, int upstreamDepth, int downstreamDepth, String queryFilter, boolean deleted)

openmetadata-service/src/main/java/org/openmetadata/service/search/lineage/AbstractLineageGraphBuilder.java

@@ -183,15 +183,26 @@ protected void cacheResult(SearchLineageRequest request, SearchLineageResult res
    * Should be called when entity is updated or lineage edges change.
    */
   protected void invalidateCache(String fqn) {
-    if (!config.isEnableCaching()) {
+    invalidateLineageCacheForFqn(fqn);
+  }
+
+  /**
+   * Drops every cached lineage graph whose root, nodes, or edges reference the given FQN.
+   * Called when a lineage edge touching this FQN is added or deleted.
+   */
+  public void invalidateLineageCacheForFqn(String fqn) {
+    if (!config.isEnableCaching() || nullOrEmpty(fqn)) {
       return;
     }
+    cache.invalidateIfGraphContains(fqn);
+  }
 
-    // Note: This is a simplified invalidation.
-    // Full implementation would need to invalidate all cache entries
-    // that involve this FQN (as source or in the graph).
-    // For now, we rely on TTL-based expiration.
-    LOG.debug("Cache invalidation requested for fqn={} (TTL-based expiration active)", fqn);
+  /** Drops the entire lineage graph cache. */
+  public void invalidateAllLineageCache() {
+    if (!config.isEnableCaching()) {
+      return;
+    }
+    cache.invalidateAll();
   }
 
   /**

openmetadata-service/src/main/java/org/openmetadata/service/search/lineage/GuavaLineageGraphCache.java

@@ -129,6 +129,55 @@ public void invalidateAll() {
     LOG.info("Cache INVALIDATE_ALL: Cleared all {} entries", cache.size());
   }
 
+  @Override
+  public void invalidateIfGraphContains(String fqn) {
+    if (fqn == null || fqn.isEmpty() || cache.size() == 0) {
+      return;
+    }
+    java.util.List<LineageCacheKey> toEvict = new java.util.ArrayList<>();
+    for (java.util.Map.Entry<LineageCacheKey, SearchLineageResult> entry :
+        cache.asMap().entrySet()) {
+      if (graphReferencesFqn(entry.getKey(), entry.getValue(), fqn)) {
+        toEvict.add(entry.getKey());
+      }
+    }
+    if (!toEvict.isEmpty()) {
+      cache.invalidateAll(toEvict);
+      LOG.debug("Cache INVALIDATE_FQN fqn={} evicted {} entries", fqn, toEvict.size());
+    }
+  }
+
+  private boolean graphReferencesFqn(LineageCacheKey key, SearchLineageResult result, String fqn) {
+    if (fqn.equals(key.getFqn())) {
+      return true;
+    }
+    if (result == null) {
+      return false;
+    }
+    if (result.getNodes() != null && result.getNodes().containsKey(fqn)) {
+      return true;
+    }
+    return edgeMapReferencesFqn(result.getUpstreamEdges(), fqn)
+        || edgeMapReferencesFqn(result.getDownstreamEdges(), fqn);
+  }
+
+  private boolean edgeMapReferencesFqn(
+      java.util.Map<String, org.openmetadata.schema.api.lineage.EsLineageData> edges, String fqn) {
+    if (edges == null || edges.isEmpty()) {
+      return false;
+    }
+    for (org.openmetadata.schema.api.lineage.EsLineageData edge : edges.values()) {
+      if (edge.getFromEntity() != null
+          && fqn.equals(edge.getFromEntity().getFullyQualifiedName())) {
+        return true;
+      }
+      if (edge.getToEntity() != null && fqn.equals(edge.getToEntity().getFullyQualifiedName())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   @Override
   public CacheStats getStats() {
     return cache.stats();

openmetadata-service/src/main/java/org/openmetadata/service/search/lineage/LineageGraphCache.java

@@ -52,6 +52,17 @@ public interface LineageGraphCache {
    */
   void invalidateAll();
 
+  /**
+   * Invalidates every cached graph whose root FQN, nodes map, or edge endpoints reference
+   * the given FQN. Used after lineage edges touching the FQN are added or deleted so stale
+   * graphs are not served.
+   *
+   * @param fqn Fully qualified name of the entity whose graphs should be evicted
+   */
+  default void invalidateIfGraphContains(String fqn) {
+    invalidateAll();
+  }
+
   /**
    * Gets cache statistics for monitoring and metrics.
    *

openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java

@@ -415,6 +415,14 @@ public SearchLineageResult searchPlatformLineage(
     return lineageGraphBuilder.getPlatformLineage(index, queryFilter, deleted);
   }
 
+  @Override
+  public void invalidateLineageCache(String fqn) {
+    if (lineageGraphBuilder == null) {
+      return;
+    }
+    lineageGraphBuilder.invalidateLineageCacheForFqn(fqn);
+  }
+
   @Override
   public Response searchEntityRelationship(
       String fqn, int upstreamDepth, int downstreamDepth, String queryFilter, boolean deleted)

openmetadata-service/src/test/java/org/openmetadata/service/search/lineage/AbstractLineageGraphBuilderTest.java

@@ -815,5 +815,73 @@ public <T> List<T> sortEntitiesByDepthThenName(
         java.util.function.Function<T, String> fqnExtractor) {
       return super.sortEntitiesByDepthThenName(entities, depthExtractor, fqnExtractor);
     }
+
+    public void cacheResultForTest(SearchLineageRequest request, SearchLineageResult result) {
+      super.cacheResult(request, result);
+    }
+
+    public java.util.Optional<SearchLineageResult> checkCacheForTest(SearchLineageRequest request) {
+      return super.checkCache(request);
+    }
+  }
+
+  @Test
+  void invalidateLineageCacheForFqnDropsEntryKeyedByThatFqn() {
+    SearchLineageRequest request =
+        new SearchLineageRequest().withFqn("svc.db.target").withUpstreamDepth(3);
+    SearchLineageResult result = new SearchLineageResult().withNodes(new HashMap<>());
+
+    builder.cacheResultForTest(request, result);
+    assertTrue(builder.checkCacheForTest(request).isPresent(), "Precondition: entry cached");
+
+    builder.invalidateLineageCacheForFqn("svc.db.target");
+
+    assertFalse(
+        builder.checkCacheForTest(request).isPresent(),
+        "Entry whose root FQN matches must be evicted");
+  }
+
+  @Test
+  void invalidateLineageCacheForFqnDropsEntryWhereFqnAppearsInNodes() {
+    SearchLineageRequest request =
+        new SearchLineageRequest().withFqn("svc.db.root").withUpstreamDepth(3);
+    SearchLineageResult result = new SearchLineageResult().withNodes(new HashMap<>());
+    result.getNodes().put("svc.db.target", new NodeInformation().withNodeDepth(1));
+
+    builder.cacheResultForTest(request, result);
+
+    builder.invalidateLineageCacheForFqn("svc.db.target");
+
+    assertFalse(
+        builder.checkCacheForTest(request).isPresent(),
+        "Entry containing the FQN in its nodes must be evicted");
+  }
+
+  @Test
+  void invalidateLineageCacheForFqnLeavesUnrelatedEntriesAlone() {
+    SearchLineageRequest request =
+        new SearchLineageRequest().withFqn("svc.db.kept").withUpstreamDepth(3);
+    SearchLineageResult result = new SearchLineageResult().withNodes(new HashMap<>());
+    result.getNodes().put("svc.db.kept", new NodeInformation().withNodeDepth(0));
+
+    builder.cacheResultForTest(request, result);
+
+    builder.invalidateLineageCacheForFqn("svc.db.unrelated");
+
+    assertTrue(
+        builder.checkCacheForTest(request).isPresent(),
+        "Unrelated cache entry must survive FQN-targeted invalidation");
+  }
+
+  @Test
+  void invalidateLineageCacheForFqnIgnoresNullOrBlank() {
+    SearchLineageRequest request =
+        new SearchLineageRequest().withFqn("svc.db.kept").withUpstreamDepth(3);
+    builder.cacheResultForTest(request, new SearchLineageResult().withNodes(new HashMap<>()));
+
+    builder.invalidateLineageCacheForFqn(null);
+    builder.invalidateLineageCacheForFqn("");
+
+    assertTrue(builder.checkCacheForTest(request).isPresent());
   }
 }

openmetadata-service/src/test/java/org/openmetadata/service/search/lineage/GuavaLineageGraphCacheTest.java

@@ -20,6 +20,8 @@
 import java.util.Optional;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.openmetadata.schema.api.lineage.EsLineageData;
+import org.openmetadata.schema.api.lineage.RelationshipRef;
 import org.openmetadata.schema.api.lineage.SearchLineageRequest;
 import org.openmetadata.schema.api.lineage.SearchLineageResult;
 
@@ -216,4 +218,91 @@ private SearchLineageResult createMockResult(int nodeCount) {
 
     return result;
   }
+
+  @Test
+  public void testInvalidateIfGraphContainsEvictsEntryKeyedByFqn() {
+    LineageCacheKey targetKey =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.target"));
+    LineageCacheKey otherKey =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.other"));
+
+    cache.put(targetKey, createMockResult(5));
+    cache.put(otherKey, createMockResult(5));
+
+    cache.invalidateIfGraphContains("svc.db.target");
+
+    assertFalse(cache.get(targetKey).isPresent(), "Root-FQN-matching entry must be evicted");
+    assertTrue(cache.get(otherKey).isPresent(), "Unrelated entry must survive");
+  }
+
+  @Test
+  public void testInvalidateIfGraphContainsEvictsWhenFqnAppearsAsNode() {
+    LineageCacheKey rootKey =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.root"));
+    LineageCacheKey otherKey =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.elsewhere"));
+
+    SearchLineageResult resultWithTarget = createMockResult(0);
+    resultWithTarget.getNodes().put("svc.db.target", null);
+
+    cache.put(rootKey, resultWithTarget);
+    cache.put(otherKey, createMockResult(5));
+
+    cache.invalidateIfGraphContains("svc.db.target");
+
+    assertFalse(
+        cache.get(rootKey).isPresent(), "Entry whose nodes contain the FQN must be evicted");
+    assertTrue(cache.get(otherKey).isPresent(), "Unrelated entry must survive");
+  }
+
+  @Test
+  public void testInvalidateIfGraphContainsEvictsWhenFqnAppearsAsEdgeEndpoint() {
+    LineageCacheKey rootKey =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.root"));
+
+    SearchLineageResult resultWithEdge = createMockResult(0);
+    EsLineageData edge =
+        new EsLineageData()
+            .withDocId("svc.db.root->svc.db.downstream")
+            .withFromEntity(new RelationshipRef().withFullyQualifiedName("svc.db.root"))
+            .withToEntity(new RelationshipRef().withFullyQualifiedName("svc.db.downstream"));
+    resultWithEdge.getDownstreamEdges().put("edge1", edge);
+
+    cache.put(rootKey, resultWithEdge);
+
+    cache.invalidateIfGraphContains("svc.db.downstream");
+
+    assertFalse(
+        cache.get(rootKey).isPresent(),
+        "Entry whose edge endpoint matches the FQN must be evicted");
+  }
+
+  @Test
+  public void testInvalidateIfGraphContainsIsNoOpForNullOrBlankFqn() {
+    LineageCacheKey key =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.kept"));
+    cache.put(key, createMockResult(5));
+
+    cache.invalidateIfGraphContains(null);
+    cache.invalidateIfGraphContains("");
+
+    assertTrue(cache.get(key).isPresent(), "Null/blank FQN must not evict anything");
+  }
+
+  @Test
+  public void testInvalidateIfGraphContainsLeavesUnrelatedEntries() {
+    LineageCacheKey keyA =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.a"));
+    LineageCacheKey keyB =
+        LineageCacheKey.fromRequest(new SearchLineageRequest().withFqn("svc.db.b"));
+
+    cache.put(keyA, createMockResult(5));
+    cache.put(keyB, createMockResult(5));
+
+    cache.invalidateIfGraphContains("svc.db.unrelated");
+
+    assertTrue(cache.get(keyA).isPresent());
+    assertTrue(cache.get(keyB).isPresent());
+    assertEquals(2, cache.size());
+  }
 }