fix(lineage): add cache key truncation for search lineage key (datahu…

…b-project#7671)

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

@@ -4,6 +4,10 @@
 import com.linkedin.metadata.graph.LineageDirection;
 import lombok.Data;
 
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.time.temporal.TemporalUnit;
+
 
 @Data
 public class EntityLineageResultCacheKey {
@@ -12,4 +16,28 @@ public class EntityLineageResultCacheKey {
   private final Long startTimeMillis;
   private final Long endTimeMillis;
   private final Integer maxHops;
+
+  // Force use of static method outside of package (for tests)
+  EntityLineageResultCacheKey(Urn sourceUrn, LineageDirection direction, Long startTimeMillis, Long endTimeMillis, Integer maxHops) {
+    this.sourceUrn = sourceUrn;
+    this.direction = direction;
+    this.startTimeMillis = startTimeMillis;
+    this.endTimeMillis = endTimeMillis;
+    this.maxHops = maxHops;
+  }
+
+  public static EntityLineageResultCacheKey from(Urn sourceUrn, LineageDirection direction, Long startTimeMillis,
+                                                 Long endTimeMillis, Integer maxHops) {
+    return EntityLineageResultCacheKey.from(sourceUrn, direction, startTimeMillis, endTimeMillis, maxHops, ChronoUnit.DAYS);
+  }
+
+  public static EntityLineageResultCacheKey from(Urn sourceUrn, LineageDirection direction, Long startTimeMillis,
+                                                 Long endTimeMillis, Integer maxHops, TemporalUnit resolution) {
+    long endOffset = resolution.getDuration().getSeconds() * 1000;
+    return new EntityLineageResultCacheKey(sourceUrn, direction,
+        startTimeMillis == null ? null
+            : Instant.ofEpochMilli(startTimeMillis).truncatedTo(resolution).toEpochMilli(),
+        endTimeMillis == null ? null : Instant.ofEpochMilli(endTimeMillis + endOffset).truncatedTo(resolution).toEpochMilli(),
+        maxHops);
+  }
 }

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

@@ -122,7 +122,8 @@ public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull
     }
 
     // Cache multihop result for faster performance
-    final EntityLineageResultCacheKey cacheKey = new EntityLineageResultCacheKey(sourceUrn, direction, startTimeMillis, endTimeMillis, maxHops);
+    final EntityLineageResultCacheKey cacheKey = EntityLineageResultCacheKey.from(sourceUrn, direction, startTimeMillis,
+            endTimeMillis, maxHops);
     CachedEntityLineageResult cachedLineageResult = null;
 
     if (cacheEnabled) {

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

@@ -0,0 +1,25 @@
+package com.linkedin.metadata.search;
+
+import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
+import org.testng.annotations.Test;
+
+import static org.testng.AssertJUnit.assertEquals;
+
+
+public class LineageSearchResultCacheKeyTest extends AbstractTestNGSpringContextTests {
+
+  @Test
+  public void testNulls() {
+    // ensure no NPE
+    assertEquals(new EntityLineageResultCacheKey(null, null, null, null, null),
+            EntityLineageResultCacheKey.from(null, null, null, null, null));
+  }
+
+  @Test
+  public void testDateTruncation() {
+    // expect start of day milli
+    assertEquals(new EntityLineageResultCacheKey(null, null, 1679529600000L, 1679616000000L, null),
+            EntityLineageResultCacheKey.from(null, null, 1679530293000L, 1679530293001L,
+                    null));
+  }
+}