[EGG Migration] NPE fix for dual read (#556)

* Bug fix

* Checkstyle

---------

Co-authored-by: Rakhi Agrawal <rakagrawal@linkedin.com>

Author: rakhiagr

restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java

@@ -119,8 +119,9 @@ protected Task<VALUE> get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional
       // The assumption is main GMS must have this entity.
       // If entity only has routing aspect, resourceNotFoundException will be thrown.
       final URN urn = toUrn(id);
+      BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO = getShadowReadLocalDAO();
       if (!getLocalDAO().exists(urn)) {
-        if (getShadowReadLocalDAO() != null && getShadowReadLocalDAO().exists(urn)) {
+        if (shadowReadLocalDAO != null && shadowReadLocalDAO.exists(urn)) {
           // If the entity exists in shadow DAO, we can return an empty value.
           log.warn("Entity {} exists in shadow DAO but not in local DAO. Ignoring shadow-only data.", urn);
         }
@@ -571,15 +572,16 @@ private List<INTERNAL_ASPECT_UNION> getInternalAspectsFromLocalDao(URN urn, Set<
         .map(aspectClass -> new AspectKey<>(aspectClass, urn, LATEST_VERSION))
         .collect(Collectors.toSet());
 
-    if (getShadowReadLocalDAO() == null) {
+    BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO = getShadowReadLocalDAO();
+    if (shadowReadLocalDAO == null) {
       return getLocalDAO().get(keys)
           .values()
           .stream()
           .filter(java.util.Optional::isPresent)
           .map(aspect -> ModelUtils.newAspectUnion(_internalAspectUnionClass, aspect.get()))
           .collect(Collectors.toList());
     }
-    return getInternalAspectsWithShadowComparison(keys);
+    return getInternalAspectsWithShadowComparison(keys, shadowReadLocalDAO);
   }
 
   /**
@@ -588,12 +590,13 @@ private List<INTERNAL_ASPECT_UNION> getInternalAspectsFromLocalDao(URN urn, Set<
    * @param keys Aspect keys to be retrieved from shadow DAO
    * @return A list of internal aspects.
    */
-  private List<INTERNAL_ASPECT_UNION> getInternalAspectsWithShadowComparison(Set<AspectKey<URN, ? extends RecordTemplate>> keys) {
+  private List<INTERNAL_ASPECT_UNION> getInternalAspectsWithShadowComparison(
+      Set<AspectKey<URN, ? extends RecordTemplate>> keys, BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO) {
 
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> localResults =
         getLocalDAO().get(keys);
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> shadowResults =
-        getShadowReadLocalDAO().get(keys);
+        shadowReadLocalDAO.get(keys);
 
     return keys.stream()
         .map(key -> {

restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java

@@ -234,8 +234,9 @@ protected Task<VALUE> get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional
 
     return RestliUtils.toTask(() -> {
       final URN urn = toUrn(id);
+      BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO = getShadowReadLocalDAO();
       if (!getLocalDAO().exists(urn)) {
-        if (getShadowReadLocalDAO() != null && getShadowReadLocalDAO().exists(urn)) {
+        if (shadowReadLocalDAO != null && shadowReadLocalDAO.exists(urn)) {
           log.warn("Entity {} exists in shadow DAO but not in local DAO. Ignoring shadow-only data.", urn);
         }
         throw RestliUtils.resourceNotFoundException();
@@ -571,7 +572,8 @@ public Task<ASSET> getAsset(@ActionParam(PARAM_URN) @Nonnull String urnString,
       return RestliUtils.toTask(() -> {
         final URN urn = parseUrnParam(urnString);
 
-        if (getShadowReadLocalDAO() == null) {
+        BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO = getShadowReadLocalDAO();
+        if (shadowReadLocalDAO == null) {
           if (!getLocalDAO().exists(urn)) {
             throw RestliUtils.resourceNotFoundException();
           }
@@ -589,7 +591,7 @@ public Task<ASSET> getAsset(@ActionParam(PARAM_URN) @Nonnull String urnString,
 
           return ModelUtils.newAsset(_assetClass, urn, aspects);
         }
-        return getAssetWithShadowComparison(urn, aspectNames);
+        return getAssetWithShadowComparison(urn, aspectNames, shadowReadLocalDAO);
       });
     } catch (ModelValidationException e) {
       throw RestliUtils.invalidArgumentsException(e.getMessage());
@@ -610,10 +612,10 @@ public Task<ASSET> getAsset(@ActionParam(PARAM_URN) @Nonnull String urnString,
    * @return an asset assembled from the resolved aspects
    * @throws RestLiServiceException if the entity does not exist in the local DAO
    */
-  private ASSET getAssetWithShadowComparison(@Nonnull URN urn, @Nullable String[] aspectNames) {
+  private ASSET getAssetWithShadowComparison(@Nonnull URN urn, @Nullable String[] aspectNames, BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO) {
 
     if (!getLocalDAO().exists(urn)) {
-      if (getShadowReadLocalDAO().exists(urn)) {
+      if (shadowReadLocalDAO.exists(urn)) {
         log.warn("Entity {} exists in shadow DAO but not in local DAO. Ignoring shadow-only data.", urn);
       }
       throw RestliUtils.resourceNotFoundException();
@@ -628,7 +630,7 @@ private ASSET getAssetWithShadowComparison(@Nonnull URN urn, @Nullable String[]
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> localResults =
         getLocalDAO().get(keys);
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> shadowResults =
-        getShadowReadLocalDAO().get(keys);
+        shadowReadLocalDAO.get(keys);
 
     // Collect aspects
     List<UnionTemplate> aspects = new ArrayList<>();
@@ -1230,7 +1232,8 @@ private Map<URN, List<UnionTemplate>> getUrnAspectMap(@Nonnull Collection<URN> u
     final Map<URN, List<UnionTemplate>> urnAspectsMap =
         urns.stream().collect(Collectors.toMap(Function.identity(), urn -> new ArrayList<>()));
 
-    if (getShadowReadLocalDAO() == null) {
+    BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO = getShadowReadLocalDAO();
+    if (shadowReadLocalDAO == null) {
       if (isInternalModelsEnabled) {
         getLocalDAO().get(keys)
             .forEach((key, aspect) -> aspect.ifPresent(metadata -> urnAspectsMap.get(key.getUrn())
@@ -1242,22 +1245,22 @@ private Map<URN, List<UnionTemplate>> getUrnAspectMap(@Nonnull Collection<URN> u
       }
       return urnAspectsMap;
     } else {
-      return getUrnAspectMapFromShadowDao(urns, keys, isInternalModelsEnabled);
+      return getUrnAspectMapFromShadowDao(urns, keys, isInternalModelsEnabled, shadowReadLocalDAO);
     }
   }
 
   @Nonnull
   private Map<URN, List<UnionTemplate>> getUrnAspectMapFromShadowDao(
       @Nonnull Collection<URN> urns,
       @Nonnull Set<AspectKey<URN, ? extends RecordTemplate>> keys,
-      boolean isInternalModelsEnabled) {
+      boolean isInternalModelsEnabled,
+      @Nonnull BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowReadLocalDAO) {
 
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> localResults =
         getLocalDAO().get(keys);
 
-    BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> shadowDao = getShadowReadLocalDAO();
     Map<AspectKey<URN, ? extends RecordTemplate>, java.util.Optional<? extends RecordTemplate>> shadowResults =
-        shadowDao.get(keys);
+        shadowReadLocalDAO.get(keys);
 
     final Map<URN, List<UnionTemplate>> urnAspectsMap =
         urns.stream().collect(Collectors.toMap(Function.identity(), urn -> new ArrayList<>()));