Skip to content

Commit 4e3d565

Browse files
davidkyledavidkyle
committed
[ML] Job in Index: Stop and preview datafeed (elastic#34605)
1 parent f8614a1 commit 4e3d565

File tree

3 files changed

+84
-146
lines changed

3 files changed

+84
-146
lines changed

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPreviewDatafeedAction.java‎

Lines changed: 39 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -9,22 +9,20 @@
99
import org.elasticsearch.action.support.ActionFilters;
1010
import org.elasticsearch.action.support.HandledTransportAction;
1111
import org.elasticsearch.client.Client;
12-
import org.elasticsearch.cluster.service.ClusterService;
1312
import org.elasticsearch.common.bytes.BytesArray;
1413
import org.elasticsearch.common.inject.Inject;
1514
import org.elasticsearch.common.settings.Settings;
1615
import org.elasticsearch.tasks.Task;
1716
import org.elasticsearch.threadpool.ThreadPool;
1817
import org.elasticsearch.transport.TransportService;
1918
import org.elasticsearch.xpack.core.ClientHelper;
20-
import org.elasticsearch.xpack.core.ml.MlMetadata;
2119
import org.elasticsearch.xpack.core.ml.action.PreviewDatafeedAction;
2220
import org.elasticsearch.xpack.core.ml.datafeed.ChunkingConfig;
2321
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
2422
import org.elasticsearch.xpack.core.ml.datafeed.extractor.DataExtractor;
25-
import org.elasticsearch.xpack.core.ml.job.config.Job;
26-
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
2723
import org.elasticsearch.xpack.ml.datafeed.extractor.DataExtractorFactory;
24+
import org.elasticsearch.xpack.ml.datafeed.persistence.DatafeedConfigProvider;
25+
import org.elasticsearch.xpack.ml.job.persistence.JobConfigProvider;
2826

2927
import java.io.BufferedReader;
3028
import java.io.InputStream;
@@ -39,51 +37,56 @@ public class TransportPreviewDatafeedAction extends HandledTransportAction<Previ
3937

4038
private final ThreadPool threadPool;
4139
private final Client client;
42-
private final ClusterService clusterService;
40+
private final JobConfigProvider jobConfigProvider;
41+
private final DatafeedConfigProvider datafeedConfigProvider;
4342

4443
@Inject
4544
public TransportPreviewDatafeedAction(Settings settings, ThreadPool threadPool, TransportService transportService,
46-
ActionFilters actionFilters, Client client, ClusterService clusterService) {
45+
ActionFilters actionFilters, Client client, JobConfigProvider jobConfigProvider,
46+
DatafeedConfigProvider datafeedConfigProvider) {
4747
super(settings, PreviewDatafeedAction.NAME, transportService, actionFilters,
4848
(Supplier<PreviewDatafeedAction.Request>) PreviewDatafeedAction.Request::new);
4949
this.threadPool = threadPool;
5050
this.client = client;
51-
this.clusterService = clusterService;
51+
this.jobConfigProvider = jobConfigProvider;
52+
this.datafeedConfigProvider = datafeedConfigProvider;
5253
}
5354

5455
@Override
5556
protected void doExecute(Task task, PreviewDatafeedAction.Request request, ActionListener<PreviewDatafeedAction.Response> listener) {
56-
MlMetadata mlMetadata = MlMetadata.getMlMetadata(clusterService.state());
57-
DatafeedConfig datafeed = mlMetadata.getDatafeed(request.getDatafeedId());
58-
if (datafeed == null) {
59-
throw ExceptionsHelper.missingDatafeedException(request.getDatafeedId());
60-
}
61-
Job job = mlMetadata.getJobs().get(datafeed.getJobId());
62-
if (job == null) {
63-
throw ExceptionsHelper.missingJobException(datafeed.getJobId());
64-
}
6557

66-
DatafeedConfig.Builder previewDatafeed = buildPreviewDatafeed(datafeed);
67-
Map<String, String> headers = threadPool.getThreadContext().getHeaders().entrySet().stream()
68-
.filter(e -> ClientHelper.SECURITY_HEADER_FILTERS.contains(e.getKey()))
69-
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
70-
previewDatafeed.setHeaders(headers);
71-
// NB: this is using the client from the transport layer, NOT the internal client.
72-
// This is important because it means the datafeed search will fail if the user
73-
// requesting the preview doesn't have permission to search the relevant indices.
74-
DataExtractorFactory.create(client, previewDatafeed.build(), job, new ActionListener<DataExtractorFactory>() {
75-
@Override
76-
public void onResponse(DataExtractorFactory dataExtractorFactory) {
77-
DataExtractor dataExtractor = dataExtractorFactory.newExtractor(0, Long.MAX_VALUE);
78-
threadPool.generic().execute(() -> previewDatafeed(dataExtractor, listener));
79-
}
80-
81-
@Override
82-
public void onFailure(Exception e) {
83-
listener.onFailure(e);
84-
}
85-
});
58+
datafeedConfigProvider.getDatafeedConfig(request.getDatafeedId(), ActionListener.wrap(
59+
datafeedConfigBuilder -> {
60+
DatafeedConfig datafeedConfig = datafeedConfigBuilder.build();
61+
jobConfigProvider.getJob(datafeedConfig.getJobId(), ActionListener.wrap(
62+
jobBuilder -> {
63+
DatafeedConfig.Builder previewDatafeed = buildPreviewDatafeed(datafeedConfig);
64+
Map<String, String> headers = threadPool.getThreadContext().getHeaders().entrySet().stream()
65+
.filter(e -> ClientHelper.SECURITY_HEADER_FILTERS.contains(e.getKey()))
66+
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
67+
previewDatafeed.setHeaders(headers);
68+
// NB: this is using the client from the transport layer, NOT the internal client.
69+
// This is important because it means the datafeed search will fail if the user
70+
// requesting the preview doesn't have permission to search the relevant indices.
71+
DataExtractorFactory.create(client, previewDatafeed.build(), jobBuilder.build(),
72+
new ActionListener<DataExtractorFactory>() {
73+
@Override
74+
public void onResponse(DataExtractorFactory dataExtractorFactory) {
75+
DataExtractor dataExtractor = dataExtractorFactory.newExtractor(0, Long.MAX_VALUE);
76+
threadPool.generic().execute(() -> previewDatafeed(dataExtractor, listener));
77+
}
8678

79+
@Override
80+
public void onFailure(Exception e) {
81+
listener.onFailure(e);
82+
}
83+
});
84+
},
85+
listener::onFailure
86+
));
87+
},
88+
listener::onFailure
89+
));
8790
}
8891

8992
/** Visible for testing */

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportStopDatafeedAction.java‎

Lines changed: 33 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -22,19 +22,17 @@
2222
import org.elasticsearch.common.util.concurrent.AbstractRunnable;
2323
import org.elasticsearch.common.util.concurrent.AtomicArray;
2424
import org.elasticsearch.discovery.MasterNotDiscoveredException;
25+
import org.elasticsearch.persistent.PersistentTasksCustomMetaData;
26+
import org.elasticsearch.persistent.PersistentTasksService;
2527
import org.elasticsearch.tasks.Task;
2628
import org.elasticsearch.threadpool.ThreadPool;
2729
import org.elasticsearch.transport.TransportService;
28-
import org.elasticsearch.xpack.core.ml.MlMetadata;
2930
import org.elasticsearch.xpack.core.ml.MlTasks;
3031
import org.elasticsearch.xpack.core.ml.action.StopDatafeedAction;
31-
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
3232
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedState;
33-
import org.elasticsearch.xpack.core.ml.job.messages.Messages;
3433
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
35-
import org.elasticsearch.persistent.PersistentTasksCustomMetaData;
36-
import org.elasticsearch.persistent.PersistentTasksService;
3734
import org.elasticsearch.xpack.ml.MachineLearning;
35+
import org.elasticsearch.xpack.ml.datafeed.persistence.DatafeedConfigProvider;
3836

3937
import java.io.IOException;
4038
import java.util.ArrayList;
@@ -50,35 +48,35 @@ public class TransportStopDatafeedAction extends TransportTasksAction<TransportS
5048

5149
private final ThreadPool threadPool;
5250
private final PersistentTasksService persistentTasksService;
51+
private final DatafeedConfigProvider datafeedConfigProvider;
5352

5453
@Inject
5554
public TransportStopDatafeedAction(Settings settings, TransportService transportService, ThreadPool threadPool,
5655
ActionFilters actionFilters, ClusterService clusterService,
57-
PersistentTasksService persistentTasksService) {
56+
PersistentTasksService persistentTasksService, DatafeedConfigProvider datafeedConfigProvider) {
5857
super(settings, StopDatafeedAction.NAME, clusterService, transportService, actionFilters,
5958
StopDatafeedAction.Request::new, StopDatafeedAction.Response::new, MachineLearning.UTILITY_THREAD_POOL_NAME);
6059
this.threadPool = threadPool;
6160
this.persistentTasksService = persistentTasksService;
61+
this.datafeedConfigProvider = datafeedConfigProvider;
62+
6263
}
6364

6465
/**
65-
* Resolve the requested datafeeds and add their IDs to one of the list
66-
* arguments depending on datafeed state.
66+
* Sort the datafeed IDs the their task state and add to one
67+
* of the list arguments depending on the state.
6768
*
68-
* @param request The stop datafeed request
69-
* @param mlMetadata ML Metadata
69+
* @param expandedDatafeedIds The expanded set of IDs
7070
* @param tasks Persistent task meta data
7171
* @param startedDatafeedIds Started datafeed ids are added to this list
7272
* @param stoppingDatafeedIds Stopping datafeed ids are added to this list
7373
*/
74-
static void resolveDataFeedIds(StopDatafeedAction.Request request, MlMetadata mlMetadata,
75-
PersistentTasksCustomMetaData tasks,
76-
List<String> startedDatafeedIds,
77-
List<String> stoppingDatafeedIds) {
74+
static void sortDatafeedIdsByTaskState(Set<String> expandedDatafeedIds,
75+
PersistentTasksCustomMetaData tasks,
76+
List<String> startedDatafeedIds,
77+
List<String> stoppingDatafeedIds) {
7878

79-
Set<String> expandedDatafeedIds = mlMetadata.expandDatafeedIds(request.getDatafeedId(), request.allowNoDatafeeds());
8079
for (String expandedDatafeedId : expandedDatafeedIds) {
81-
validateDatafeedTask(expandedDatafeedId, mlMetadata);
8280
addDatafeedTaskIdAccordingToState(expandedDatafeedId, MlTasks.getDatafeedState(expandedDatafeedId, tasks),
8381
startedDatafeedIds, stoppingDatafeedIds);
8482
}
@@ -102,20 +100,6 @@ private static void addDatafeedTaskIdAccordingToState(String datafeedId,
102100
}
103101
}
104102

105-
/**
106-
* Validate the stop request.
107-
* Throws an {@code ResourceNotFoundException} if there is no datafeed
108-
* with id {@code datafeedId}
109-
* @param datafeedId The datafeed Id
110-
* @param mlMetadata ML meta data
111-
*/
112-
static void validateDatafeedTask(String datafeedId, MlMetadata mlMetadata) {
113-
DatafeedConfig datafeed = mlMetadata.getDatafeed(datafeedId);
114-
if (datafeed == null) {
115-
throw new ResourceNotFoundException(Messages.getMessage(Messages.DATAFEED_NOT_FOUND, datafeedId));
116-
}
117-
}
118-
119103
@Override
120104
protected void doExecute(Task task, StopDatafeedAction.Request request, ActionListener<StopDatafeedAction.Response> listener) {
121105
final ClusterState state = clusterService.state();
@@ -130,23 +114,27 @@ protected void doExecute(Task task, StopDatafeedAction.Request request, ActionLi
130114
new ActionListenerResponseHandler<>(listener, StopDatafeedAction.Response::new));
131115
}
132116
} else {
133-
MlMetadata mlMetadata = MlMetadata.getMlMetadata(state);
134-
PersistentTasksCustomMetaData tasks = state.getMetaData().custom(PersistentTasksCustomMetaData.TYPE);
117+
datafeedConfigProvider.expandDatafeedIds(request.getDatafeedId(), request.allowNoDatafeeds(), ActionListener.wrap(
118+
expandedIds -> {
119+
PersistentTasksCustomMetaData tasks = state.getMetaData().custom(PersistentTasksCustomMetaData.TYPE);
135120

136-
List<String> startedDatafeeds = new ArrayList<>();
137-
List<String> stoppingDatafeeds = new ArrayList<>();
138-
resolveDataFeedIds(request, mlMetadata, tasks, startedDatafeeds, stoppingDatafeeds);
139-
if (startedDatafeeds.isEmpty() && stoppingDatafeeds.isEmpty()) {
140-
listener.onResponse(new StopDatafeedAction.Response(true));
141-
return;
142-
}
143-
request.setResolvedStartedDatafeedIds(startedDatafeeds.toArray(new String[startedDatafeeds.size()]));
121+
List<String> startedDatafeeds = new ArrayList<>();
122+
List<String> stoppingDatafeeds = new ArrayList<>();
123+
sortDatafeedIdsByTaskState(expandedIds, tasks, startedDatafeeds, stoppingDatafeeds);
124+
if (startedDatafeeds.isEmpty() && stoppingDatafeeds.isEmpty()) {
125+
listener.onResponse(new StopDatafeedAction.Response(true));
126+
return;
127+
}
128+
request.setResolvedStartedDatafeedIds(startedDatafeeds.toArray(new String[startedDatafeeds.size()]));
144129

145-
if (request.isForce()) {
146-
forceStopDatafeed(request, listener, tasks, startedDatafeeds);
147-
} else {
148-
normalStopDatafeed(task, request, listener, tasks, startedDatafeeds, stoppingDatafeeds);
149-
}
130+
if (request.isForce()) {
131+
forceStopDatafeed(request, listener, tasks, startedDatafeeds);
132+
} else {
133+
normalStopDatafeed(task, request, listener, tasks, startedDatafeeds, stoppingDatafeeds);
134+
}
135+
},
136+
listener::onFailure
137+
));
150138
}
151139
}
152140

0 commit comments

Comments
 (0)