[ML] Return all Datafeeds with GET Anomaly Detector (#84759) (#84770)

Fixes a bug where only the first 10 datafeeds were included Backport of #84759
elastic · Mar 8, 2022 · c59fabc · c59fabc
1 parent 79c1cfe
commit c59fabc
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 3 deletions.
diff --git a/docs/changelog/84759.yaml b/docs/changelog/84759.yaml
@@ -0,0 +1,5 @@
+pr: 84759
+summary: Return all Datafeeds with GET Anomaly Detector
+area: Machine Learning
+type: bug
+issues: []
diff --git a/...rnalClusterTest/java/org/elasticsearch/xpack/ml/integration/DatafeedConfigProviderIT.java b/...rnalClusterTest/java/org/elasticsearch/xpack/ml/integration/DatafeedConfigProviderIT.java
@@ -25,9 +25,11 @@
 import org.hamcrest.core.IsInstanceOf;
 import org.junit.Before;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -432,6 +434,31 @@ public void testFindDatafeedIdsForJobIds() throws Exception {
         assertThat(datafeedIdsHolder.get(), contains("bar-1", "foo-1"));
     }
 
+    public void testFindDatafeedIdsForJobIds_ManyJobs() throws Exception {
+        var jobIds = new ArrayList<String>();
+        var dfIds = new HashSet<String>();
+        for (int i = 0; i < 13; i++) {
+            String id = Integer.toString(i);
+            var dfId = "df-" + id;
+            var jobId = "j-" + id;
+            putDatafeedConfig(createDatafeedConfig(dfId, jobId), Collections.emptyMap());
+            dfIds.add(dfId);
+            jobIds.add(jobId);
+        }
+
+        client().admin().indices().prepareRefresh(MlConfigIndex.indexName()).get();
+
+        AtomicReference<Set<String>> datafeedIdsHolder = new AtomicReference<>();
+        AtomicReference<Exception> exceptionHolder = new AtomicReference<>();
+
+        blockingCall(
+            actionListener -> datafeedConfigProvider.findDatafeedIdsForJobIds(jobIds, actionListener),
+            datafeedIdsHolder,
+            exceptionHolder
+        );
+        assertEquals(dfIds, datafeedIdsHolder.get());
+    }
+
     public void testFindDatafeedsForJobIds() throws Exception {
         putDatafeedConfig(createDatafeedConfig("foo-1", "j1"), Collections.emptyMap());
         putDatafeedConfig(createDatafeedConfig("foo-2", "j2"), Collections.emptyMap());
@@ -467,6 +494,29 @@ public void testFindDatafeedsForJobIds() throws Exception {
         assertThat(datafeedMapHolder.get().get("j1").getId(), equalTo("foo-1"));
     }
 
+    public void testFindDatafeedsForJobIds_ManyJobs() throws Exception {
+        var jobIds = new ArrayList<String>();
+        for (int i = 0; i < 13; i++) {
+            String id = Integer.toString(i);
+            var dfId = "df-" + id;
+            var jobId = "j-" + id;
+            putDatafeedConfig(createDatafeedConfig(dfId, jobId), Collections.emptyMap());
+            jobIds.add(jobId);
+        }
+
+        client().admin().indices().prepareRefresh(MlConfigIndex.indexName()).get();
+
+        AtomicReference<Map<String, DatafeedConfig.Builder>> datafeedMapHolder = new AtomicReference<>();
+        AtomicReference<Exception> exceptionHolder = new AtomicReference<>();
+
+        blockingCall(
+            actionListener -> datafeedConfigProvider.findDatafeedsByJobIds(jobIds, actionListener),
+            datafeedMapHolder,
+            exceptionHolder
+        );
+        assertThat(datafeedMapHolder.get().entrySet(), hasSize(jobIds.size()));
+    }
+
     public void testHeadersAreOverwritten() throws Exception {
         String dfId = "df-with-headers";
         DatafeedConfig.Builder configWithUnrelatedHeaders = createDatafeedConfig(dfId, "j1");

diff --git a/...src/main/java/org/elasticsearch/xpack/ml/datafeed/persistence/DatafeedConfigProvider.java b/...src/main/java/org/elasticsearch/xpack/ml/datafeed/persistence/DatafeedConfigProvider.java
@@ -182,11 +182,11 @@ public void onFailure(Exception e) {
     public void findDatafeedIdsForJobIds(Collection<String> jobIds, ActionListener<Set<String>> listener) {
         SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(buildDatafeedJobIdsQuery(jobIds));
         sourceBuilder.fetchSource(false);
+        sourceBuilder.size(jobIds.size());
         sourceBuilder.docValueField(DatafeedConfig.ID.getPreferredName(), null);
 
         SearchRequest searchRequest = client.prepareSearch(MlConfigIndex.indexName())
             .setIndicesOptions(IndicesOptions.lenientExpandOpen())
-            .setSize(jobIds.size())
             .setSource(sourceBuilder)
             .request();
 
@@ -213,8 +213,7 @@ public void findDatafeedIdsForJobIds(Collection<String> jobIds, ActionListener<S
     public void findDatafeedsByJobIds(Collection<String> jobIds, ActionListener<Map<String, DatafeedConfig.Builder>> listener) {
         SearchRequest searchRequest = client.prepareSearch(MlConfigIndex.indexName())
             .setIndicesOptions(IndicesOptions.lenientExpandOpen())
-            .setSize(jobIds.size())
-            .setSource(new SearchSourceBuilder().query(buildDatafeedJobIdsQuery(jobIds)))
+            .setSource(new SearchSourceBuilder().query(buildDatafeedJobIdsQuery(jobIds)).size(jobIds.size()))
             .request();
 
         executeAsyncWithOrigin(