From 6a17c7156f28a37aac9032017b5683be822ed524 Mon Sep 17 00:00:00 2001
From: Chris Collins <chriscollins3456@gmail.com>
Date: Tue, 9 Jul 2024 19:12:08 -0400
Subject: [PATCH] fix(lineage) Fix lineage source/dest filtering with explored
 per hop limit

---
 .../graph/elastic/ESGraphQueryDAO.java        | 30 ++++++++++++++-----
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
index 30b688761d584..e21c4e15b4fc9 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java
@@ -177,13 +177,18 @@ private SearchResponse executeGroupByLineageSearchQuery(
     sourceFilterQuery.minimumShouldMatch(1);
     validEdges.stream()
         .filter(pair -> RelationshipDirection.OUTGOING.equals(pair.getValue().getDirection()))
-        .forEach(pair -> sourceFilterQuery.should(getAggregationFilter(pair)));
+        .forEach(
+            pair ->
+                sourceFilterQuery.should(
+                    getAggregationFilter(pair, RelationshipDirection.OUTGOING)));
 
     BoolQueryBuilder destFilterQuery = QueryBuilders.boolQuery();
     destFilterQuery.minimumShouldMatch(1);
     validEdges.stream()
         .filter(pair -> RelationshipDirection.INCOMING.equals(pair.getValue().getDirection()))
-        .forEach(pair -> destFilterQuery.should(getAggregationFilter(pair)));
+        .forEach(
+            pair ->
+                destFilterQuery.should(getAggregationFilter(pair, RelationshipDirection.INCOMING)));
 
     FilterAggregationBuilder sourceRelationshipTypeFilters =
         AggregationBuilders.filter(FILTER_BY_SOURCE_RELATIONSHIP, sourceFilterQuery);
@@ -226,17 +231,28 @@ private SearchResponse executeGroupByLineageSearchQuery(
     }
   }
 
-  private BoolQueryBuilder getAggregationFilter(Pair<String, EdgeInfo> pair) {
+  private BoolQueryBuilder getAggregationFilter(
+      Pair<String, EdgeInfo> pair, RelationshipDirection direction) {
     BoolQueryBuilder subFilter = QueryBuilders.boolQuery();
     TermQueryBuilder relationshipTypeTerm =
-        QueryBuilders.termQuery(RELATIONSHIP_TYPE, pair.getValue().getType());
+        QueryBuilders.termQuery(RELATIONSHIP_TYPE, pair.getValue().getType()).caseInsensitive(true);
     subFilter.must(relationshipTypeTerm);
+
+    String sourceType;
+    String destinationType;
+    if (direction.equals(RelationshipDirection.OUTGOING)) {
+      sourceType = pair.getKey();
+      destinationType = pair.getValue().getOpposingEntityType();
+    } else {
+      sourceType = pair.getValue().getOpposingEntityType();
+      destinationType = pair.getKey();
+    }
+
     TermQueryBuilder sourceTypeTerm =
-        QueryBuilders.termQuery(SOURCE + ".entityType", pair.getKey());
+        QueryBuilders.termQuery(SOURCE + ".entityType", sourceType).caseInsensitive(true);
     subFilter.must(sourceTypeTerm);
     TermQueryBuilder destinationTypeTerm =
-        QueryBuilders.termQuery(
-            DESTINATION + ".entityType", pair.getValue().getOpposingEntityType());
+        QueryBuilders.termQuery(DESTINATION + ".entityType", destinationType).caseInsensitive(true);
     subFilter.must(destinationTypeTerm);
     return subFilter;
   }