From 9e311bebcbfa0692a9671c1afad3e751ff72f6f0 Mon Sep 17 00:00:00 2001 From: Benjamin Trent Date: Fri, 30 Jul 2021 09:03:20 -0400 Subject: [PATCH] [ML] ensure tasks API executes through ML_ORIGIN (#75880) --- .../xpack/ml/action/TransportSetUpgradeModeAction.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportSetUpgradeModeAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportSetUpgradeModeAction.java index 4bdce5d3bb5c3..928c2ba50cfe7 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportSetUpgradeModeAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportSetUpgradeModeAction.java @@ -17,6 +17,7 @@ import org.elasticsearch.action.support.master.AcknowledgedResponse; import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction; import org.elasticsearch.client.Client; +import org.elasticsearch.client.OriginSettingClient; import org.elasticsearch.cluster.AckedClusterStateUpdateTask; import org.elasticsearch.cluster.ClusterState; import org.elasticsearch.cluster.block.ClusterBlockException; @@ -50,7 +51,6 @@ import static org.elasticsearch.ExceptionsHelper.rethrowAndSuppress; import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; -import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; import static org.elasticsearch.xpack.core.ml.MlTasks.AWAITING_UPGRADE; import static org.elasticsearch.xpack.core.ml.MlTasks.DATAFEED_TASK_NAME; import static org.elasticsearch.xpack.core.ml.MlTasks.JOB_TASK_NAME; @@ -67,7 +67,7 @@ public class TransportSetUpgradeModeAction extends AcknowledgedTransportMasterNo private final PersistentTasksClusterService persistentTasksClusterService; private final PersistentTasksService persistentTasksService; private final ClusterService clusterService; - private final Client client; + private final OriginSettingClient client; @Inject public TransportSetUpgradeModeAction(TransportService transportService, ThreadPool threadPool, ClusterService clusterService, @@ -78,7 +78,7 @@ public TransportSetUpgradeModeAction(TransportService transportService, ThreadPo indexNameExpressionResolver, ThreadPool.Names.SAME); this.persistentTasksClusterService = persistentTasksClusterService; this.clusterService = clusterService; - this.client = client; + this.client = new OriginSettingClient(client, ML_ORIGIN); this.persistentTasksService = persistentTasksService; } @@ -289,7 +289,7 @@ private void isolateDatafeeds(PersistentTasksCustomMetadata tasksCustomMetadata, datafeedsToIsolate.forEach(datafeedId -> { IsolateDatafeedAction.Request isolationRequest = new IsolateDatafeedAction.Request(datafeedId); isolateDatafeedsExecutor.add(isolateListener -> - executeAsyncWithOrigin(client, ML_ORIGIN, IsolateDatafeedAction.INSTANCE, isolationRequest, isolateListener) + client.execute(IsolateDatafeedAction.INSTANCE, isolationRequest, isolateListener) ); });