diff --git a/metadata-io/src/main/java/com/linkedin/metadata/dataHubUsage/DataHubUsageServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/dataHubUsage/DataHubUsageServiceImpl.java index 4cd3d754dc..9e6a360af9 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/dataHubUsage/DataHubUsageServiceImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/dataHubUsage/DataHubUsageServiceImpl.java @@ -112,7 +112,8 @@ public class DataHubUsageServiceImpl implements DataHubUsageService { response.count(searchHits.getHits().length); response.total((int) searchHits.getTotalHits().value); String nextScrollId = null; - if (searchHits.getHits().length == analyticsSearchRequest.getSize()) { + if (searchHits.getHits().length == analyticsSearchRequest.getSize() + && searchHits.getHits().length > 0) { Object[] sort = searchHits.getHits()[searchHits.getHits().length - 1].getSortValues(); nextScrollId = new SearchAfterWrapper(sort, null, 0L).toScrollId(); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java index 0ee7524414..a2ab660d93 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java @@ -364,7 +364,7 @@ public class ElasticSearchGraphService implements GraphService, ElasticSearchInd SearchHit[] searchHits = response.getHits().getHits(); // Only return next scroll ID if there are more results, indicated by full size results String nextScrollId = null; - if (searchHits.length == count) { + if (searchHits.length == count && searchHits.length > 0) { Object[] sort = searchHits[searchHits.length - 1].getSortValues(); nextScrollId = new SearchAfterWrapper(sort, null, 0L).toScrollId(); } 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 37de8dee88..8c3ec7e384 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 @@ -452,7 +452,7 @@ public class SearchRequestHandler extends BaseRequestHandler { SearchHit[] searchHits = searchResponse.getHits().getHits(); // Only return next scroll ID if there are more results, indicated by full size results String nextScrollId = null; - if (searchHits.length == size) { + if (searchHits.length == size && searchHits.length > 0) { Object[] sort = searchHits[searchHits.length - 1].getSortValues(); long expirationTimeMs = 0L; if (keepAlive != null && supportsPointInTime) { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java index 2c603d9dd4..0814882c82 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.graph.elastic; import static io.datahubproject.test.search.SearchTestUtils.TEST_GRAPH_SERVICE_CONFIG; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; @@ -254,6 +255,52 @@ public class ElasticSearchGraphServiceTest { assertTrue(result.getEntities().isEmpty()); } + @Test + public void testScrollRelatedEntitiesWithZeroCount() { + // Test the edge case where count=0 is passed, which should not cause + // ArrayIndexOutOfBoundsException when accessing searchHits[searchHits.length - 1] + OperationContext mockOpContext = TestOperationContexts.systemContextNoValidate(); + GraphFilters mockGraphFilters = GraphFilters.ALL; + List mockSortCriteria = Collections.emptyList(); + + String scrollId = "test-scroll-id"; + String keepAlive = "1m"; + int count = 0; // This is the edge case we're testing + + // Create mock search response with empty hits + SearchResponse mockResponse = mock(SearchResponse.class); + SearchHits mockHits = mock(SearchHits.class); + when(mockResponse.getHits()).thenReturn(mockHits); + + // Create empty search hits array (simulating no results) + SearchHit[] hits = new SearchHit[0]; + when(mockHits.getHits()).thenReturn(hits); + when(mockHits.getTotalHits()).thenReturn(new TotalHits(0L, TotalHits.Relation.EQUAL_TO)); + + when(mockReadDAO.getSearchResponse(any(), any(), any(), any(), any(), anyInt())) + .thenReturn(mockResponse); + when(mockReadDAO.getGraphServiceConfig()).thenReturn(TEST_GRAPH_SERVICE_CONFIG); + + // This should not throw ArrayIndexOutOfBoundsException + RelatedEntitiesScrollResult result = + test.scrollRelatedEntities( + mockOpContext, + mockGraphFilters, + mockSortCriteria, + scrollId, + keepAlive, + count, + null, + null); + + // Verify the result + assertNotNull(result); + assertEquals(result.getNumResults(), 0); + assertEquals(result.getPageSize(), 0); + assertTrue(result.getEntities().isEmpty()); + assertNull(result.getScrollId()); // No scroll ID since no results + } + @Test public void testRaw() { // Create test edge tuples diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java index 16a11ae51a..4f0916287e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java @@ -1383,6 +1383,43 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { assertFalse(result.hasScrollId()); } + @Test + public void testExtractScrollResultWithZeroCount() { + // Test the edge case where count=0 is passed, which should not cause + // ArrayIndexOutOfBoundsException + SearchRequestHandler handler = + SearchRequestHandler.getBuilder( + operationContext, + TestEntitySpecBuilder.getSpec(), + testQueryConfig, + null, + QueryFilterRewriteChain.EMPTY, + TEST_SEARCH_SERVICE_CONFIG); + + SearchResponse mockResponse = mock(SearchResponse.class); + SearchHits mockHits = mock(SearchHits.class); + + // Create empty search hits array (simulating no results) + SearchHit[] hits = new SearchHit[0]; + + when(mockResponse.getHits()).thenReturn(mockHits); + when(mockHits.getTotalHits()).thenReturn(new TotalHits(0L, TotalHits.Relation.EQUAL_TO)); + when(mockHits.getHits()).thenReturn(hits); + when(mockResponse.getAggregations()).thenReturn(null); + when(mockResponse.getSuggest()).thenReturn(null); + when(mockResponse.pointInTimeId()).thenReturn("test-pit-id"); + + // This should not throw ArrayIndexOutOfBoundsException + ScrollResult result = + handler.extractScrollResult(operationContext, mockResponse, null, "5m", 0, true); + + // Verify the result + assertNotNull(result); + assertEquals(result.getPageSize().intValue(), 0); + assertEquals(result.getNumEntities().intValue(), 0); + assertFalse(result.hasScrollId()); // No scroll ID since no results + } + // Helper method to create scroll results with specific sizes private ScrollResult verifyScrollResultSize( SearchRequestHandler handler,