From aabb3e907de705278b611120258c33747616732a Mon Sep 17 00:00:00 2001
From: Jesse Jia <jejia@linkedin.com>
Date: Wed, 18 Sep 2024 12:28:36 -0700
Subject: [PATCH] Fetch urns from search result metadata (#432)

Co-authored-by: Jesse Jia <jejia@jejia-mn2.linkedin.biz>
---
 .../restli/BaseSearchableEntityResource.java    | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java
index 8eeb9db23..a84e027a1 100644
--- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java
+++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java
@@ -23,6 +23,7 @@
 import com.linkedin.restli.server.annotations.PagingContextParam;
 import com.linkedin.restli.server.annotations.QueryParam;
 import com.linkedin.restli.server.annotations.RestMethod;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -200,10 +201,18 @@ public Task<AutoCompleteResult> autocomplete(@ActionParam(PARAM_QUERY) @Nonnull
   public CollectionResult<VALUE, SearchResultMetadata> getSearchQueryCollectionResult(@Nonnull SearchResult<DOCUMENT> searchResult,
       @Nullable String[] aspectNames, boolean isInternalModelsEnabled) {
 
-    final List<URN> matchedUrns = searchResult.getDocumentList()
-        .stream()
-        .map(d -> (URN) ModelUtils.getUrnFromDocument(d))
-        .collect(Collectors.toList());
+    List<URN> matchedUrns = new ArrayList<>();
+
+    // First try to fetch the urns from document field. Then try to fetch urns from SearchResultMetadata.
+    if (searchResult.getDocumentList() != null && !searchResult.getDocumentList().isEmpty()) {
+      matchedUrns = searchResult.getDocumentList()
+          .stream()
+          .map(d -> (URN) ModelUtils.getUrnFromDocument(d))
+          .collect(Collectors.toList());
+    } else if (searchResult.getSearchResultMetadata().hasUrns()) {
+      matchedUrns = searchResult.getSearchResultMetadata().getUrns().stream().map(urn -> (URN) urn).collect(Collectors.toList());
+    }
+
     final Map<URN, VALUE> urnValueMap =
         getInternalNonEmpty(matchedUrns, parseAspectsParam(aspectNames, isInternalModelsEnabled),
             isInternalModelsEnabled);