From 27392f9f81e35c95e24588ad5c26235ec0f4b2d3 Mon Sep 17 00:00:00 2001 From: Indy Prentice Date: Mon, 24 Jul 2023 15:39:44 -0300 Subject: [PATCH] refactor(search): Support searching multiple entities in search() as in scroll() (#8461) Co-authored-by: Indy Prentice --- .../metadata/client/JavaEntityClient.java | 4 +- .../metadata/search/SearchService.java | 6 +- .../AllEntitiesSearchAggregator.java | 2 +- .../client/CachingEntitySearchService.java | 14 +-- .../elasticsearch/ElasticSearchService.java | 10 +- .../elasticsearch/query/ESSearchDAO.java | 16 +-- .../com/linkedin/metadata/ESTestUtils.java | 15 ++- .../ElasticSearchServiceTest.java | 19 ++-- .../fixtures/SampleDataFixtureTests.java | 98 +++++++++++-------- .../SearchLineageDataFixtureTests.java | 6 +- .../hook/siblings/SiblingAssociationHook.java | 2 +- .../siblings/SiblingAssociationHookTest.java | 2 +- .../boot/steps/RestoreGlossaryIndices.java | 4 +- .../steps/RestoreGlossaryIndicesTest.java | 9 +- .../resources/entity/EntityResource.java | 2 +- .../metadata/search/EntitySearchService.java | 8 +- .../src/main/java/mock/MockEntitySpec.java | 4 +- 17 files changed, 125 insertions(+), 96 deletions(-) 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 index 911ab993e5789..e816700e583c7 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java @@ -277,7 +277,7 @@ public SearchResult search(@Nonnull String entity, @Nonnull String input, @Nullable SearchFlags searchFlags) throws RemoteInvocationException { - return ValidationUtils.validateSearchResult(_entitySearchService.search(entity, input, newFilter(requestFilters), + return ValidationUtils.validateSearchResult(_entitySearchService.search(List.of(entity), input, newFilter(requestFilters), null, start, count, searchFlags), _entityService); } @@ -329,7 +329,7 @@ public SearchResult search( @Nullable SearchFlags searchFlags) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( - _entitySearchService.search(entity, input, filter, sortCriterion, start, count, searchFlags), _entityService); + _entitySearchService.search(List.of(entity), input, filter, sortCriterion, start, count, searchFlags), _entityService); } @Nonnull 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 index f3c0d14bddac3..a14ea8d36b0c1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java @@ -44,7 +44,7 @@ public Map docCountPerEntity(@Nonnull List entityNames) { * Gets a list of documents that match given search request. The results are aggregated and filters are applied to the * search hits and not the aggregation results. * - * @param entityName name of the entity + * @param entityNames names of the entity * @param input the search input text * @param postFilters the request map with fields and values as filters to be applied to search hits * @param sortCriterion {@link SortCriterion} to be applied to search results @@ -54,10 +54,10 @@ public Map docCountPerEntity(@Nonnull List entityNames) { * @return a {@link SearchResult} that contains a list of matched documents and related search result metadata */ @Nonnull - public SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + public SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags) { SearchResult result = - _cachingEntitySearchService.search(entityName, input, postFilters, sortCriterion, from, size, searchFlags, null); + _cachingEntitySearchService.search(entityNames, input, postFilters, sortCriterion, from, size, searchFlags, null); try { return result.copy().setEntities(new SearchEntityArray(_searchRanker.rank(result.getEntities()))); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java b/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java index ee93edaf2480c..1af94141366e1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java @@ -171,7 +171,7 @@ private Map getSearchResultsForEachEntity(@Nonnull List new Pair<>(entity, - _cachingEntitySearchService.search(entity, input, postFilters, sortCriterion, queryFrom, querySize, searchFlags, facets))) + _cachingEntitySearchService.search(List.of(entity), input, postFilters, sortCriterion, queryFrom, querySize, searchFlags, facets))) .stream() .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java index 56f6fed3ad9d2..f698e28c0be6d 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/client/CachingEntitySearchService.java @@ -53,7 +53,7 @@ public class CachingEntitySearchService { * @return a {@link SearchResult} containing the requested batch of search results */ public SearchResult search( - @Nonnull String entityName, + @Nonnull List entityNames, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @@ -61,7 +61,7 @@ public SearchResult search( int size, @Nullable SearchFlags flags, @Nullable List facets) { - return getCachedSearchResults(entityName, query, filters, sortCriterion, from, size, flags, facets); + return getCachedSearchResults(entityNames, query, filters, sortCriterion, from, size, flags, facets); } /** @@ -141,7 +141,7 @@ public ScrollResult scroll( * This lets us have batches that return a variable number of results (we have no idea which batch the "from" "size" page corresponds to) */ public SearchResult getCachedSearchResults( - @Nonnull String entityName, + @Nonnull List entityNames, @Nonnull String query, @Nullable Filter filters, @Nullable SortCriterion sortCriterion, @@ -152,9 +152,9 @@ public SearchResult getCachedSearchResults( return new CacheableSearcher<>( cacheManager.getCache(ENTITY_SEARCH_SERVICE_SEARCH_CACHE_NAME), batchSize, - querySize -> getRawSearchResults(entityName, query, filters, sortCriterion, querySize.getFrom(), + querySize -> getRawSearchResults(entityNames, query, filters, sortCriterion, querySize.getFrom(), querySize.getSize(), flags, facets), - querySize -> Sextet.with(entityName, query, filters != null ? toJsonString(filters) : null, + querySize -> Sextet.with(entityNames, query, filters != null ? toJsonString(filters) : null, sortCriterion != null ? toJsonString(sortCriterion) : null, facets, querySize), flags, enableCache).getSearchResults(from, size); } @@ -272,7 +272,7 @@ public ScrollResult getCachedScrollResults( * Executes the expensive search query using the {@link EntitySearchService} */ private SearchResult getRawSearchResults( - final String entityName, + final List entityNames, final String input, final Filter filters, final SortCriterion sortCriterion, @@ -280,7 +280,7 @@ private SearchResult getRawSearchResults( final int count, @Nullable final SearchFlags searchFlags, @Nullable final List facets) { - return entitySearchService.search(entityName, input, filters, sortCriterion, start, count, searchFlags, facets); + return entitySearchService.search(entityNames, input, filters, sortCriterion, start, count, searchFlags, facets); } /** diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java index 0680e6d74edf3..ce7b44c715d6b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java @@ -106,18 +106,18 @@ public void appendRunId(@Nonnull String entityName, @Nonnull Urn urn, @Nullable @Nonnull @Override - public SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + public SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags) { - return search(entityName, input, postFilters, sortCriterion, from, size, searchFlags, null); + return search(entityNames, input, postFilters, sortCriterion, from, size, searchFlags, null); } @Nonnull - public SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + public SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags, @Nullable List facets) { log.debug(String.format( "Searching FullText Search documents entityName: %s, input: %s, postFilters: %s, sortCriterion: %s, from: %s, size: %s", - entityName, input, postFilters, sortCriterion, from, size)); - return esSearchDAO.search(entityName, input, postFilters, sortCriterion, from, size, searchFlags, facets); + entityNames, input, postFilters, sortCriterion, from, size)); + return esSearchDAO.search(entityNames, input, postFilters, sortCriterion, from, size, searchFlags, facets); } @Nonnull diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java index 51e08763cd7c8..57e8967c83985 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java @@ -79,7 +79,7 @@ public long docCount(@Nonnull String entityName) { @Nonnull @WithSpan - private SearchResult executeAndExtract(@Nonnull EntitySpec entitySpec, @Nonnull SearchRequest searchRequest, + private SearchResult executeAndExtract(@Nonnull List entitySpec, @Nonnull SearchRequest searchRequest, @Nullable Filter filter, int from, int size) { long id = System.currentTimeMillis(); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "executeAndExtract_search").time()) { @@ -181,20 +181,22 @@ private ScrollResult executeAndExtract(@Nonnull List entitySpecs, @N * @return a {@link SearchResult} that contains a list of matched documents and related search result metadata */ @Nonnull - public SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + public SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags, @Nullable List facets) { final String finalInput = input.isEmpty() ? "*" : input; Timer.Context searchRequestTimer = MetricUtils.timer(this.getClass(), "searchRequest").time(); - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + List entitySpecs = entityNames.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList()); Filter transformedFilters = transformFilterForEntities(postFilters, indexConvention); // Step 1: construct the query final SearchRequest searchRequest = SearchRequestHandler - .getBuilder(entitySpec, searchConfiguration, customSearchConfiguration) + .getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration) .getSearchRequest(finalInput, transformedFilters, sortCriterion, from, size, searchFlags, facets); - searchRequest.indices(indexConvention.getIndexName(entitySpec)); + searchRequest.indices(entityNames.stream() + .map(indexConvention::getEntityIndexName) + .toArray(String[]::new)); searchRequestTimer.stop(); // Step 2: execute the query and extract results, validated against document model as well - return executeAndExtract(entitySpec, searchRequest, transformedFilters, from, size); + return executeAndExtract(entitySpecs, searchRequest, transformedFilters, from, size); } /** @@ -217,7 +219,7 @@ public SearchResult filter(@Nonnull String entityName, @Nullable Filter filters, .getFilterRequest(transformedFilters, sortCriterion, from, size); searchRequest.indices(indexConvention.getIndexName(entitySpec)); - return executeAndExtract(entitySpec, searchRequest, transformedFilters, from, size); + return executeAndExtract(List.of(entitySpec), searchRequest, transformedFilters, from, size); } /** diff --git a/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java index c1297866edcc4..79496888650e1 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java @@ -68,15 +68,24 @@ private ESTestUtils() { .collect(Collectors.toList()); } - public static SearchResult search(SearchService searchService, String query) { - return search(searchService, query, null); + public static SearchResult searchAcrossEntities(SearchService searchService, String query) { + return searchAcrossEntities(searchService, query, null); } - public static SearchResult search(SearchService searchService, String query, @Nullable List facets) { + public static SearchResult searchAcrossEntities(SearchService searchService, String query, @Nullable List facets) { return searchService.searchAcrossEntities(SEARCHABLE_ENTITIES, query, null, null, 0, 100, new SearchFlags().setFulltext(true).setSkipCache(true), facets); } + public static SearchResult search(SearchService searchService, String query) { + return search(searchService, SEARCHABLE_ENTITIES, query); + } + + public static SearchResult search(SearchService searchService, List entities, String query) { + return searchService.search(entities, query, null, null, 0, 100, + new SearchFlags().setFulltext(true).setSkipCache(true)); + } + public static ScrollResult scroll(SearchService searchService, String query, int batchSize, @Nullable String scrollId) { return searchService.scrollAcrossEntities(SEARCHABLE_ENTITIES, query, null, null, scrollId, "3m", batchSize, new SearchFlags().setFulltext(true).setSkipCache(true)); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java index 0e8d881b70791..9a6d2dc6fc1fa 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java @@ -23,6 +23,7 @@ import com.linkedin.metadata.search.elasticsearch.update.ESWriteDAO; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import java.util.List; import org.elasticsearch.client.RestHighLevelClient; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.annotation.Import; @@ -93,7 +94,7 @@ private ElasticSearchService buildService() { @Test public void testElasticSearchServiceStructuredQuery() throws Exception { - SearchResult searchResult = _elasticSearchService.search(ENTITY_NAME, "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); + SearchResult searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); BrowseResult browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); @@ -110,10 +111,10 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - searchResult = _elasticSearchService.search(ENTITY_NAME, "foreignKey:Node", null, null, 0, 10, new SearchFlags().setFulltext(false)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "foreignKey:Node", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); @@ -135,7 +136,7 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); @@ -152,7 +153,7 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); @@ -162,7 +163,7 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { @Test public void testElasticSearchServiceFulltext() throws Exception { - SearchResult searchResult = _elasticSearchService.search(ENTITY_NAME, "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + SearchResult searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); Urn urn = new TestEntityUrn("test", "urn1", "VALUE_1"); @@ -175,7 +176,7 @@ public void testElasticSearchServiceFulltext() throws Exception { _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); @@ -192,7 +193,7 @@ public void testElasticSearchServiceFulltext() throws Exception { _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); @@ -203,7 +204,7 @@ public void testElasticSearchServiceFulltext() throws Exception { _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(_bulkProcessor); - searchResult = _elasticSearchService.search(ENTITY_NAME, "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); + searchResult = _elasticSearchService.search(List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 0); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java index 817cf5aa5b37b..f4a8400fb005c 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java @@ -54,6 +54,7 @@ import java.util.stream.IntStream; import java.util.stream.Stream; +import static com.linkedin.metadata.Constants.*; import static com.linkedin.metadata.ESTestUtils.*; import static com.linkedin.metadata.search.elasticsearch.query.request.SearchQueryBuilder.STRUCTURED_QUERY_PREFIX; import static com.linkedin.metadata.utils.SearchUtil.*; @@ -189,10 +190,10 @@ public void testDatasetHasTags() throws IOException { @Test public void testFixtureInitialization() { assertNotNull(searchService); - SearchResult noResult = search(searchService, "no results"); + SearchResult noResult = searchAcrossEntities(searchService, "no results"); assertEquals(0, noResult.getEntities().size()); - final SearchResult result = search(searchService, "test"); + final SearchResult result = searchAcrossEntities(searchService, "test"); Map expectedTypes = Map.of( "dataset", 13, @@ -238,7 +239,7 @@ public void testDataPlatform() { .build(); expected.forEach((key, value) -> { - SearchResult result = search(searchService, key); + SearchResult result = searchAcrossEntities(searchService, key); assertEquals(result.getEntities().size(), value.intValue(), String.format("Unexpected data platform `%s` hits.", key)); // max is 100 without pagination }); @@ -254,14 +255,14 @@ public void testUrn() { "urn:li:mlFeature:(test_feature_table_all_feature_dtypes,test_BOOL_LIST_feature)", "urn:li:mlModel:(urn:li:dataPlatform:science,scienceModel,PROD)" ).forEach(query -> - assertTrue(search(searchService, query).getEntities().size() >= 1, + assertTrue(searchAcrossEntities(searchService, query).getEntities().size() >= 1, String.format("Unexpected >1 urn result for `%s`", query)) ); } @Test public void testExactTable() { - SearchResult results = search(searchService, "stg_customers"); + SearchResult results = searchAcrossEntities(searchService, "stg_customers"); assertEquals(results.getEntities().size(), 1, "Unexpected single urn result for `stg_customers`"); assertEquals(results.getEntities().get(0).getEntity().toString(), "urn:li:dataset:(urn:li:dataPlatform:dbt,cypress_project.jaffle_shop.stg_customers,PROD)"); @@ -278,7 +279,7 @@ public void testStemming() { testSets.forEach(testSet -> { Integer expectedResults = null; for (String testQuery : testSet) { - SearchResult results = search(searchService, testQuery); + SearchResult results = searchAcrossEntities(searchService, testQuery); assertTrue(results.hasEntities() && !results.getEntities().isEmpty(), String.format("Expected search results for `%s`", testQuery)); @@ -296,7 +297,7 @@ public void testStemmingOverride() throws IOException { Set testSet = Set.of("customer", "customers"); Set results = testSet.stream() - .map(test -> search(searchService, test)) + .map(test -> searchAcrossEntities(searchService, test)) .collect(Collectors.toSet()); results.forEach(r -> assertTrue(r.hasEntities() && !r.getEntities().isEmpty(), "Expected search results")); @@ -349,7 +350,7 @@ public void testDelimitedSynonym() throws IOException { "customer acquisition cost" ); List resultCounts = testSet.stream().map(q -> { - SearchResult result = search(searchService, q); + SearchResult result = searchAcrossEntities(searchService, q); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), "Expected search results for: " + q); return result.getEntities().size(); @@ -382,7 +383,7 @@ public void testUrnSynonym() throws IOException { "big query" ); List results = testSet.stream().map(query -> { - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), "Expected search results for: " + query); return result; }).collect(Collectors.toList()); @@ -621,7 +622,7 @@ public void testSmokeTestQueries() { ); Map results = expectedFulltextMinimums.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> search(searchService, entry.getKey()))); + .collect(Collectors.toMap(Map.Entry::getKey, entry -> searchAcrossEntities(searchService, entry.getKey()))); results.forEach((key, value) -> { Integer actualCount = value.getEntities().size(); @@ -691,7 +692,7 @@ public void testUnderscore() throws IOException { @Test public void testFacets() { Set expectedFacets = Set.of("entity", "typeNames", "platform", "origin", "tags"); - SearchResult testResult = search(searchService, "cypress"); + SearchResult testResult = searchAcrossEntities(searchService, "cypress"); expectedFacets.forEach(facet -> { assertTrue(testResult.getMetadata().getAggregations().stream().anyMatch(agg -> agg.getName().equals(facet)), String.format("Failed to find facet `%s` in %s", facet, @@ -703,7 +704,7 @@ public void testFacets() { @Test public void testNestedAggregation() { Set expectedFacets = Set.of("platform"); - SearchResult testResult = search(searchService, "cypress", List.copyOf(expectedFacets)); + SearchResult testResult = searchAcrossEntities(searchService, "cypress", List.copyOf(expectedFacets)); assertEquals(testResult.getMetadata().getAggregations().size(), 1); expectedFacets.forEach(facet -> { assertTrue(testResult.getMetadata().getAggregations().stream().anyMatch(agg -> agg.getName().equals(facet)), @@ -713,7 +714,7 @@ public void testNestedAggregation() { }); expectedFacets = Set.of("platform", "typeNames", "_entityType", "entity"); - SearchResult testResult2 = search(searchService, "cypress", List.copyOf(expectedFacets)); + SearchResult testResult2 = searchAcrossEntities(searchService, "cypress", List.copyOf(expectedFacets)); assertEquals(testResult2.getMetadata().getAggregations().size(), 4); expectedFacets.forEach(facet -> { assertTrue(testResult2.getMetadata().getAggregations().stream().anyMatch(agg -> agg.getName().equals(facet)), @@ -723,11 +724,11 @@ public void testNestedAggregation() { }); String singleNestedFacet = String.format("_entityType%sowners", AGGREGATION_SEPARATOR_CHAR); expectedFacets = Set.of(singleNestedFacet); - SearchResult testResultSingleNested = search(searchService, "cypress", List.copyOf(expectedFacets)); + SearchResult testResultSingleNested = searchAcrossEntities(searchService, "cypress", List.copyOf(expectedFacets)); assertEquals(testResultSingleNested.getMetadata().getAggregations().size(), 1); expectedFacets = Set.of("platform", singleNestedFacet, "typeNames", "origin"); - SearchResult testResultNested = search(searchService, "cypress", List.copyOf(expectedFacets)); + SearchResult testResultNested = searchAcrossEntities(searchService, "cypress", List.copyOf(expectedFacets)); assertEquals(testResultNested.getMetadata().getAggregations().size(), 4); expectedFacets.forEach(facet -> { assertTrue(testResultNested.getMetadata().getAggregations().stream().anyMatch(agg -> agg.getName().equals(facet)), @@ -811,6 +812,19 @@ public void testScrollAcrossEntities() throws IOException { assertEquals(totalResults, 8); } + @Test + public void testSearchAcrossMultipleEntities() { + String query = "logging_events"; + SearchResult result = search(searchService, query); + assertEquals((int) result.getNumEntities(), 8); + result = search(searchService, List.of(DATASET_ENTITY_NAME, DATA_JOB_ENTITY_NAME), query); + assertEquals((int) result.getNumEntities(), 8); + result = search(searchService, List.of(DATASET_ENTITY_NAME), query); + assertEquals((int) result.getNumEntities(), 4); + result = search(searchService, List.of(DATA_JOB_ENTITY_NAME), query); + assertEquals((int) result.getNumEntities(), 4); + } + @Test public void testQuotedAnalyzer() throws IOException { AnalyzeRequest request = AnalyzeRequest.withIndexAnalyzer( @@ -875,7 +889,7 @@ public void testFragmentUrns() { ); testSet.forEach(query -> { - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected partial urn search results", query)); @@ -941,7 +955,7 @@ public void testPlatformTest() { @Test public void testStructQueryFieldMatch() { String query = STRUCTURED_QUERY_PREFIX + "name: customers"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -954,7 +968,7 @@ public void testStructQueryFieldMatch() { @Test public void testStructQueryFieldPrefixMatch() { String query = STRUCTURED_QUERY_PREFIX + "name: customers*"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -967,7 +981,7 @@ public void testStructQueryFieldPrefixMatch() { @Test public void testStructQueryCustomPropertiesKeyPrefix() { String query = STRUCTURED_QUERY_PREFIX + "customProperties: node_type=*"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -980,7 +994,7 @@ public void testStructQueryCustomPropertiesKeyPrefix() { @Test public void testStructQueryCustomPropertiesMatch() { String query = STRUCTURED_QUERY_PREFIX + "customProperties: node_type=model"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -998,7 +1012,7 @@ public void testCustomPropertiesQuoted() { ); Map results = expectedResults.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> search(searchService, entry.getKey()))); + .collect(Collectors.toMap(Map.Entry::getKey, entry -> searchAcrossEntities(searchService, entry.getKey()))); results.forEach((key, value) -> { Integer actualCount = value.getEntities().size(); @@ -1012,7 +1026,7 @@ public void testCustomPropertiesQuoted() { @Test public void testStructQueryFieldPaths() { String query = STRUCTURED_QUERY_PREFIX + "fieldPaths: customer_id"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1025,7 +1039,7 @@ public void testStructQueryFieldPaths() { @Test public void testStructQueryBoolean() { String query = STRUCTURED_QUERY_PREFIX + "editedFieldTags:urn\\:li\\:tag\\:Legacy OR tags:urn\\:li\\:tag\\:testTag"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1035,7 +1049,7 @@ public void testStructQueryBoolean() { assertEquals(result.getEntities().size(), 2); query = STRUCTURED_QUERY_PREFIX + "editedFieldTags:urn\\:li\\:tag\\:Legacy"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1045,7 +1059,7 @@ public void testStructQueryBoolean() { assertEquals(result.getEntities().size(), 1); query = STRUCTURED_QUERY_PREFIX + "tags:urn\\:li\\:tag\\:testTag"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1058,7 +1072,7 @@ public void testStructQueryBoolean() { @Test public void testStructQueryBrowsePaths() { String query = STRUCTURED_QUERY_PREFIX + "browsePaths:*/dbt/*"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1071,7 +1085,7 @@ public void testStructQueryBrowsePaths() { @Test public void testOr() { String query = "stg_customers | logging_events"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1079,7 +1093,7 @@ public void testOr() { assertEquals(result.getEntities().size(), 9); query = "stg_customers"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1087,7 +1101,7 @@ public void testOr() { assertEquals(result.getEntities().size(), 1); query = "logging_events"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1098,7 +1112,7 @@ public void testOr() { @Test public void testNegate() { String query = "logging_events -bckp"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1106,7 +1120,7 @@ public void testNegate() { assertEquals(result.getEntities().size(), 7); query = "logging_events"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1117,7 +1131,7 @@ public void testNegate() { @Test public void testPrefix() { String query = "bigquery"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1125,7 +1139,7 @@ public void testPrefix() { assertEquals(result.getEntities().size(), 8); query = "big*"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1136,7 +1150,7 @@ public void testPrefix() { @Test public void testParens() { String query = "dbt | (bigquery + covid19)"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1144,7 +1158,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 11); query = "dbt"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1152,7 +1166,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 9); query = "bigquery + covid19"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1160,7 +1174,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 2); query = "bigquery"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1168,7 +1182,7 @@ public void testParens() { assertEquals(result.getEntities().size(), 8); query = "covid19"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1179,7 +1193,7 @@ public void testParens() { @Test public void testPrefixVsExact() { String query = "\"customers\""; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1197,7 +1211,7 @@ public void testPrefixVsExact() { public void testPrefixVsExactCaseSensitivity() { List insensitiveExactMatches = List.of("testExactMatchCase", "testexactmatchcase", "TESTEXACTMATCHCASE"); for (String query : insensitiveExactMatches) { - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); @@ -1214,7 +1228,7 @@ public void testPrefixVsExactCaseSensitivity() { @Test public void testColumnExactMatch() { String query = "unit_data"; - SearchResult result = search(searchService, query); + SearchResult result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), @@ -1227,7 +1241,7 @@ public void testColumnExactMatch() { "Expected table name exact match first"); query = "special_column_only_present_here_info"; - result = search(searchService, query); + result = searchAcrossEntities(searchService, query); assertTrue(result.hasEntities() && !result.getEntities().isEmpty(), String.format("%s - Expected search results", query)); assertTrue(result.getEntities().stream().noneMatch(e -> e.getMatchedFields().isEmpty()), diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SearchLineageDataFixtureTests.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SearchLineageDataFixtureTests.java index f726eba547b99..55f7d4618f479 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SearchLineageDataFixtureTests.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SearchLineageDataFixtureTests.java @@ -2,6 +2,7 @@ import com.linkedin.common.urn.Urn; import com.linkedin.metadata.ESSearchLineageFixture; +import com.linkedin.metadata.ESTestUtils; import com.linkedin.metadata.search.LineageSearchResult; import com.linkedin.metadata.search.LineageSearchService; import com.linkedin.metadata.search.SearchResult; @@ -14,7 +15,6 @@ import java.net.URISyntaxException; -import static com.linkedin.metadata.ESTestUtils.search; import static com.linkedin.metadata.ESTestUtils.lineage; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; @@ -35,10 +35,10 @@ public class SearchLineageDataFixtureTests extends AbstractTestNGSpringContextTe @Test public void testFixtureInitialization() { assertNotNull(searchService); - SearchResult noResult = search(searchService, "no results"); + SearchResult noResult = ESTestUtils.searchAcrossEntities(searchService, "no results"); assertEquals(noResult.getEntities().size(), 0); - SearchResult result = search(searchService, "e3859789eed1cef55288b44f016ee08290d9fd08973e565c112d8"); + SearchResult result = ESTestUtils.searchAcrossEntities(searchService, "e3859789eed1cef55288b44f016ee08290d9fd08973e565c112d8"); assertEquals(result.getEntities().size(), 1); assertEquals(result.getEntities().get(0).getEntity().toString(), diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java index 780927d21417f..2be719ed263ea 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java @@ -141,7 +141,7 @@ public void invoke(@Nonnull MetadataChangeLog event) { private void handleEntityKeyEvent(DatasetUrn datasetUrn) { Filter entitiesWithYouAsSiblingFilter = createFilterForEntitiesWithYouAsSibling(datasetUrn); final SearchResult searchResult = _searchService.search( - "dataset", + List.of(DATASET_ENTITY_NAME), "*", entitiesWithYouAsSiblingFilter, null, diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java index 3eb21125bbaf2..5fb2cfaaef2d1 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java @@ -255,7 +255,7 @@ public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { Mockito.when( _mockSearchService.search( - anyString(), anyString(), any(), any(), anyInt(), anyInt(), eq(new SearchFlags().setFulltext(false) + any(), anyString(), any(), any(), anyInt(), anyInt(), eq(new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)) )).thenReturn(returnSearchResult); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java index 486961c2c1f07..097dcfdfdf52e 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java @@ -74,7 +74,7 @@ public ExecutionMode getExecutionMode() { private int getAndRestoreTermAspectIndices(int start, AuditStamp auditStamp, AspectSpec termAspectSpec) throws Exception { SearchResult termsResult = - _entitySearchService.search(Constants.GLOSSARY_TERM_ENTITY_NAME, "", null, + _entitySearchService.search(List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), "", null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); List termUrns = termsResult.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList()); @@ -116,7 +116,7 @@ null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) } private int getAndRestoreNodeAspectIndices(int start, AuditStamp auditStamp, AspectSpec nodeAspectSpec) throws Exception { - SearchResult nodesResult = _entitySearchService.search(Constants.GLOSSARY_NODE_ENTITY_NAME, "", + SearchResult nodesResult = _entitySearchService.search(List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), "", null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); List nodeUrns = nodesResult.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList()); diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java index 88c63110ee63e..d56fbed07f890 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java @@ -21,6 +21,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.mxe.MetadataChangeProposal; +import java.util.List; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -40,7 +41,7 @@ private void mockGetTermInfo(Urn glossaryTermUrn, EntitySearchService mockSearch termInfoAspects.put(Constants.GLOSSARY_TERM_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(new GlossaryTermInfo().setName("test").data()))); Map termInfoResponses = new HashMap<>(); termInfoResponses.put(glossaryTermUrn, new EntityResponse().setUrn(glossaryTermUrn).setAspects(new EnvelopedAspectMap(termInfoAspects))); - Mockito.when(mockSearchService.search(Constants.GLOSSARY_TERM_ENTITY_NAME, "", null, null, 0, 1000, + Mockito.when(mockSearchService.search(List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), "", null, null, 0, 1000, new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true))) .thenReturn(new SearchResult().setNumEntities(1).setEntities(new SearchEntityArray(ImmutableList.of(new SearchEntity().setEntity(glossaryTermUrn))))); Mockito.when(mockService.getEntitiesV2( @@ -55,7 +56,7 @@ private void mockGetNodeInfo(Urn glossaryNodeUrn, EntitySearchService mockSearch nodeInfoAspects.put(Constants.GLOSSARY_NODE_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(new GlossaryNodeInfo().setName("test").data()))); Map nodeInfoResponses = new HashMap<>(); nodeInfoResponses.put(glossaryNodeUrn, new EntityResponse().setUrn(glossaryNodeUrn).setAspects(new EnvelopedAspectMap(nodeInfoAspects))); - Mockito.when(mockSearchService.search(Constants.GLOSSARY_NODE_ENTITY_NAME, "", null, null, 0, 1000, + Mockito.when(mockSearchService.search(List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), "", null, null, 0, 1000, new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true))) .thenReturn(new SearchResult().setNumEntities(1).setEntities(new SearchEntityArray(ImmutableList.of(new SearchEntity().setEntity(glossaryNodeUrn))))); Mockito.when(mockService.getEntitiesV2( @@ -221,10 +222,10 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.GLOSSARY_TERM_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.GLOSSARY_NODE_ENTITY_NAME); - Mockito.verify(mockSearchService, Mockito.times(0)).search(Constants.GLOSSARY_TERM_ENTITY_NAME, + Mockito.verify(mockSearchService, Mockito.times(0)).search(List.of(Constants.GLOSSARY_TERM_ENTITY_NAME), "", null, null, 0, 1000, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); - Mockito.verify(mockSearchService, Mockito.times(0)).search(Constants.GLOSSARY_NODE_ENTITY_NAME, + Mockito.verify(mockSearchService, Mockito.times(0)).search(List.of(Constants.GLOSSARY_NODE_ENTITY_NAME), "", null, null, 0, 1000, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java index b8fd785eaad0f..f6dedfb9a07c6 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java @@ -332,7 +332,7 @@ public Task search(@ActionParam(PARAM_ENTITY) @Nonnull String enti () -> { final SearchResult result; // This API is not used by the frontend for search bars so we default to structured - result = _entitySearchService.search(entityName, input, filter, sortCriterion, start, count, searchFlags); + result = _entitySearchService.search(List.of(entityName), input, filter, sortCriterion, start, count, searchFlags); return validateSearchResult(result, _entityService); }, MetricRegistry.name(this.getClass(), "search")); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java index de5bdb62f201b..a46b58aabfb0b 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java @@ -64,7 +64,7 @@ public interface EntitySearchService { * Safe for non-structured, user input, queries with an attempt to provide some advanced features * Impl * - * @param entityName name of the entity + * @param entityNames names of the entities * @param input the search input text * @param postFilters the request map with fields and values as filters to be applied to search hits * @param sortCriterion {@link SortCriterion} to be applied to search results @@ -74,7 +74,7 @@ public interface EntitySearchService { * @return a {@link SearchResult} that contains a list of matched documents and related search result metadata */ @Nonnull - SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags); /** @@ -84,7 +84,7 @@ SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable * Safe for non-structured, user input, queries with an attempt to provide some advanced features * Impl * - * @param entityName name of the entity + * @param entityNames names of the entities * @param input the search input text * @param postFilters the request map with fields and values as filters to be applied to search hits * @param sortCriterion {@link SortCriterion} to be applied to search results @@ -95,7 +95,7 @@ SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable * @return a {@link SearchResult} that contains a list of matched documents and related search result metadata */ @Nonnull - SearchResult search(@Nonnull String entityName, @Nonnull String input, @Nullable Filter postFilters, + SearchResult search(@Nonnull List entityNames, @Nonnull String input, @Nullable Filter postFilters, @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags, @Nullable List facets); /** diff --git a/mock-entity-registry/src/main/java/mock/MockEntitySpec.java b/mock-entity-registry/src/main/java/mock/MockEntitySpec.java index 097e0845504a1..f43c1f7fd6613 100644 --- a/mock-entity-registry/src/main/java/mock/MockEntitySpec.java +++ b/mock-entity-registry/src/main/java/mock/MockEntitySpec.java @@ -25,6 +25,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static com.linkedin.metadata.Constants.*; @@ -88,7 +89,8 @@ public AspectSpec createAspectSpec(T type, String nam @Override public List getAspectSpecs() { - return Collections.emptyList(); + return ASPECT_TYPE_MAP.keySet().stream().map(name -> createAspectSpec(ASPECT_TYPE_MAP.get(name), name)).collect( + Collectors.toList()); } @Override