Adding shadowDao for read requests (#541)

* Adding shadowDao for read requests

* Adding shadowDao for read requests

* Fix keys

* Added logs

---------

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

Author: rakhiagr

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

@@ -567,11 +567,53 @@ private List<INTERNAL_ASPECT_UNION> getInternalAspectsFromLocalDao(URN urn, Set<
         .map(aspectClass -> new AspectKey<>(aspectClass, urn, LATEST_VERSION))
         .collect(Collectors.toSet());
 
-    return getLocalDAO().get(keys)
-        .values()
-        .stream()
-        .filter(java.util.Optional::isPresent)
-        .map(aspect -> ModelUtils.newAspectUnion(_internalAspectUnionClass, aspect.get()))
+    if (getShadowReadLocalDAO() == 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);
+  }
+
+  /**
+   * Get internal aspect values from shadow DAO for specified aspect keys.
+   * This method is used to retrieve aspects from shadow DAO when the local DAO does not have the shadow read capability.
+   * @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) {
+
+    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);
+
+    return keys.stream()
+        .map(key -> {
+          java.util.Optional<? extends RecordTemplate> local = localResults.getOrDefault(key, java.util.Optional.empty());
+          java.util.Optional<? extends RecordTemplate> shadow = shadowResults.getOrDefault(key, java.util.Optional.empty());
+
+          RecordTemplate valueToUse = null;
+          if (shadow.isPresent() && local.isPresent() && !Objects.equals(local.get(), shadow.get())) {
+            log.warn("Aspect mismatch for URN {} and aspect {}: local = {}, shadow = {}",
+                key.getUrn(), key.getAspectClass().getSimpleName(), local.get(), shadow.get());
+            valueToUse = local.get();
+          } else if (shadow.isPresent()) {
+            log.warn("Only shadow value present for URN {} and aspect {}", key.getUrn(), key.getAspectClass().getSimpleName());
+            valueToUse = shadow.get();
+          } else if (local.isPresent()) {
+            log.info("Only local value present for URN {} and aspect {}. Using local.", key.getUrn(), key.getAspectClass().getSimpleName());
+            valueToUse = local.get();
+          }
+
+          return valueToUse != null
+              ? ModelUtils.newAspectUnion(_internalAspectUnionClass, valueToUse)
+              : null;
+        })
+        .filter(Objects::nonNull)
         .collect(Collectors.toList());
   }
 

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

@@ -48,11 +48,13 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.metadata.dao.BaseReadDAO.*;
 import static com.linkedin.metadata.dao.utils.IngestionUtils.*;
@@ -73,6 +75,7 @@
  * @param <INTERNAL_ASPECT_UNION> must be a valid internal aspect union type supported by the internal snapshot
  * @param <ASSET> must be a valid asset type defined in com.linkedin.metadata.asset
  */
+@Slf4j
 public abstract class BaseEntityResource<
     // @formatter:off
     KEY,
@@ -172,6 +175,11 @@ protected BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> getShadowLocalDAO() {
     return null; // override in resource class only if needed
   }
 
+  @Nullable
+  protected BaseLocalDAO<INTERNAL_ASPECT_UNION, URN> getShadowReadLocalDAO() {
+    return null; // override in resource class only if needed
+  }
+
   /**
    * Creates an URN from its string representation.
    */
@@ -1142,18 +1150,70 @@ 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 (isInternalModelsEnabled) {
-      getLocalDAO().get(keys)
-          .forEach((key, aspect) -> aspect.ifPresent(metadata -> urnAspectsMap.get(key.getUrn())
-              .add(ModelUtils.newAspectUnion(_internalAspectUnionClass, metadata))));
+    if (getShadowReadLocalDAO() == null) {
+      if (isInternalModelsEnabled) {
+        getLocalDAO().get(keys)
+            .forEach((key, aspect) -> aspect.ifPresent(metadata -> urnAspectsMap.get(key.getUrn())
+                .add(ModelUtils.newAspectUnion(_internalAspectUnionClass, metadata))));
+      } else {
+        getLocalDAO().get(keys)
+            .forEach((key, aspect) -> aspect.ifPresent(metadata -> urnAspectsMap.get(key.getUrn())
+                .add(ModelUtils.newAspectUnion(_aspectUnionClass, metadata))));
+      }
+      return urnAspectsMap;
     } else {
-      getLocalDAO().get(keys)
-          .forEach((key, aspect) -> aspect.ifPresent(
-              metadata -> urnAspectsMap.get(key.getUrn()).add(ModelUtils.newAspectUnion(_aspectUnionClass, metadata))));
+      return getUrnAspectMapFromShadowDao(urns, keys, isInternalModelsEnabled);
     }
+  }
+
+  @Nonnull
+  private Map<URN, List<UnionTemplate>> getUrnAspectMapFromShadowDao(
+      @Nonnull Collection<URN> urns,
+      @Nonnull Set<AspectKey<URN, ? extends RecordTemplate>> keys,
+      boolean isInternalModelsEnabled) {
+
+    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);
+
+    final Map<URN, List<UnionTemplate>> urnAspectsMap =
+        urns.stream().collect(Collectors.toMap(Function.identity(), urn -> new ArrayList<>()));
+
+    keys.forEach(key -> {
+      java.util.Optional<? extends RecordTemplate> localValue = localResults.getOrDefault(key, java.util.Optional.empty());
+      java.util.Optional<? extends RecordTemplate> shadowValue = shadowResults.getOrDefault(key, java.util.Optional.empty());
+
+      RecordTemplate valueToUse = null;
+
+      if (localValue.isPresent() && shadowValue.isPresent()) {
+        if (!Objects.equals(localValue.get(), shadowValue.get())) {
+          log.warn("Aspect mismatch for URN {} and aspect {}: local = {}, shadow = {}",
+              key.getUrn(), key.getAspectClass().getSimpleName(),
+              localValue.get(), shadowValue.get());
+          valueToUse = localValue.get(); // fallback to local if there's mismatch
+        } else {
+          valueToUse = shadowValue.get(); // match → use shadow
+        }
+      } else if (shadowValue.isPresent()) {
+        valueToUse = shadowValue.get();
+      } else if (localValue.isPresent()) {
+        valueToUse = localValue.get();
+      }
+
+      if (valueToUse != null) {
+        urnAspectsMap.get(key.getUrn()).add(ModelUtils.newAspectUnion(
+            (Class<? extends ASPECT_UNION>) (isInternalModelsEnabled ? _internalAspectUnionClass : _aspectUnionClass),
+            valueToUse));
+      }
+    });
+
     return urnAspectsMap;
   }
 
+
   @Nonnull
   private SNAPSHOT newSnapshot(@Nonnull URN urn, @Nonnull List<UnionTemplate> aspects) {
     return ModelUtils.newSnapshot(_snapshotClass, urn, aspects);