feat(gms): Add support for platform-based browse (datahub-project#9376)

Co-authored-by: John Joyce <john@Johns-MBP.attlocal.net>

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/featureflags/FeatureFlags.java

@@ -12,6 +12,7 @@ public class FeatureFlags {
   private boolean readOnlyModeEnabled = false;
   private boolean showSearchFiltersV2 = false;
   private boolean showBrowseV2 = false;
+  private boolean platformBrowseV2 = false;
   private PreProcessHooks preProcessHooks;
   private boolean showAcrylInfo = false;
   private boolean showAccessManagement = false;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java

@@ -2,14 +2,16 @@
 
 import static com.linkedin.datahub.graphql.Constants.BROWSE_PATH_V2_DELIMITER;
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
-import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.resolveView;
+import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*;
 
+import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.BrowseResultGroupV2;
 import com.linkedin.datahub.graphql.generated.BrowseResultMetadata;
 import com.linkedin.datahub.graphql.generated.BrowseResultsV2;
 import com.linkedin.datahub.graphql.generated.BrowseV2Input;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.resolvers.EntityTypeMapper;
 import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
 import com.linkedin.datahub.graphql.resolvers.search.SearchUtils;
@@ -43,8 +45,8 @@ public class BrowseV2Resolver implements DataFetcher<CompletableFuture<BrowseRes
   public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environment) {
     final QueryContext context = environment.getContext();
     final BrowseV2Input input = bindArgument(environment.getArgument("input"), BrowseV2Input.class);
-    final String entityName = EntityTypeMapper.getName(input.getType());
 
+    final List<String> entityNames = getEntityNames(input);
     final int start = input.getStart() != null ? input.getStart() : DEFAULT_START;
     final int count = input.getCount() != null ? input.getCount() : DEFAULT_COUNT;
     final String query = input.getQuery() != null ? input.getQuery() : "*";
@@ -70,7 +72,7 @@ public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environmen
 
             BrowseResultV2 browseResults =
                 _entityClient.browseV2(
-                    entityName,
+                    entityNames,
                     pathStr,
                     maybeResolvedView != null
                         ? SearchUtils.combineFilters(
@@ -87,6 +89,18 @@ public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environmen
         });
   }
 
+  public static List<String> getEntityNames(BrowseV2Input input) {
+    List<EntityType> entityTypes;
+    if (input.getTypes() != null && input.getTypes().size() > 0) {
+      entityTypes = input.getTypes();
+    } else if (input.getType() != null) {
+      entityTypes = ImmutableList.of(input.getType());
+    } else {
+      entityTypes = BROWSE_ENTITY_TYPES;
+    }
+    return entityTypes.stream().map(EntityTypeMapper::getName).collect(Collectors.toList());
+  }
+
   private BrowseResultsV2 mapBrowseResults(BrowseResultV2 browseResults) {
     BrowseResultsV2 results = new BrowseResultsV2();
     results.setTotal(browseResults.getNumGroups());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java

@@ -175,6 +175,7 @@ public CompletableFuture<AppConfig> get(final DataFetchingEnvironment environmen
             .setShowAcrylInfo(_featureFlags.isShowAcrylInfo())
             .setShowAccessManagement(_featureFlags.isShowAccessManagement())
             .setNestedDomainsEnabled(_featureFlags.isNestedDomainsEnabled())
+            .setPlatformBrowseV2(_featureFlags.isPlatformBrowseV2())
             .build();
 
     appConfig.setFeatureFlags(featureFlagsConfig);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java

@@ -92,6 +92,20 @@ private SearchUtils() {}
           EntityType.NOTEBOOK,
           EntityType.DATA_PRODUCT);
 
+  /** Entities that are part of browse by default */
+  public static final List<EntityType> BROWSE_ENTITY_TYPES =
+      ImmutableList.of(
+          EntityType.DATASET,
+          EntityType.DASHBOARD,
+          EntityType.CHART,
+          EntityType.CONTAINER,
+          EntityType.MLMODEL,
+          EntityType.MLMODEL_GROUP,
+          EntityType.MLFEATURE_TABLE,
+          EntityType.DATA_FLOW,
+          EntityType.DATA_JOB,
+          EntityType.NOTEBOOK);
+
   /** A prioritized list of source filter types used to generate quick filters */
   public static final List<String> PRIORITIZED_SOURCE_ENTITY_TYPES =
       Stream.of(

datahub-graphql-core/src/main/resources/app.graphql

@@ -437,6 +437,11 @@ type FeatureFlagsConfig {
   """
   showBrowseV2: Boolean!
 
+  """
+  Whether browse v2 is platform mode, which means that platforms are displayed instead of entity types at the root.
+  """
+  platformBrowseV2: Boolean!
+
   """
   Whether we should show CTAs in the UI related to moving to Managed DataHub by Acryl.
   """

datahub-graphql-core/src/main/resources/search.graphql

@@ -1176,9 +1176,14 @@ Input required for browse queries
 """
 input BrowseV2Input {
   """
-  The browse entity type
+  The browse entity type - deprecated use types instead
   """
-  type: EntityType!
+  type: EntityType
+
+  """
+  The browse entity type - deprecated use types instead. If not provided, all types will be used.
+  """
+  types: [EntityType!]
 
   """
   The browse path V2 - a list with each entry being part of the browse path V2

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java

@@ -249,7 +249,7 @@ private static EntityClient initMockEntityClient(
     EntityClient client = Mockito.mock(EntityClient.class);
     Mockito.when(
             client.browseV2(
-                Mockito.eq(entityName),
+                Mockito.eq(ImmutableList.of(entityName)),
                 Mockito.eq(path),
                 Mockito.eq(filter),
                 Mockito.eq(query),

datahub-web-react/src/appConfigContext.tsx

@@ -50,6 +50,7 @@ export const DEFAULT_APP_CONFIG = {
         showAcrylInfo: false,
         showAccessManagement: false,
         nestedDomainsEnabled: true,
+        platformBrowseV2: false,
     },
 };
 

datahub-web-react/src/graphql/app.graphql

@@ -65,6 +65,7 @@ query appConfig {
             showAcrylInfo
             showAccessManagement
             nestedDomainsEnabled
+            platformBrowseV2
         }
     }
 }

metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java

@@ -235,6 +235,30 @@ public BrowseResultV2 browseV2(
     return _entitySearchService.browseV2(entityName, path, filter, input, start, count);
   }
 
+  /**
+   * Gets browse V2 snapshot of a given path
+   *
+   * @param entityNames entities being browsed
+   * @param path path being browsed
+   * @param filter browse filter
+   * @param input search query
+   * @param start start offset of first group
+   * @param count max number of results requested
+   * @throws RemoteInvocationException
+   */
+  @Nonnull
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entityNames,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count,
+      @Nonnull Authentication authentication) {
+    // TODO: cache browseV2 results
+    return _entitySearchService.browseV2(entityNames, path, filter, input, start, count);
+  }
+
   @SneakyThrows
   @Deprecated
   public void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication)

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java

@@ -210,6 +210,18 @@ public BrowseResultV2 browseV2(
     return esBrowseDAO.browseV2(entityName, path, filter, input, start, count);
   }
 
+  @Nonnull
+  @Override
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entityNames,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count) {
+    return esBrowseDAO.browseV2(entityNames, path, filter, input, start, count);
+  }
+
   @Nonnull
   @Override
   public List<String> getBrowsePaths(@Nonnull String entityName, @Nonnull Urn urn) {

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java

@@ -427,6 +427,44 @@ public BrowseResultV2 browseV2(
     }
   }
 
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entities,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count) {
+    try {
+      final SearchResponse groupsResponse;
+
+      try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) {
+        final String finalInput = input.isEmpty() ? "*" : input;
+        groupsResponse =
+            client.search(
+                constructGroupsSearchRequestBrowseAcrossEntities(
+                    entities, path, filter, finalInput),
+                RequestOptions.DEFAULT);
+      }
+
+      final BrowseGroupsResultV2 browseGroupsResult =
+          extractGroupsResponseV2(groupsResponse, path, start, count);
+      final int numGroups = browseGroupsResult.getTotalGroups();
+
+      return new BrowseResultV2()
+          .setMetadata(
+              new BrowseResultMetadata()
+                  .setTotalNumEntities(browseGroupsResult.getTotalNumEntities())
+                  .setPath(path))
+          .setGroups(new BrowseResultGroupV2Array(browseGroupsResult.getGroups()))
+          .setNumGroups(numGroups)
+          .setFrom(start)
+          .setPageSize(count);
+    } catch (Exception e) {
+      log.error("Browse Across Entities query failed: " + e.getMessage());
+      throw new ESQueryException("Browse Across Entities query failed: ", e);
+    }
+  }
+
   @Nonnull
   private SearchRequest constructGroupsSearchRequestV2(
       @Nonnull String entityName,
@@ -448,6 +486,33 @@ private SearchRequest constructGroupsSearchRequestV2(
     return searchRequest;
   }
 
+  @Nonnull
+  private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities(
+      @Nonnull List<String> entities,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input) {
+
+    List<EntitySpec> entitySpecs =
+        entities.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList());
+
+    String[] indexArray =
+        entities.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new);
+
+    final SearchRequest searchRequest = new SearchRequest(indexArray);
+    final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+    searchSourceBuilder.size(0);
+    searchSourceBuilder.query(
+        buildQueryStringBrowseAcrossEntities(
+            entitySpecs,
+            path,
+            SearchUtil.transformFilterForEntities(filter, indexConvention),
+            input));
+    searchSourceBuilder.aggregation(buildAggregationsV2(path));
+    searchRequest.source(searchSourceBuilder);
+    return searchRequest;
+  }
+
   /**
    * Extracts the name of group from path.
    *
@@ -494,6 +559,32 @@ private QueryBuilder buildQueryStringV2(
     return queryBuilder;
   }
 
+  @Nonnull
+  private QueryBuilder buildQueryStringBrowseAcrossEntities(
+      @Nonnull List<EntitySpec> entitySpecs,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input) {
+    final int browseDepthVal = getPathDepthV2(path);
+
+    final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
+
+    QueryBuilder query =
+        SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
+            .getQuery(input, false);
+    queryBuilder.must(query);
+
+    if (!path.isEmpty()) {
+      queryBuilder.filter(QueryBuilders.matchQuery(BROWSE_PATH_V2, path));
+    }
+
+    queryBuilder.filter(QueryBuilders.rangeQuery(BROWSE_PATH_V2_DEPTH).gt(browseDepthVal));
+
+    queryBuilder.filter(SearchRequestHandler.getFilterQuery(filter));
+
+    return queryBuilder;
+  }
+
   @Nonnull
   private AggregationBuilder buildAggregationsV2(@Nonnull String path) {
     final String currentLevel = ESUtils.escapeReservedCharacters(path) + "␟.*";

metadata-service/configuration/src/main/resources/application.yml

@@ -317,6 +317,7 @@ featureFlags:
   showAccessManagement: ${SHOW_ACCESS_MANAGEMENT:false} #Whether we should show AccessManagement tab in the datahub UI.
   showSearchFiltersV2: ${SHOW_SEARCH_FILTERS_V2:true} # Enables showing the search filters V2 experience.
   showBrowseV2: ${SHOW_BROWSE_V2:true} # Enables showing the browse v2 sidebar experience.
+  platformBrowseV2: ${PLATFORM_BROWSE_V2:false} # Enables the platform browse experience, instead of the entity-oriented browse default.
   preProcessHooks:
     uiEnabled: ${PRE_PROCESS_HOOKS_UI_ENABLED:true} # Circumvents Kafka for processing index updates for UI changes sourced from GraphQL to avoid processing delays
   showAcrylInfo: ${SHOW_ACRYL_INFO:false} # Show different CTAs within DataHub around moving to Managed DataHub. Set to true for the demo site.

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java

@@ -153,6 +153,28 @@ public BrowseResultV2 browseV2(
       @Nonnull Authentication authentication)
       throws RemoteInvocationException;
 
+  /**
+   * Gets browse snapshot of a given path
+   *
+   * @param entityNames entities being browsed
+   * @param path path being browsed
+   * @param filter browse filter
+   * @param input search query
+   * @param start start offset of first group
+   * @param count max number of results requested
+   * @throws RemoteInvocationException
+   */
+  @Nonnull
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entityNames,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count,
+      @Nonnull Authentication authentication)
+      throws RemoteInvocationException;
+
   @Deprecated
   public void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication)
       throws RemoteInvocationException;

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java

@@ -381,6 +381,20 @@ public BrowseResultV2 browseV2(
     throw new NotImplementedException("BrowseV2 is not implemented in Restli yet");
   }
 
+  @Nonnull
+  @Override
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entityNames,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count,
+      @Nonnull Authentication authentication)
+      throws RemoteInvocationException {
+    throw new NotImplementedException("BrowseV2 is not implemented in Restli yet");
+  }
+
   public void update(@Nonnull final Entity entity, @Nonnull final Authentication authentication)
       throws RemoteInvocationException {
     EntitiesDoIngestRequestBuilder requestBuilder =

metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java

@@ -207,6 +207,25 @@ public BrowseResultV2 browseV2(
       int start,
       int count);
 
+  /**
+   * Gets browse snapshot of a given path
+   *
+   * @param entityNames set of entities being browsed
+   * @param path path being browsed
+   * @param filter browse filter
+   * @param input search query
+   * @param start start offset of first group
+   * @param count max number of results requested
+   */
+  @Nonnull
+  public BrowseResultV2 browseV2(
+      @Nonnull List<String> entityNames,
+      @Nonnull String path,
+      @Nullable Filter filter,
+      @Nonnull String input,
+      int start,
+      int count);
+
   /**
    * Gets a list of paths for a given urn.
    *