Properly cast generic urn to entity specific urn (#457)

* Add batchGet for BaseAspectV2Resource * fix * fix * Properly cast generic Urn to entity specific urn --------- Co-authored-by: Derek Pham <[email protected]>
linkedin · Oct 23, 2024 · 5bbe44c · 5bbe44c
1 parent b762dac
commit 5bbe44c
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/...li-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java b/...li-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.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -210,7 +211,7 @@ public CollectionResult<VALUE, SearchResultMetadata> getSearchQueryCollectionRes
           .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());
+      matchedUrns = searchResult.getSearchResultMetadata().getUrns().stream().map(this::toEntitySpecificUrn).collect(Collectors.toList());
     }
 
     final Map<URN, VALUE> urnValueMap =
@@ -223,4 +224,12 @@ public CollectionResult<VALUE, SearchResultMetadata> getSearchQueryCollectionRes
         searchResult.getSearchResultMetadata().setUrns(new UrnArray(existingUrns.stream().map(urn -> (Urn) urn).collect(Collectors.toList())))
     );
   }
+
+  private URN toEntitySpecificUrn(Urn urn) {
+     try {
+       return (URN) (_urnClass.getMethod("createFromUrn", Urn.class).invoke(null, urn));
+     } catch (InvocationTargetException | IllegalAccessException | NoSuchMethodException e) {
+       throw new RuntimeException(e);
+     }
+  }
 }