feat(elasticsearch): Elasticsearch improvements

build.gradle

@@ -9,6 +9,7 @@ buildscript {
   ext.neo4jVersion = '4.4.9'
   ext.graphQLJavaVersion = '19.0'
   ext.testContainersVersion = '1.17.4'
+  ext.elasticsearchVersion = '7.10.2'
   ext.jacksonVersion = '2.13.4'
   ext.jettyVersion = '9.4.46.v20220331'
   ext.playVersion = '2.8.18'
@@ -85,8 +86,8 @@ project.ext.externalDependency = [
     'dropwizardMetricsJmx': 'io.dropwizard.metrics:metrics-jmx:4.2.3',
     'ebean': 'io.ebean:ebean:11.33.3',
     'ebeanAgent': 'io.ebean:ebean-agent:11.27.1',
-    'elasticSearchRest': 'org.elasticsearch.client:elasticsearch-rest-high-level-client:7.9.3',
-    'elasticSearchTransport': 'org.elasticsearch.client:transport:7.9.3',
+    'elasticSearchRest': 'org.elasticsearch.client:elasticsearch-rest-high-level-client:' + elasticsearchVersion,
+    'elasticSearchTransport': 'org.elasticsearch.client:transport:' + elasticsearchVersion,
     'findbugsAnnotations': 'com.google.code.findbugs:annotations:3.0.1',
     'graphqlJava': 'com.graphql-java:graphql-java:' + graphQLJavaVersion,
     'graphqlJavaScalars': 'com.graphql-java:graphql-java-extended-scalars:' + graphQLJavaVersion,

datahub-frontend/build.gradle

@@ -73,14 +73,13 @@ distributions {
 }
 
 docker {
-  name "${docker_registry}/${docker_repo}:${version}"
+  name "${docker_registry}/${docker_repo}:v${version}"
+  version "v${version}"
   dockerfile file("${rootProject.projectDir}/docker/${docker_dir}/Dockerfile")
   files fileTree(rootProject.projectDir) {
     include 'docker/monitoring/*'
     include "docker/${docker_dir}/*"
   }
-  tag("-${version}", "${docker_registry}/${docker_repo}:v${version}")
-  tag("-version", "${docker_registry}/${docker_repo}:v${version}")
 
   // platform('linux/arm64', 'linux/amd64')
   buildx(true)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java

@@ -57,7 +57,7 @@ public CompletableFuture<ListAccessTokenResult> get(DataFetchingEnvironment envi
 
           final SearchResult searchResult = _entityClient.search(Constants.ACCESS_TOKEN_ENTITY_NAME, "",
               buildFilter(filters, Collections.emptyList()), sortCriterion, start, count,
-              getAuthentication(environment));
+              getAuthentication(environment), true);
 
           final List<AccessTokenMetadata> tokens = searchResult.getEntities().stream().map(entity -> {
             final AccessTokenMetadata metadata = new AccessTokenMetadata();

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/ListDomainsResolver.java

@@ -62,7 +62,8 @@ public CompletableFuture<ListDomainsResult> get(final DataFetchingEnvironment en
                   new SortCriterion().setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
                   start,
                   count,
-                  context.getAuthentication());
+                  context.getAuthentication(),
+                  true);
 
           // Now that we have entities we can bind this to a result.
           final ListDomainsResult result = new ListDomainsResult();

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java

@@ -55,10 +55,11 @@ public CompletableFuture<ListGroupsResult> get(final DataFetchingEnvironment env
           // First, get all group Urns.
           final SearchResult gmsResult =
               _entityClient.search(CORP_GROUP_ENTITY_NAME,
-                  query,
-                  null,
-                  new SortCriterion().setField(CORP_GROUP_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
-                  start, count, context.getAuthentication());
+                      query,
+                      null,
+                      new SortCriterion().setField(CORP_GROUP_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
+                      start, count, context.getAuthentication(),
+                      true);
 
           // Then, get hydrate all groups.
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(CORP_GROUP_ENTITY_NAME,

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java

@@ -65,13 +65,14 @@ public CompletableFuture<ListSecretsResult> get(final DataFetchingEnvironment en
         try {
           // First, get all secrets
           final SearchResult gmsResult = _entityClient.search(
-              Constants.SECRETS_ENTITY_NAME,
-              query,
-              null,
-              new SortCriterion().setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
-              start,
-              count,
-              context.getAuthentication());
+                  Constants.SECRETS_ENTITY_NAME,
+                  query,
+                  null,
+                  new SortCriterion().setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
+                  start,
+                  count,
+                  context.getAuthentication(),
+                  true);
 
           // Then, resolve all secrets
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java

@@ -60,7 +60,8 @@ public CompletableFuture<ListIngestionSourcesResult> get(final DataFetchingEnvir
               Collections.emptyMap(),
               start,
               count,
-              context.getAuthentication());
+              context.getAuthentication(),
+                  true);
 
           // Then, resolve all ingestion sources
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java

@@ -51,7 +51,7 @@ public CompletableFuture<ListPostsResult> get(final DataFetchingEnvironment envi
 
         // First, get all Post Urns.
         final SearchResult gmsResult = _entityClient.search(POST_ENTITY_NAME, query, null, sortCriterion, start, count,
-            context.getAuthentication());
+            context.getAuthentication(), true);
 
         // Then, get and hydrate all Posts.
         final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(POST_ENTITY_NAME,

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java

@@ -50,7 +50,7 @@ public CompletableFuture<ListRolesResult> get(final DataFetchingEnvironment envi
         // First, get all role Urns.
         final SearchResult gmsResult =
             _entityClient.search(DATAHUB_ROLE_ENTITY_NAME, query, Collections.emptyMap(), start, count,
-                context.getAuthentication());
+                context.getAuthentication(), true);
 
         // Then, get and hydrate all users.
         final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(DATAHUB_ROLE_ENTITY_NAME,

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

@@ -50,21 +50,21 @@ public static CompletableFuture<AutoCompleteMultipleResults> batchGetAutocomplet
             input.getField(),
             input.getQuery(),
             input.getFilters(),
-            input.getLimit()) + " "
-            + e.getMessage());
+            input.getLimit()), e);
         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(autoCompletesFuture.stream()
-              .map(CompletableFuture::join)
-              .filter(
-                  autoCompleteResultForEntity ->
-                      autoCompleteResultForEntity.getSuggestions() != null && autoCompleteResultForEntity.getSuggestions().size() > 0
-              )
-              .collect(Collectors.toList()));
+          List<AutoCompleteResultForEntity> suggestions = autoCompletesFuture.stream()
+                  .map(CompletableFuture::join)
+                  .filter(
+                          autoCompleteResultForEntity ->
+                                  autoCompleteResultForEntity.getSuggestions() != null && autoCompleteResultForEntity.getSuggestions().size() > 0
+                  )
+                  .collect(Collectors.toList());
+          result.setSuggestions(suggestions);
           return result;
         });
   }

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

@@ -45,7 +45,7 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
             input.getQuery(), input.getFilters(), input.getOrFilters(), start, count);
         return UrnSearchResultsMapper.map(
             _entityClient.search(entityName, sanitizedQuery, ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()), null, start,
-                count, ResolverUtils.getAuthentication(environment)));
+                count, ResolverUtils.getAuthentication(environment), true));
       } catch (Exception e) {
         log.error("Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}",
             input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), start, count);

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

@@ -53,6 +53,7 @@ private SearchUtils() {
           EntityType.DATASET,
           EntityType.DASHBOARD,
           EntityType.CHART,
+          EntityType.CONTAINER,
           EntityType.MLMODEL,
           EntityType.MLMODEL_GROUP,
           EntityType.MLFEATURE_TABLE,

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java

@@ -58,7 +58,8 @@ public CompletableFuture<ListTestsResult> get(final DataFetchingEnvironment envi
               Collections.emptyMap(),
               start,
               count,
-              context.getAuthentication());
+              context.getAuthentication(),
+                  true);
 
           // Now that we have entities we can bind this to a result.
           final ListTestsResult result = new ListTestsResult();

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java

@@ -53,7 +53,8 @@ public CompletableFuture<ListUsersResult> get(final DataFetchingEnvironment envi
         try {
           // First, get all policy Urns.
           final SearchResult gmsResult =
-              _entityClient.search(CORP_USER_ENTITY_NAME, query, Collections.emptyMap(), start, count, context.getAuthentication());
+              _entityClient.search(CORP_USER_ENTITY_NAME, query, Collections.emptyMap(), start, count,
+                      context.getAuthentication(), true);
 
           // Then, get hydrate all users.
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(CORP_USER_ENTITY_NAME,

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListGlobalViewsResolver.java

@@ -67,13 +67,14 @@ public CompletableFuture<ListViewsResult> get(final DataFetchingEnvironment envi
       try {
 
         final SearchResult gmsResult = _entityClient.search(
-            Constants.DATAHUB_VIEW_ENTITY_NAME,
-            query,
-            buildFilters(),
-            DEFAULT_SORT_CRITERION,
-            start,
-            count,
-            context.getAuthentication());
+                Constants.DATAHUB_VIEW_ENTITY_NAME,
+                query,
+                buildFilters(),
+                DEFAULT_SORT_CRITERION,
+                start,
+                count,
+                context.getAuthentication(),
+                true);
 
         final ListViewsResult result = new ListViewsResult();
         result.setStart(gmsResult.getFrom());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/view/ListMyViewsResolver.java

@@ -75,7 +75,8 @@ public CompletableFuture<ListViewsResult> get(final DataFetchingEnvironment envi
                 DEFAULT_SORT_CRITERION,
                 start,
                 count,
-                context.getAuthentication());
+                context.getAuthentication(),
+                true);
 
         final ListViewsResult result = new ListViewsResult();
         result.setStart(gmsResult.getFrom());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java

@@ -139,11 +139,12 @@ public SearchResults search(@Nonnull String query,
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
         final SearchResult searchResult = _entityClient.search(
             "chart",
-            query,
-            facetFilters,
-            start,
-            count,
-            context.getAuthentication()
+                query,
+                facetFilters,
+                start,
+                count,
+                context.getAuthentication(),
+                true
         );
         return UrnSearchResultsMapper.map(searchResult);
     }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java

@@ -3,13 +3,22 @@
 import com.google.common.collect.ImmutableSet;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
 import com.linkedin.datahub.graphql.generated.Container;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.FacetFilterInput;
+import com.linkedin.datahub.graphql.generated.SearchResults;
+import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
+import com.linkedin.datahub.graphql.types.SearchableEntityType;
 import com.linkedin.datahub.graphql.types.container.mappers.ContainerMapper;
+import com.linkedin.datahub.graphql.types.mappers.AutoCompleteResultsMapper;
+import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.AutoCompleteResult;
+import com.linkedin.metadata.search.SearchResult;
 import graphql.execution.DataFetcherResult;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
@@ -20,9 +29,11 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 
-public class ContainerType implements com.linkedin.datahub.graphql.types.EntityType<Container, String> {
+public class ContainerType implements SearchableEntityType<Container, String>,
+        com.linkedin.datahub.graphql.types.EntityType<Container, String> {
 
   static final Set<String> ASPECTS_TO_FETCH = ImmutableSet.of(
       Constants.DATA_PLATFORM_INSTANCE_ASPECT_NAME,
@@ -38,6 +49,9 @@ public class ContainerType implements com.linkedin.datahub.graphql.types.EntityT
       Constants.DOMAINS_ASPECT_NAME,
       Constants.DEPRECATION_ASPECT_NAME
   );
+
+  private static final Set<String> FACET_FIELDS = ImmutableSet.of("origin", "platform");
+  private static final String ENTITY_NAME = "container";
   private final EntityClient _entityClient;
 
   public ContainerType(final EntityClient entityClient)  {
@@ -95,4 +109,27 @@ private Urn getUrn(final String urnStr) {
       throw new RuntimeException(String.format("Failed to convert urn string %s into Urn", urnStr));
     }
   }
+
+  @Override
+  public SearchResults search(@Nonnull String query,
+                              @Nullable List<FacetFilterInput> filters,
+                              int start,
+                              int count,
+                              @Nonnull final QueryContext context) throws Exception {
+    final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
+    final SearchResult searchResult = _entityClient.search(ENTITY_NAME, query, facetFilters, start, count,
+            context.getAuthentication(), true);
+    return UrnSearchResultsMapper.map(searchResult);
+  }
+
+  @Override
+  public AutoCompleteResults autoComplete(@Nonnull String query,
+                                          @Nullable String field,
+                                          @Nullable List<FacetFilterInput> filters,
+                                          int limit,
+                                          @Nonnull final QueryContext context) throws Exception {
+    final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
+    final AutoCompleteResult result = _entityClient.autoComplete(ENTITY_NAME, query, facetFilters, limit, context.getAuthentication());
+    return AutoCompleteResultsMapper.map(result);
+  }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java

@@ -103,7 +103,7 @@ public SearchResults search(@Nonnull String query,
                                 @Nonnull final QueryContext context) throws Exception {
         final SearchResult
             searchResult = _entityClient.search("corpGroup", query, Collections.emptyMap(), start, count,
-            context.getAuthentication());
+            context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java

@@ -106,7 +106,7 @@ public SearchResults search(@Nonnull String query,
                                 int count,
                                 @Nonnull final QueryContext context) throws Exception {
         final SearchResult searchResult = _entityClient.search("corpuser", query, Collections.emptyMap(), start, count,
-            context.getAuthentication());
+            context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java

@@ -139,7 +139,8 @@ public SearchResults search(@Nonnull String query,
                                 int count,
                                 @Nonnull QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
-        final SearchResult searchResult = _entityClient.search("dashboard", query, facetFilters, start, count, context.getAuthentication());
+        final SearchResult searchResult = _entityClient.search("dashboard", query, facetFilters, start, count,
+                context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java

@@ -134,7 +134,8 @@ public SearchResults search(@Nonnull String query,
                                 int count,
                                 @Nonnull final QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
-        final SearchResult searchResult = _entityClient.search("dataFlow", query, facetFilters, start, count, context.getAuthentication());
+        final SearchResult searchResult = _entityClient.search("dataFlow", query, facetFilters, start, count,
+                context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java

@@ -135,7 +135,7 @@ public SearchResults search(@Nonnull String query,
                                 @Nonnull final QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
         final SearchResult searchResult = _entityClient.search(
-            "dataJob", query, facetFilters, start, count, context.getAuthentication());
+            "dataJob", query, facetFilters, start, count, context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java

@@ -154,7 +154,8 @@ public SearchResults search(@Nonnull String query,
                                 int count,
                                 @Nonnull final QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
-        final SearchResult searchResult = _entityClient.search(ENTITY_NAME, query, facetFilters, start, count, context.getAuthentication());
+        final SearchResult searchResult = _entityClient.search(ENTITY_NAME, query, facetFilters, start, count,
+                context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/GlossaryTermType.java

@@ -112,7 +112,7 @@ public SearchResults search(@Nonnull String query,
                                 @Nonnull final QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
         final SearchResult searchResult = _entityClient.search(
-            "glossaryTerm", query, facetFilters, start, count, context.getAuthentication());
+            "glossaryTerm", query, facetFilters, start, count, context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/MLFeatureTableType.java

@@ -97,7 +97,8 @@ public SearchResults search(@Nonnull String query,
                                 int count,
                                 @Nonnull final QueryContext context) throws Exception {
         final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
-        final SearchResult searchResult = _entityClient.search("mlFeatureTable", query, facetFilters, start, count, context.getAuthentication());
+        final SearchResult searchResult = _entityClient.search("mlFeatureTable", query, facetFilters, start, count,
+                context.getAuthentication(), true);
         return UrnSearchResultsMapper.map(searchResult);
     }