From 99004939b65da928ef8cc32dbe0ccf0c5dafe072 Mon Sep 17 00:00:00 2001
From: Sarat Vemulapalli <vemsarat@amazon.com>
Date: Fri, 30 Oct 2020 00:39:36 -0400
Subject: [PATCH] Stashing context for Stats API to allow users to query from
 RestAPI (#300)

---
 .../StatsAnomalyDetectorTransportAction.java          | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/StatsAnomalyDetectorTransportAction.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/StatsAnomalyDetectorTransportAction.java
index 30bf21b5..463b6d1b 100644
--- a/src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/StatsAnomalyDetectorTransportAction.java
+++ b/src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/StatsAnomalyDetectorTransportAction.java
@@ -19,6 +19,8 @@
 import java.util.Map;
 import java.util.Set;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.search.SearchRequest;
@@ -27,6 +29,7 @@
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
@@ -38,6 +41,7 @@
 import com.amazon.opendistroforelasticsearch.ad.util.MultiResponsesDelegateActionListener;
 
 public class StatsAnomalyDetectorTransportAction extends HandledTransportAction<ADStatsRequest, StatsAnomalyDetectorResponse> {
+    private final Logger logger = LogManager.getLogger(StatsAnomalyDetectorTransportAction.class);
 
     private final Client client;
     private final ADStats adStats;
@@ -60,7 +64,12 @@ public StatsAnomalyDetectorTransportAction(
 
     @Override
     protected void doExecute(Task task, ADStatsRequest request, ActionListener<StatsAnomalyDetectorResponse> listener) {
-        getStats(client, listener, request);
+        try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            getStats(client, listener, request);
+        } catch (Exception e) {
+            logger.error(e);
+            listener.onFailure(e);
+        }
     }
 
     /**