From 45629788746c184317e5c30cae73cd2aa72c112f Mon Sep 17 00:00:00 2001
From: "Singh, Himanshu" <hisingh@visa.com>
Date: Tue, 19 Mar 2024 19:30:29 +0530
Subject: [PATCH] Fix : Business Attribute SearchableRef Search Depth

---
 .../elasticsearch/indexbuilder/MappingsBuilder.java       | 8 +++++++-
 .../elasticsearch/query/request/SearchFieldConfig.java    | 4 +++-
 .../elasticsearch/query/request/SearchQueryBuilder.java   | 4 +++-
 .../search/transformer/SearchDocumentTransformer.java     | 5 +++--
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java
index 8b2eb5a781701..8696d5a1d557c 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/MappingsBuilder.java
@@ -332,7 +332,13 @@ private static Map<String, Object> getMappingForSearchableRefField(
         .forEach(
             entitySearchableRefFieldSpec ->
                 mappingForField.putAll(
-                    getMappingForSearchableRefField(entitySearchableRefFieldSpec, depth - 1)));
+                    getMappingForSearchableRefField(
+                        entitySearchableRefFieldSpec,
+                        Math.min(
+                            depth - 1,
+                            entitySearchableRefFieldSpec
+                                .getSearchableRefAnnotation()
+                                .getDepth()))));
     mappingForField.put("urn", getMappingsForUrn());
     mappingForProperty.put("properties", mappingForField);
     mappings.put(
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java
index a7ba78230e9de..7415c6e5ce5aa 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchFieldConfig.java
@@ -145,12 +145,14 @@ public static Set<SearchFieldConfig> detectSubFieldType(
             aspectSpec.getSearchableRefFieldSpecs()) {
           String refFieldName = searchableRefFieldSpec.getSearchableRefAnnotation().getFieldName();
           refFieldName = fieldName + "." + refFieldName;
+          int newDepth =
+              Math.min(depth - 1, searchableRefFieldSpec.getSearchableRefAnnotation().getDepth());
           final float refBoost =
               (float) searchableRefFieldSpec.getSearchableRefAnnotation().getBoostScore()
                   * boostScore;
           fieldConfigs.addAll(
               detectSubFieldType(
-                  searchableRefFieldSpec, depth - 1, entityRegistry, refBoost, refFieldName));
+                  searchableRefFieldSpec, newDepth, entityRegistry, refBoost, refFieldName));
         }
       }
     }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java
index 52067003eef14..1aa92b71b7a52 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java
@@ -682,9 +682,11 @@ public Map<String, SearchableAnnotation.FieldType> getAllFieldTypeFromSearchable
             aspectSpec.getSearchableRefFieldSpecs()) {
           String refFieldName = searchableRefFieldSpec.getSearchableRefAnnotation().getFieldName();
           refFieldName = fieldName + "." + refFieldName;
+          int newDepth =
+              Math.min(depth - 1, searchableRefFieldSpec.getSearchableRefAnnotation().getDepth());
           fieldNameMap.putAll(
               getAllFieldTypeFromSearchableRef(
-                  searchableRefFieldSpec, depth - 1, entityRegistry, refFieldName));
+                  searchableRefFieldSpec, newDepth, entityRegistry, refFieldName));
         }
       }
     }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java
index 1f9487be28e91..9651198bb7984 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java
@@ -502,6 +502,7 @@ private Optional<JsonNode> getNodeForRef(
                 String fieldName = spec.getSearchableRefAnnotation().getFieldName();
                 boolean isArray = spec.isArray();
                 if (!value.isEmpty()) {
+                  int newDepth = Math.min(depth - 1, spec.getSearchableRefAnnotation().getDepth());
                   if (isArray) {
                     ArrayNode arrayNode = JsonNodeFactory.instance.arrayNode();
                     value
@@ -509,7 +510,7 @@ private Optional<JsonNode> getNodeForRef(
                         .forEach(
                             val ->
                                 getNodeForRef(
-                                        depth - 1,
+                                        newDepth,
                                         val,
                                         spec.getSearchableRefAnnotation().getFieldType())
                                     .ifPresent(arrayNode::add));
@@ -517,7 +518,7 @@ private Optional<JsonNode> getNodeForRef(
                   } else {
                     Optional<JsonNode> node =
                         getNodeForRef(
-                            depth - 1,
+                            newDepth,
                             value.get(0),
                             spec.getSearchableRefAnnotation().getFieldType());
                     if (node.isPresent()) {