Skip to content

Commit 8b8e3e3

Browse files
davidkyledavidkyle
committed
Store headers for put/update datafeed
1 parent d70d73a commit 8b8e3e3

File tree

11 files changed

+118
-68
lines changed

11 files changed

+118
-68
lines changed

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MlMetadata.java

+6-7
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import org.elasticsearch.common.io.stream.StreamInput;
2121
import org.elasticsearch.common.io.stream.StreamOutput;
2222
import org.elasticsearch.common.io.stream.Writeable;
23-
import org.elasticsearch.common.util.concurrent.ThreadContext;
2423
import org.elasticsearch.common.xcontent.ObjectParser;
2524
import org.elasticsearch.common.xcontent.ToXContent;
2625
import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -293,7 +292,7 @@ public Builder deleteJob(String jobId, PersistentTasksCustomMetaData tasks) {
293292
return this;
294293
}
295294

296-
public Builder putDatafeed(DatafeedConfig datafeedConfig, ThreadContext threadContext) {
295+
public Builder putDatafeed(DatafeedConfig datafeedConfig, Map<String, String> headers) {
297296
if (datafeeds.containsKey(datafeedConfig.getId())) {
298297
throw new ResourceAlreadyExistsException("A datafeed with id [" + datafeedConfig.getId() + "] already exists");
299298
}
@@ -302,13 +301,13 @@ public Builder putDatafeed(DatafeedConfig datafeedConfig, ThreadContext threadCo
302301
Job job = jobs.get(jobId);
303302
DatafeedJobValidator.validate(datafeedConfig, job);
304303

305-
if (threadContext != null) {
304+
if (headers.isEmpty() == false) {
306305
// Adjust the request, adding security headers from the current thread context
307306
DatafeedConfig.Builder builder = new DatafeedConfig.Builder(datafeedConfig);
308-
Map<String, String> headers = threadContext.getHeaders().entrySet().stream()
307+
Map<String, String> securityHeaders = headers.entrySet().stream()
309308
.filter(e -> ClientHelper.SECURITY_HEADER_FILTERS.contains(e.getKey()))
310309
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
311-
builder.setHeaders(headers);
310+
builder.setHeaders(securityHeaders);
312311
datafeedConfig = builder.build();
313312
}
314313

@@ -328,15 +327,15 @@ private void checkJobIsAvailableForDatafeed(String jobId) {
328327
}
329328
}
330329

331-
public Builder updateDatafeed(DatafeedUpdate update, PersistentTasksCustomMetaData persistentTasks, ThreadContext threadContext) {
330+
public Builder updateDatafeed(DatafeedUpdate update, PersistentTasksCustomMetaData persistentTasks, Map<String, String> headers) {
332331
String datafeedId = update.getId();
333332
DatafeedConfig oldDatafeedConfig = datafeeds.get(datafeedId);
334333
if (oldDatafeedConfig == null) {
335334
throw ExceptionsHelper.missingDatafeedException(datafeedId);
336335
}
337336
checkDatafeedIsStopped(() -> Messages.getMessage(Messages.DATAFEED_CANNOT_UPDATE_IN_CURRENT_STATE, datafeedId,
338337
DatafeedState.STARTED), datafeedId, persistentTasks);
339-
DatafeedConfig newDatafeedConfig = update.apply(oldDatafeedConfig, threadContext);
338+
DatafeedConfig newDatafeedConfig = update.apply(oldDatafeedConfig, headers);
340339
if (newDatafeedConfig.getJobId().equals(oldDatafeedConfig.getJobId()) == false) {
341340
checkJobIsAvailableForDatafeed(newDatafeedConfig.getJobId());
342341
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/datafeed/DatafeedUpdate.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,7 @@ ChunkingConfig getChunkingConfig() {
264264
* Applies the update to the given {@link DatafeedConfig}
265265
* @return a new {@link DatafeedConfig} that contains the update
266266
*/
267-
public DatafeedConfig apply(DatafeedConfig datafeedConfig, ThreadContext threadContext) {
267+
public DatafeedConfig apply(DatafeedConfig datafeedConfig, Map<String, String> headers) {
268268
if (id.equals(datafeedConfig.getId()) == false) {
269269
throw new IllegalArgumentException("Cannot apply update to datafeedConfig with different id");
270270
}
@@ -301,12 +301,12 @@ public DatafeedConfig apply(DatafeedConfig datafeedConfig, ThreadContext threadC
301301
builder.setChunkingConfig(chunkingConfig);
302302
}
303303

304-
if (threadContext != null) {
304+
if (headers.isEmpty() == false) {
305305
// Adjust the request, adding security headers from the current thread context
306-
Map<String, String> headers = threadContext.getHeaders().entrySet().stream()
306+
Map<String, String> securityHeaders = headers.entrySet().stream()
307307
.filter(e -> ClientHelper.SECURITY_HEADER_FILTERS.contains(e.getKey()))
308308
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
309-
builder.setHeaders(headers);
309+
builder.setHeaders(securityHeaders);
310310
}
311311

312312
return builder.build();

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/datafeed/DatafeedUpdateTests.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ public void testApply_failBecauseTargetDatafeedHasDifferentId() {
114114

115115
public void testApply_givenEmptyUpdate() {
116116
DatafeedConfig datafeed = DatafeedConfigTests.createRandomizedDatafeedConfig("foo");
117-
DatafeedConfig updatedDatafeed = new DatafeedUpdate.Builder(datafeed.getId()).build().apply(datafeed, null);
117+
DatafeedConfig updatedDatafeed = new DatafeedUpdate.Builder(datafeed.getId()).build().apply(datafeed, Collections.emptyMap());
118118
assertThat(datafeed, equalTo(updatedDatafeed));
119119
}
120120

@@ -125,7 +125,7 @@ public void testApply_givenPartialUpdate() {
125125

126126
DatafeedUpdate.Builder updated = new DatafeedUpdate.Builder(datafeed.getId());
127127
updated.setScrollSize(datafeed.getScrollSize() + 1);
128-
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, null);
128+
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, Collections.emptyMap());
129129

130130
DatafeedConfig.Builder expectedDatafeed = new DatafeedConfig.Builder(datafeed);
131131
expectedDatafeed.setScrollSize(datafeed.getScrollSize() + 1);
@@ -149,7 +149,7 @@ public void testApply_givenFullUpdateNoAggregations() {
149149
update.setScrollSize(8000);
150150
update.setChunkingConfig(ChunkingConfig.newManual(TimeValue.timeValueHours(1)));
151151

152-
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, null);
152+
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, Collections.emptyMap());
153153

154154
assertThat(updatedDatafeed.getJobId(), equalTo("bar"));
155155
assertThat(updatedDatafeed.getIndices(), equalTo(Collections.singletonList("i_2")));
@@ -175,7 +175,7 @@ public void testApply_givenAggregations() {
175175
update.setAggregations(new AggregatorFactories.Builder().addAggregator(
176176
AggregationBuilders.histogram("a").interval(300000).field("time").subAggregation(maxTime)));
177177

178-
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, null);
178+
DatafeedConfig updatedDatafeed = update.build().apply(datafeed, Collections.emptyMap());
179179

180180
assertThat(updatedDatafeed.getIndices(), equalTo(Collections.singletonList("i_1")));
181181
assertThat(updatedDatafeed.getTypes(), equalTo(Collections.singletonList("t_1")));

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

+8-6
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.elasticsearch.xpack.core.security.support.Exceptions;
4242

4343
import java.io.IOException;
44+
import java.util.Map;
4445

4546
public class TransportPutDatafeedAction extends TransportMasterNodeAction<PutDatafeedAction.Request, PutDatafeedAction.Response> {
4647

@@ -95,15 +96,15 @@ protected void masterOperation(PutDatafeedAction.Request request, ClusterState s
9596

9697
client.execute(HasPrivilegesAction.INSTANCE, privRequest, privResponseListener);
9798
} else {
98-
putDatafeed(request, listener);
99+
putDatafeed(request, threadPool.getThreadContext().getHeaders(), listener);
99100
}
100101
}
101102

102103
private void handlePrivsResponse(String username, PutDatafeedAction.Request request,
103104
HasPrivilegesResponse response,
104105
ActionListener<PutDatafeedAction.Response> listener) throws IOException {
105106
if (response.isCompleteMatch()) {
106-
putDatafeed(request, listener);
107+
putDatafeed(request, threadPool.getThreadContext().getHeaders(), listener);
107108
} else {
108109
XContentBuilder builder = JsonXContent.contentBuilder();
109110
builder.startObject();
@@ -120,7 +121,8 @@ private void handlePrivsResponse(String username, PutDatafeedAction.Request requ
120121
}
121122
}
122123

123-
private void putDatafeed(PutDatafeedAction.Request request, ActionListener<PutDatafeedAction.Response> listener) {
124+
private void putDatafeed(PutDatafeedAction.Request request, Map<String, String> headers,
125+
ActionListener<PutDatafeedAction.Response> listener) {
124126

125127
clusterService.submitStateUpdateTask(
126128
"put-datafeed-" + request.getDatafeed().getId(),
@@ -136,16 +138,16 @@ protected PutDatafeedAction.Response newResponse(boolean acknowledged) {
136138

137139
@Override
138140
public ClusterState execute(ClusterState currentState) {
139-
return putDatafeed(request, currentState);
141+
return putDatafeed(request, headers, currentState);
140142
}
141143
});
142144
}
143145

144-
private ClusterState putDatafeed(PutDatafeedAction.Request request, ClusterState clusterState) {
146+
private ClusterState putDatafeed(PutDatafeedAction.Request request, Map<String, String> headers, ClusterState clusterState) {
145147
XPackPlugin.checkReadyForXPackCustomMetadata(clusterState);
146148
MlMetadata currentMetadata = MlMetadata.getMlMetadata(clusterState);
147149
MlMetadata newMetadata = new MlMetadata.Builder(currentMetadata)
148-
.putDatafeed(request.getDatafeed(), threadPool.getThreadContext()).build();
150+
.putDatafeed(request.getDatafeed(), headers).build();
149151
return ClusterState.builder(clusterState).metaData(
150152
MetaData.builder(clusterState.getMetaData()).putCustom(MLMetadataField.TYPE, newMetadata).build())
151153
.build();

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

+5-1
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedUpdate;
2828
import org.elasticsearch.persistent.PersistentTasksCustomMetaData;
2929

30+
import java.util.Map;
31+
3032
public class TransportUpdateDatafeedAction extends TransportMasterNodeAction<UpdateDatafeedAction.Request, PutDatafeedAction.Response> {
3133

3234
@Inject
@@ -50,6 +52,8 @@ protected PutDatafeedAction.Response newResponse() {
5052
@Override
5153
protected void masterOperation(UpdateDatafeedAction.Request request, ClusterState state,
5254
ActionListener<PutDatafeedAction.Response> listener) {
55+
final Map<String, String> headers = threadPool.getThreadContext().getHeaders();
56+
5357
clusterService.submitStateUpdateTask("update-datafeed-" + request.getUpdate().getId(),
5458
new AckedClusterStateUpdateTask<PutDatafeedAction.Response>(request, listener) {
5559
private volatile DatafeedConfig updatedDatafeed;
@@ -69,7 +73,7 @@ public ClusterState execute(ClusterState currentState) {
6973
PersistentTasksCustomMetaData persistentTasks =
7074
currentState.getMetaData().custom(PersistentTasksCustomMetaData.TYPE);
7175
MlMetadata newMetadata = new MlMetadata.Builder(currentMetadata)
72-
.updateDatafeed(update, persistentTasks, threadPool.getThreadContext()).build();
76+
.updateDatafeed(update, persistentTasks, headers).build();
7377
updatedDatafeed = newMetadata.getDatafeed(update.getId());
7478
return ClusterState.builder(currentState).metaData(
7579
MetaData.builder(currentState.getMetaData()).putCustom(MLMetadataField.TYPE, newMetadata).build()).build();

0 commit comments

Comments
 (0)