Skip to content

Commit

Permalink
feat(elasticsearch): Elasticsearch improvements (datahub-project#6894)
Browse files Browse the repository at this point in the history
  • Loading branch information
david-leifker authored and danielcmessias committed Feb 1, 2023
1 parent 1a7fe36 commit 61339fe
Show file tree
Hide file tree
Showing 246 changed files with 7,111 additions and 2,115 deletions.
5 changes: 3 additions & 2 deletions build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 2 additions & 3 deletions datahub-frontend/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
});
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ private SearchUtils() {
EntityType.DATASET,
EntityType.DASHBOARD,
EntityType.CHART,
EntityType.CONTAINER,
EntityType.MLMODEL,
EntityType.MLMODEL_GROUP,
EntityType.MLFEATURE_TABLE,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand All @@ -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) {
Expand Down Expand Up @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
Loading

0 comments on commit 61339fe

Please sign in to comment.