feat(autocomplete): Show recent searches + improved autocomplete

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -12,6 +12,8 @@
 import com.linkedin.datahub.graphql.generated.AggregationMetadata;
 import com.linkedin.datahub.graphql.generated.Aspect;
 import com.linkedin.datahub.graphql.generated.Assertion;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResultForEntity;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
 import com.linkedin.datahub.graphql.generated.BrowseResults;
 import com.linkedin.datahub.graphql.generated.Chart;
 import com.linkedin.datahub.graphql.generated.ChartInfo;
@@ -730,7 +732,20 @@ private void configureGenericEntityResolvers(final RuntimeWiring.Builder builder
                         (env) -> ((ListDomainsResult) env.getSource()).getDomains().stream()
                             .map(Domain::getUrn)
                             .collect(Collectors.toList())))
+            )
+            .type("AutoCompleteResults", typeWiring -> typeWiring
+                .dataFetcher("entities",
+                    new EntityTypeBatchResolver(
+                        new ArrayList<>(entityTypes),
+                        (env) -> ((AutoCompleteResults) env.getSource()).getEntities()))
+            )
+            .type("AutoCompleteResultForEntity", typeWiring -> typeWiring
+                .dataFetcher("entities",
+                    new EntityTypeBatchResolver(
+                        new ArrayList<>(entityTypes),
+                        (env) -> ((AutoCompleteResultForEntity) env.getSource()).getEntities()))
             );
+        ;
     }
 
     /**

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

@@ -18,6 +18,7 @@
 import org.slf4j.LoggerFactory;
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*;
 import static org.apache.commons.lang3.StringUtils.isBlank;
 
 /**
@@ -27,11 +28,9 @@ public class AutoCompleteForMultipleResolver implements DataFetcher<CompletableF
 
     private static final Logger _logger = LoggerFactory.getLogger(AutoCompleteForMultipleResolver.class.getName());
 
-    private final List<SearchableEntityType<?>> _searchableEntities;
     private final Map<EntityType, SearchableEntityType<?>> _typeToEntity;
 
     public AutoCompleteForMultipleResolver(@Nonnull final List<SearchableEntityType<?>> searchableEntities) {
-        _searchableEntities = searchableEntities;
         _typeToEntity = searchableEntities.stream().collect(Collectors.toMap(
             SearchableEntityType::type,
             entity -> entity
@@ -51,10 +50,18 @@ public CompletableFuture<AutoCompleteMultipleResults> get(DataFetchingEnvironmen
 
         List<EntityType> types = input.getTypes();
         if (types != null && types.size() > 0) {
-            return AutocompleteUtils.batchGetAutocompleteResults(types.stream().map(type -> _typeToEntity.get(type)).collect(
-                Collectors.toList()), sanitizedQuery, input, environment);
+            return AutocompleteUtils.batchGetAutocompleteResults(
+                types.stream().map(_typeToEntity::get).collect(Collectors.toList()),
+                sanitizedQuery,
+                input,
+                environment);
         }
 
-        return AutocompleteUtils.batchGetAutocompleteResults(_searchableEntities, sanitizedQuery, input, environment);
+        // By default, autocomplete only against the set of Searchable Entity Types.
+        return AutocompleteUtils.batchGetAutocompleteResults(
+            AUTO_COMPLETE_ENTITY_TYPES.stream().map(_typeToEntity::get).collect(Collectors.toList()),
+            sanitizedQuery,
+            input,
+            environment);
     }
 }

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

@@ -7,7 +7,7 @@
 import com.linkedin.datahub.graphql.types.SearchableEntityType;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
@@ -30,35 +30,35 @@ public static CompletableFuture<AutoCompleteMultipleResults> batchGetAutocomplet
   ) {
     final int limit = input.getLimit() != null ? input.getLimit() : DEFAULT_LIMIT;
 
-    final CompletableFuture<AutoCompleteResultForEntity>[] autoCompletesFuture = entities.stream().map(entity -> {
-      return CompletableFuture.supplyAsync(() -> {
-        try {
-          final AutoCompleteResults searchResult = entity.autoComplete(
-              sanitizedQuery,
-              input.getField(),
-              input.getFilters(),
-              limit,
-              environment.getContext()
-          );
-          final AutoCompleteResultForEntity autoCompleteResultForEntity =
-              new AutoCompleteResultForEntity(entity.type(), searchResult.getSuggestions());
-          return autoCompleteResultForEntity;
-        } catch (Exception e) {
-          _logger.error("Failed to execute autocomplete all: "
-              + String.format("field %s, query %s, filters: %s, limit: %s",
-              input.getField(),
-              input.getQuery(),
-              input.getFilters(),
-              input.getLimit()) + " "
-              + e.getMessage());
-          return new AutoCompleteResultForEntity(entity.type(), new ArrayList<>());
-        }
-      });
-    }).toArray(CompletableFuture[]::new);
-    return CompletableFuture.allOf(autoCompletesFuture)
+    final List<CompletableFuture<AutoCompleteResultForEntity>> autoCompletesFuture = entities.stream().map(entity -> CompletableFuture.supplyAsync(() -> {
+      try {
+        final AutoCompleteResults searchResult = entity.autoComplete(
+            sanitizedQuery,
+            input.getField(),
+            input.getFilters(),
+            limit,
+            environment.getContext()
+        );
+        return new AutoCompleteResultForEntity(
+            entity.type(),
+            searchResult.getSuggestions(),
+            searchResult.getEntities()
+        );
+      } catch (Exception e) {
+        _logger.error("Failed to execute autocomplete all: "
+            + String.format("field %s, query %s, filters: %s, limit: %s",
+            input.getField(),
+            input.getQuery(),
+            input.getFilters(),
+            input.getLimit()) + " "
+            + e.getMessage());
+        return new AutoCompleteResultForEntity(entity.type(), Collections.emptyList(), Collections.emptyList());
+      }
+    })).collect(Collectors.toList());
+    return CompletableFuture.allOf(autoCompletesFuture.toArray(new CompletableFuture[0]))
         .thenApplyAsync((res) -> {
           AutoCompleteMultipleResults result = new AutoCompleteMultipleResults(sanitizedQuery, new ArrayList<>());
-          result.setSuggestions(Arrays.stream(autoCompletesFuture)
+          result.setSuggestions(autoCompletesFuture.stream()
               .map(CompletableFuture::join)
               .filter(
                   autoCompleteResultForEntity ->

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

@@ -9,9 +9,41 @@ public class SearchUtils {
   private SearchUtils() {
   }
 
+  /**
+   * Entities that are searched by default in Search Across Entities
+   */
   public static final List<EntityType> SEARCHABLE_ENTITY_TYPES =
-      ImmutableList.of(EntityType.DATASET, EntityType.DASHBOARD, EntityType.CHART, EntityType.MLMODEL,
-          EntityType.MLMODEL_GROUP, EntityType.MLFEATURE_TABLE, EntityType.DATA_FLOW, EntityType.DATA_JOB,
-          EntityType.GLOSSARY_TERM, EntityType.TAG, EntityType.CORP_USER, EntityType.CORP_GROUP, EntityType.CONTAINER,
+      ImmutableList.of(
+          EntityType.DATASET,
+          EntityType.DASHBOARD,
+          EntityType.CHART,
+          EntityType.MLMODEL,
+          EntityType.MLMODEL_GROUP,
+          EntityType.MLFEATURE_TABLE,
+          EntityType.DATA_FLOW,
+          EntityType.DATA_JOB,
+          EntityType.GLOSSARY_TERM,
+          EntityType.TAG,
+          EntityType.CORP_USER,
+          EntityType.CORP_GROUP,
+          EntityType.CONTAINER,
           EntityType.DOMAIN);
+
+  /**
+   * Entities that are part of autocomplete by default in Auto Complete Across Entities
+   */
+  public static final List<EntityType> AUTO_COMPLETE_ENTITY_TYPES =
+      ImmutableList.of(
+          EntityType.DATASET,
+          EntityType.DASHBOARD,
+          EntityType.CHART,
+          EntityType.MLMODEL,
+          EntityType.MLMODEL_GROUP,
+          EntityType.MLFEATURE_TABLE,
+          EntityType.DATA_FLOW,
+          EntityType.DATA_JOB,
+          EntityType.GLOSSARY_TERM,
+          EntityType.TAG,
+          EntityType.CORP_USER,
+          EntityType.CORP_GROUP);
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/AutoCompleteResultsMapper.java

@@ -1,8 +1,10 @@
 package com.linkedin.datahub.graphql.types.mappers;
 
 import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
 import com.linkedin.metadata.query.AutoCompleteResult;
 
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 
 
@@ -19,6 +21,8 @@ public AutoCompleteResults apply(@Nonnull final AutoCompleteResult input) {
         final AutoCompleteResults result = new AutoCompleteResults();
         result.setQuery(input.getQuery());
         result.setSuggestions(input.getSuggestions());
+        result.setEntities(input.getEntities().stream().map(entity -> UrnToEntityMapper.map(entity.getUrn())).collect(
+            Collectors.toList()));
         return result;
     }
 }

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

@@ -42,7 +42,12 @@ enum ScenarioType {
   """
   Recommendations to show on an Entity Profile page
   """
-  ENTITY_PROFILE
+  ENTITY_PROFILE,
+
+  """
+  Recommendations to show on the search bar when clicked
+  """
+  SEARCH_BAR
 }
 
 """

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

@@ -405,6 +405,11 @@ type AutoCompleteResultForEntity {
   The autocompletion results for specified entity type
   """
   suggestions: [String!]!
+
+  """
+  A list of entities to render in autocomplete
+  """
+  entities: [Entity!]!
 }
 
 """
@@ -435,6 +440,11 @@ type AutoCompleteResults {
 	The autocompletion results
 	"""
 	suggestions: [String!]!
+
+  """
+  A list of entities to render in autocomplete
+  """
+  entities: [Entity!]!
 }
 
 """

datahub-web-react/src/Mocks.tsx

@@ -1594,7 +1594,8 @@ export const mocks = [
                     suggestions: [
                         {
                             type: EntityType.Dataset,
-                            suggestions: ['The Great Test Dataset', 'Some other test'],
+                            suggestions: ['The Great Test Dataset', 'Some Other Dataset'],
+                            entities: [dataset1, dataset2],
                         },
                     ],
                 },
@@ -1618,7 +1619,8 @@ export const mocks = [
                     suggestions: [
                         {
                             type: EntityType.Dataset,
-                            suggestions: ['The Great Test Dataset', 'Some other test'],
+                            suggestions: ['The Great Test Dataset', 'Some Other Dataset'],
+                            entities: [dataset1, dataset2],
                         },
                     ],
                 },
@@ -1640,6 +1642,7 @@ export const mocks = [
                 autoComplete: {
                     query: 'j',
                     suggestions: ['jjoyce'],
+                    entities: [user1],
                 },
             },
         },

datahub-web-react/src/app/entity/user/User.tsx

@@ -41,9 +41,9 @@ export class UserEntity implements Entity<CorpUser> {
 
     getPathName: () => string = () => 'user';
 
-    getEntityName = () => 'User';
+    getEntityName = () => 'Person';
 
-    getCollectionName: () => string = () => 'Users';
+    getCollectionName: () => string = () => 'People';
 
     renderProfile: (urn: string) => JSX.Element = (_) => <UserProfile />;
 

datahub-web-react/src/app/home/__tests__/HomePage.test.tsx

@@ -61,8 +61,8 @@ describe('HomePage', () => {
         await waitFor(() => expect(searchInput).toBeInTheDocument());
         fireEvent.change(searchInput, { target: { value: 't' } });
 
-        await waitFor(() => expect(queryAllByText('The Great Test Dataset').length).toBeGreaterThanOrEqual(1));
-        expect(queryAllByText('Some other test').length).toBeGreaterThanOrEqual(1);
+        await waitFor(() => expect(queryAllByText('he Great Test Dataset').length).toBeGreaterThanOrEqual(1));
+        expect(queryAllByText('Some Other Dataset').length).toBeGreaterThanOrEqual(1);
     });
 
     it('renders search suggestions', async () => {