diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java index 3ba268d77d20a..a68719f5cebe3 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.search.elasticsearch.query.request; -import static com.linkedin.metadata.search.utils.ESUtils.*; import static com.linkedin.metadata.search.utils.ESUtils.NAME_SUGGESTION; import static com.linkedin.metadata.search.utils.ESUtils.applyDefaultSearchFilters; @@ -57,6 +56,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.Getter; +import lombok.Value; import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.CollectionUtils; import org.opensearch.action.search.SearchRequest; @@ -76,7 +76,7 @@ @Slf4j public class SearchRequestHandler extends BaseRequestHandler { - private static final Map, SearchRequestHandler> REQUEST_HANDLER_BY_ENTITY_NAME = + private static final Map REQUEST_HANDLER_BY_ENTITY_NAME = new ConcurrentHashMap<>(); private final List entitySpecs; private final List entityNames; @@ -144,16 +144,13 @@ public static SearchRequestHandler getBuilder( @Nullable CustomSearchConfiguration customSearchConfiguration, @Nonnull QueryFilterRewriteChain queryFilterRewriteChain, @Nonnull SearchServiceConfiguration searchServiceConfiguration) { - return REQUEST_HANDLER_BY_ENTITY_NAME.computeIfAbsent( + return getBuilder( + systemOperationContext, ImmutableList.of(entitySpec), - k -> - new SearchRequestHandler( - systemOperationContext, - entitySpec, - configs, - customSearchConfiguration, - queryFilterRewriteChain, - searchServiceConfiguration)); + configs, + customSearchConfiguration, + queryFilterRewriteChain, + searchServiceConfiguration); } public static SearchRequestHandler getBuilder( @@ -164,7 +161,12 @@ public static SearchRequestHandler getBuilder( @Nonnull QueryFilterRewriteChain queryFilterRewriteChain, @Nonnull SearchServiceConfiguration searchServiceConfiguration) { return REQUEST_HANDLER_BY_ENTITY_NAME.computeIfAbsent( - ImmutableList.copyOf(entitySpecs), + new SearchHandlerKey( + ImmutableList.copyOf(entitySpecs), + configs, + customSearchConfiguration, + queryFilterRewriteChain, + searchServiceConfiguration), k -> new SearchRequestHandler( systemOperationContext, @@ -676,4 +678,22 @@ private List extractSearchSuggestions(@Nonnull SearchResponse } return searchSuggestions; } + + /** + * Enhanced cache key implementation to prevent handler cross-contamination in tests. + * + *

Background: Flaky tests occurred because the cache key (previously just entitySpecs) didn't + * account for all configuration variants. Identical entitySpecs with different search + * configurations would incorrectly share handlers, leading to test instability. + * + *

This key ensures each unique configuration combination gets its own handler instance. + */ + @Value + private static class SearchHandlerKey { + @Nonnull private final List entitySpecs; + @Nonnull private final ElasticSearchConfiguration configs; + @Nullable private final CustomSearchConfiguration customSearchConfiguration; + @Nonnull private final QueryFilterRewriteChain queryFilterRewriteChain; + @Nonnull private final SearchServiceConfiguration searchServiceConfiguration; + } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java index 8de8e2ac1e46f..1819909436cca 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java @@ -3,9 +3,9 @@ import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion; import static io.datahubproject.test.search.SearchTestUtils.TEST_ES_SEARCH_CONFIG; import static io.datahubproject.test.search.SearchTestUtils.TEST_OS_SEARCH_CONFIG; -import static io.datahubproject.test.search.SearchTestUtils.TEST_OS_SEARCH_CONFIG_WITH_PIT; import static io.datahubproject.test.search.SearchTestUtils.TEST_SEARCH_SERVICE_CONFIG; import static io.datahubproject.test.search.SearchTestUtils.createDelegatingMappingsBuilder; +import static io.datahubproject.test.search.SearchTestUtils.getTestOsSearchConfigWithPit; import static io.datahubproject.test.search.SearchTestUtils.syncAfterWrite; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; @@ -128,9 +128,9 @@ public void setup() throws RemoteInvocationException, URISyntaxException { IndexConvention mockIndexConvention = mock(IndexConvention.class); when(mockIndexConvention.isV2EntityIndex(anyString())).thenReturn(true); settingsBuilder = new V2LegacySettingsBuilder(indexConfiguration, mockIndexConvention); - elasticSearchService = buildEntitySearchService(); + elasticSearchService = buildEntitySearchService(getSearchConfiguration()); elasticSearchService.reindexAll(operationContext, Collections.emptySet()); - pitElasticSearchService = buildPITEntitySearchService(); + pitElasticSearchService = buildPITEntitySearchService(getSearchConfiguration()); pitElasticSearchService.reindexAll(operationContext, Collections.emptySet()); cacheManager = new ConcurrentMapCacheManager(); resetSearchService(); @@ -174,8 +174,9 @@ public void wipe() throws Exception { } @Nonnull - private ElasticSearchService buildEntitySearchService() { - ElasticSearchConfiguration esConfig = TEST_OS_SEARCH_CONFIG.toBuilder().build(); + private ElasticSearchService buildEntitySearchService(SearchConfiguration searchConfiguration) { + ElasticSearchConfiguration esConfig = + TEST_OS_SEARCH_CONFIG.toBuilder().search(searchConfiguration).build(); ESSearchDAO searchDAO = new ESSearchDAO( getSearchClient(), @@ -206,8 +207,11 @@ private ElasticSearchService buildEntitySearchService() { } @Nonnull - private ElasticSearchService buildPITEntitySearchService() { - ElasticSearchConfiguration esConfig = TEST_OS_SEARCH_CONFIG_WITH_PIT.toBuilder().build(); + private ElasticSearchService buildPITEntitySearchService( + SearchConfiguration searchConfiguration) { + ElasticSearchConfiguration esConfig = + getTestOsSearchConfigWithPit(searchConfiguration).toBuilder().build(); + ESSearchDAO searchDAO = new ESSearchDAO( getSearchClient(), diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java index f8d958ed8bbdb..412b7569725d8 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java @@ -458,7 +458,7 @@ public void testExplain() { explainResponse.getId(), "urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-public-data.covid19_geotab_mobility_impact.ca_border_wait_times,PROD)"); assertTrue(explainResponse.isExists()); - assertEquals(explainResponse.getExplanation().getValue(), 18.0f); + assertEquals(explainResponse.getExplanation().getValue(), 1.25f); } @Test diff --git a/metadata-io/src/testFixtures/java/io/datahubproject/test/search/SearchTestUtils.java b/metadata-io/src/testFixtures/java/io/datahubproject/test/search/SearchTestUtils.java index 71df205cd0186..5d083c154a2ae 100644 --- a/metadata-io/src/testFixtures/java/io/datahubproject/test/search/SearchTestUtils.java +++ b/metadata-io/src/testFixtures/java/io/datahubproject/test/search/SearchTestUtils.java @@ -155,17 +155,19 @@ private SearchTestUtils() {} .build(); // Configuration with PIT enabled for search entities (for tests that specifically need PIT) - public static ElasticSearchConfiguration TEST_OS_SEARCH_CONFIG_WITH_PIT = - BASE_TEST_CONFIG.toBuilder() - .search( - BASE_TEST_CONFIG.getSearch().toBuilder() - .pointInTimeCreationEnabled(true) // Enable PIT for search entities - .graph( - BASE_TEST_CONFIG.getSearch().getGraph().toBuilder() - .pointInTimeCreationEnabled(true) // Enable graph PIT - .build()) - .build()) - .build(); + public static ElasticSearchConfiguration getTestOsSearchConfigWithPit( + SearchConfiguration searchConfiguration) { + return BASE_TEST_CONFIG.toBuilder() + .search( + searchConfiguration.toBuilder() + .pointInTimeCreationEnabled(true) // Enable PIT for search entities + .graph( + searchConfiguration.getGraph().toBuilder() + .pointInTimeCreationEnabled(true) // Enable graph PIT + .build()) + .build()) + .build(); + } public static ElasticSearchConfiguration TEST_ES_SEARCH_CONFIG = TEST_OS_SEARCH_CONFIG.toBuilder().build();