feat(views): apply views to homepage entity counts & recommendations (d…

…atahub-project#10283) Co-authored-by: gaurav2733 <[email protected]>
acryldata · Apr 18, 2024 · 91e3dc8 · 91e3dc8
1 parent a041a2e
commit 91e3dc8
Show file tree

Hide file tree

Showing 25 changed files with 224 additions and 49 deletions.
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -55,6 +55,7 @@
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.DatasetStatsSummary;
 import com.linkedin.datahub.graphql.generated.Domain;
+import com.linkedin.datahub.graphql.generated.ERModelRelationship;
 import com.linkedin.datahub.graphql.generated.ERModelRelationshipProperties;
 import com.linkedin.datahub.graphql.generated.EntityPath;
 import com.linkedin.datahub.graphql.generated.EntityRelationship;
@@ -987,8 +988,10 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher("listUsers", new ListUsersResolver(this.entityClient))
                 .dataFetcher("listGroups", new ListGroupsResolver(this.entityClient))
                 .dataFetcher(
-                    "listRecommendations", new ListRecommendationsResolver(recommendationsService))
-                .dataFetcher("getEntityCounts", new EntityCountsResolver(this.entityClient))
+                    "listRecommendations",
+                    new ListRecommendationsResolver(recommendationsService, viewService))
+                .dataFetcher(
+                    "getEntityCounts", new EntityCountsResolver(this.entityClient, viewService))
                 .dataFetcher("getAccessToken", new GetAccessTokenResolver(statefulTokenService))
                 .dataFetcher("listAccessTokens", new ListAccessTokensResolver(this.entityClient))
                 .dataFetcher(

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ResolverUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ResolverUtils.java
@@ -1,17 +1,20 @@
 package com.linkedin.datahub.graphql.resolvers;
 
+import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*;
 import static com.linkedin.metadata.Constants.*;
 
 import com.datahub.authentication.Authentication;
 import com.fasterxml.jackson.core.StreamReadConstraints;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.ValidationException;
 import com.linkedin.datahub.graphql.generated.AndFilterInput;
 import com.linkedin.datahub.graphql.generated.FacetFilterInput;
+import com.linkedin.datahub.graphql.resolvers.search.SearchUtils;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -20,7 +23,10 @@
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.utils.ESUtils;
 import com.linkedin.metadata.search.utils.QueryUtils;
+import com.linkedin.metadata.service.ViewService;
+import com.linkedin.view.DataHubViewInfo;
 import graphql.schema.DataFetchingEnvironment;
+import io.datahubproject.metadata.context.OperationContext;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -226,4 +232,14 @@ public static Filter buildFilterWithUrns(@Nonnull Set<Urn> urns, @Nullable Filte
     }
     return QueryUtils.newFilter(urnMatchCriterion);
   }
+
+  public static Filter viewFilter(
+      OperationContext opContext, ViewService viewService, String viewUrn) {
+    if (viewUrn == null) {
+      return null;
+    }
+    DataHubViewInfo viewInfo = resolveView(opContext, viewService, UrnUtils.getUrn(viewUrn));
+    Filter result = SearchUtils.combineFilters(null, viewInfo.getDefinition().getFilter());
+    return result;
+  }
 }
diff --git a/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java b/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java
@@ -8,6 +8,7 @@
 import com.linkedin.datahub.graphql.generated.EntityCountResults;
 import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper;
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.service.ViewService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.opentelemetry.extension.annotations.WithSpan;
@@ -20,8 +21,11 @@ public class EntityCountsResolver implements DataFetcher<CompletableFuture<Entit
 
   private final EntityClient _entityClient;
 
-  public EntityCountsResolver(final EntityClient entityClient) {
+  private final ViewService _viewService;
+
+  public EntityCountsResolver(final EntityClient entityClient, final ViewService viewService) {
     _entityClient = entityClient;
+    _viewService = viewService;
   }
 
   @Override
@@ -44,7 +48,8 @@ public CompletableFuture<EntityCountResults> get(final DataFetchingEnvironment e
                     context.getOperationContext(),
                     input.getTypes().stream()
                         .map(EntityTypeMapper::getName)
-                        .collect(Collectors.toList()));
+                        .collect(Collectors.toList()),
+                    viewFilter(context.getOperationContext(), _viewService, input.getViewUrn()));
 
             // bind to a result.
             List<EntityCountResult> resultList =

diff --git a/...va/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java b/...va/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java
@@ -22,6 +22,7 @@
 import com.linkedin.metadata.recommendation.EntityRequestContext;
 import com.linkedin.metadata.recommendation.RecommendationsService;
 import com.linkedin.metadata.recommendation.SearchRequestContext;
+import com.linkedin.metadata.service.ViewService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.opentelemetry.extension.annotations.WithSpan;
@@ -44,6 +45,7 @@ public class ListRecommendationsResolver
       new ListRecommendationsResult(Collections.emptyList());
 
   private final RecommendationsService _recommendationsService;
+  private final ViewService _viewService;
 
   @WithSpan
   @Override
@@ -60,6 +62,7 @@ public CompletableFuture<ListRecommendationsResult> get(DataFetchingEnvironment
                 _recommendationsService.listRecommendations(
                     context.getOperationContext(),
                     mapRequestContext(input.getRequestContext()),
+                    viewFilter(context.getOperationContext(), _viewService, input.getViewUrn()),
                     input.getLimit());
             return ListRecommendationsResult.builder()
                 .setModules(

diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -1194,6 +1194,11 @@ Input for the get entity counts endpoint
 """
 input EntityCountInput {
     types: [EntityType!]
+
+    """
+    Optional - A View to apply when generating results
+    """
+    viewUrn: String
 }
 
 """

diff --git a/datahub-graphql-core/src/main/resources/recommendation.graphql b/datahub-graphql-core/src/main/resources/recommendation.graphql
@@ -23,6 +23,11 @@ input ListRecommendationsInput {
   Max number of modules to return
   """
   limit: Int
+
+  """
+  Optional - A View to apply when generating results
+  """
+  viewUrn: String
 }
 
 """

diff --git a/datahub-web-react/src/app/home/HomePageRecommendations.tsx b/datahub-web-react/src/app/home/HomePageRecommendations.tsx
@@ -22,6 +22,7 @@ import {
 } from '../onboarding/config/HomePageOnboardingConfig';
 import { useToggleEducationStepIdsAllowList } from '../onboarding/useToggleEducationStepIdsAllowList';
 import { useBusinessAttributesFlag } from '../useAppConfig';
+import { useUserContext } from '../context/useUserContext';
 
 const PLATFORMS_MODULE_ID = 'Platforms';
 const MOST_POPULAR_MODULE_ID = 'HighUsageEntities';
@@ -105,6 +106,9 @@ export const HomePageRecommendations = ({ user }: Props) => {
     const browseEntityList = entityRegistry.getBrowseEntityTypes();
     const userUrn = user?.urn;
 
+    const userContext = useUserContext();
+    const viewUrn = userContext.localState?.selectedViewUrn;
+
     const businessAttributesFlag = useBusinessAttributesFlag();
 
     const showSimplifiedHomepage = user?.settings?.appearance?.showSimplifiedHomepage;
@@ -113,6 +117,7 @@ export const HomePageRecommendations = ({ user }: Props) => {
         variables: {
             input: {
                 types: browseEntityList,
+                viewUrn
             },
         },
     });
@@ -133,6 +138,7 @@ export const HomePageRecommendations = ({ user }: Props) => {
                     scenario,
                 },
                 limit: 10,
+                viewUrn
             },
         },
         fetchPolicy: 'no-cache',

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java
@@ -571,9 +571,11 @@ public void setWritable(@Nonnull OperationContext opContext, boolean canWrite)
   @Override
   @Nonnull
   public Map<String, Long> batchGetTotalEntityCount(
-      @Nonnull OperationContext opContext, @Nonnull List<String> entityNames)
+      @Nonnull OperationContext opContext,
+      @Nonnull List<String> entityNames,
+      @Nullable Filter filter)
       throws RemoteInvocationException {
-    return searchService.docCountPerEntity(opContext, entityNames);
+    return searchService.docCountPerEntity(opContext, entityNames, filter);
   }
 
   /** List all urns existing for a particular Entity type. */

diff --git a/...src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java b/...src/main/java/com/linkedin/metadata/recommendation/candidatesource/MostPopularSource.java
@@ -8,6 +8,7 @@
 import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.recommendation.RecommendationContent;
 import com.linkedin.metadata.recommendation.RecommendationRenderType;
 import com.linkedin.metadata.recommendation.RecommendationRequestContext;
@@ -23,6 +24,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.opensearch.action.search.SearchRequest;
@@ -99,7 +101,9 @@ public boolean isEligible(
   @Override
   @WithSpan
   public List<RecommendationContent> getRecommendations(
-      @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) {
+      @Nonnull OperationContext opContext,
+      @Nonnull RecommendationRequestContext requestContext,
+      @Nullable Filter filter) {
     SearchRequest searchRequest = buildSearchRequest(opContext);
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getMostPopular").time()) {
       final SearchResponse searchResponse =

diff --git a/.../main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java b/.../main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyEditedSource.java
@@ -8,6 +8,7 @@
 import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.recommendation.RecommendationContent;
 import com.linkedin.metadata.recommendation.RecommendationRenderType;
 import com.linkedin.metadata.recommendation.RecommendationRequestContext;
@@ -22,6 +23,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.opensearch.action.search.SearchRequest;
@@ -98,7 +100,9 @@ public boolean isEligible(
   @Override
   @WithSpan
   public List<RecommendationContent> getRecommendations(
-      @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) {
+      @Nonnull OperationContext opContext,
+      @Nonnull RecommendationRequestContext requestContext,
+      @Nullable Filter filter) {
     SearchRequest searchRequest =
         buildSearchRequest(opContext.getSessionActorContext().getActorUrn());
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlyEdited").time()) {

diff --git a/.../main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java b/.../main/java/com/linkedin/metadata/recommendation/candidatesource/RecentlyViewedSource.java
@@ -8,6 +8,7 @@
 import com.linkedin.metadata.datahubusage.DataHubUsageEventConstants;
 import com.linkedin.metadata.datahubusage.DataHubUsageEventType;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.recommendation.RecommendationContent;
 import com.linkedin.metadata.recommendation.RecommendationRenderType;
 import com.linkedin.metadata.recommendation.RecommendationRequestContext;
@@ -22,6 +23,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.opensearch.action.search.SearchRequest;
@@ -98,7 +100,9 @@ public boolean isEligible(
   @Override
   @WithSpan
   public List<RecommendationContent> getRecommendations(
-      @Nonnull OperationContext opContext, @Nonnull RecommendationRequestContext requestContext) {
+      @Nonnull OperationContext opContext,
+      @Nonnull RecommendationRequestContext requestContext,
+      @Nullable Filter filter) {
     SearchRequest searchRequest =
         buildSearchRequest(opContext.getSessionActorContext().getActorUrn());
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getRecentlyViewed").time()) {

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
@@ -40,13 +40,20 @@ public SearchService(
 
   public Map<String, Long> docCountPerEntity(
       @Nonnull OperationContext opContext, @Nonnull List<String> entityNames) {
+    return docCountPerEntity(opContext, entityNames, null);
+  }
+
+  public Map<String, Long> docCountPerEntity(
+      @Nonnull OperationContext opContext,
+      @Nonnull List<String> entityNames,
+      @Nullable Filter filter) {
     return getEntitiesToSearch(opContext, entityNames, 0).stream()
         .collect(
             Collectors.toMap(
                 Function.identity(),
                 entityName ->
                     _entityDocCountCache
-                        .getEntityDocCount(opContext)
+                        .getEntityDocCount(opContext, filter)
                         .getOrDefault(entityName.toLowerCase(), 0L)));
   }
 

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java b/metadata-io/src/main/java/com/linkedin/metadata/search/cache/EntityDocCountCache.java
@@ -3,6 +3,7 @@
 import com.google.common.base.Suppliers;
 import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration;
 import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.EntitySearchService;
 import com.linkedin.metadata.utils.ConcurrencyUtils;
 import io.datahubproject.metadata.context.OperationContext;
@@ -15,12 +16,22 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
 
 public class EntityDocCountCache {
   private final EntityRegistry entityRegistry;
   private final EntitySearchService entitySearchService;
   private final EntityDocCountCacheConfiguration config;
-  private final Map<String, Supplier<Map<String, Long>>> entityDocCounts;
+  private final Map<EntityDocCountsKey, Supplier<Map<String, Long>>> entityDocCounts;
+
+  @AllArgsConstructor
+  @EqualsAndHashCode
+  private static final class EntityDocCountsKey {
+    private final String searchContextId;
+    private final Filter filter;
+  }
 
   public EntityDocCountCache(
       EntityRegistry entityRegistry,
@@ -32,17 +43,27 @@ public EntityDocCountCache(
     this.entityDocCounts = new ConcurrentHashMap<>();
   }
 
-  private Map<String, Long> fetchEntityDocCount(@Nonnull OperationContext opContext) {
+  private Map<String, Long> fetchEntityDocCount(
+      @Nonnull OperationContext opContext, @Nullable Filter filter) {
     return ConcurrencyUtils.transformAndCollectAsync(
         entityRegistry.getEntitySpecs().keySet(),
         Function.identity(),
-        Collectors.toMap(Function.identity(), v -> entitySearchService.docCount(opContext, v)));
+        Collectors.toMap(
+            Function.identity(), v -> entitySearchService.docCount(opContext, v, filter)));
   }
 
   @WithSpan
   public Map<String, Long> getEntityDocCount(@Nonnull OperationContext opContext) {
+    return getEntityDocCount(opContext, null);
+  }
+
+  @WithSpan
+  public Map<String, Long> getEntityDocCount(
+      @Nonnull OperationContext opContext, @Nullable Filter filter) {
     return entityDocCounts
-        .computeIfAbsent(opContext.getSearchContextId(), k -> buildSupplier(opContext))
+        .computeIfAbsent(
+            new EntityDocCountsKey(opContext.getSearchContextId(), filter),
+            k -> buildSupplier(opContext, filter))
         .get();
   }
 
@@ -53,8 +74,9 @@ public List<String> getNonEmptyEntities(@Nonnull OperationContext opContext) {
         .collect(Collectors.toList());
   }
 
-  private Supplier<Map<String, Long>> buildSupplier(@Nonnull OperationContext opContext) {
+  private Supplier<Map<String, Long>> buildSupplier(
+      @Nonnull OperationContext opContext, @Nullable Filter filter) {
     return Suppliers.memoizeWithExpiration(
-        () -> fetchEntityDocCount(opContext), config.getTtlSeconds(), TimeUnit.SECONDS);
+        () -> fetchEntityDocCount(opContext, filter), config.getTtlSeconds(), TimeUnit.SECONDS);
   }
 }
diff --git a/...ata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java b/...ata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java
@@ -81,11 +81,13 @@ public void clear(@Nonnull OperationContext opContext) {
   }
 
   @Override
-  public long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName) {
+  public long docCount(
+      @Nonnull OperationContext opContext, @Nonnull String entityName, @Nullable Filter filter) {
     return esSearchDAO.docCount(
         opContext.withSearchFlags(
             flags -> applyDefaultSearchFlags(flags, null, DEFAULT_SERVICE_SEARCH_FLAGS)),
-        entityName);
+        entityName,
+        filter);
   }
 
   @Override