Use timeout in other search requests

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jan 8, 2024 · afaa06d · afaa06d
1 parent 1c5ded7
commit afaa06d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 32 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java
@@ -48,33 +48,33 @@ public abstract class AbstractSearchWorkflowAction<T extends ToXContentObject> e
     /** Search action type*/
     protected final ActionType<SearchResponse> actionType;
     /** Settings to enable FlowFramework API*/
-    protected final FlowFrameworkSettings flowFrameworkFeatureEnabledSetting;
+    protected final FlowFrameworkSettings flowFrameworkSettings;
 
     /**
      * Instantiates a new AbstractSearchWorkflowAction
      * @param urlPaths urlPaths to create routes
      * @param index index the search should be done on
      * @param clazz model class
      * @param actionType from which action abstract class is called
-     * @param flowFrameworkFeatureEnabledSetting Whether this API is enabled
+     * @param flowFrameworkSettings Whether this API is enabled
      */
     public AbstractSearchWorkflowAction(
         List<String> urlPaths,
         String index,
         Class<T> clazz,
         ActionType<SearchResponse> actionType,
-        FlowFrameworkSettings flowFrameworkFeatureEnabledSetting
+        FlowFrameworkSettings flowFrameworkSettings
     ) {
         this.urlPaths = urlPaths;
         this.index = index;
         this.clazz = clazz;
         this.actionType = actionType;
-        this.flowFrameworkFeatureEnabledSetting = flowFrameworkFeatureEnabledSetting;
+        this.flowFrameworkSettings = flowFrameworkSettings;
     }
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
-        if (!flowFrameworkFeatureEnabledSetting.isFlowFrameworkEnabled()) {
+        if (!flowFrameworkSettings.isFlowFrameworkEnabled()) {
             FlowFrameworkException ffe = new FlowFrameworkException(
                 "This API is disabled. To enable it, update the setting [" + FLOW_FRAMEWORK_ENABLED.getKey() + "] to true.",
                 RestStatus.FORBIDDEN
@@ -87,6 +87,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         searchSourceBuilder.parseXContent(request.contentOrSourceParamParser());
         searchSourceBuilder.fetchSource(getSourceContext(request, searchSourceBuilder));
         searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true);
+        searchSourceBuilder.timeout(flowFrameworkSettings.getRequestTimeout());
         SearchRequest searchRequest = new SearchRequest().source(searchSourceBuilder).indices(index);
         return channel -> client.execute(actionType, searchRequest, search(channel));
     }

diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
@@ -8,33 +8,6 @@
  */
 package org.opensearch.flowframework.transport;
 
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doNothing;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-import static org.opensearch.action.DocWriteResponse.Result.UPDATED;
-import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX;
-import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_STATE_INDEX;
-import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
-import static org.opensearch.flowframework.common.WorkflowResources.CREATE_CONNECTOR;
-import static org.opensearch.flowframework.common.WorkflowResources.DEPLOY_MODEL;
-import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
-import static org.opensearch.flowframework.common.WorkflowResources.REGISTER_REMOTE_MODEL;
-
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
-
-import org.mockito.ArgumentCaptor;
 import org.opensearch.Version;
 import org.opensearch.action.index.IndexResponse;
 import org.opensearch.action.support.ActionFilters;
@@ -59,6 +32,34 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import org.mockito.ArgumentCaptor;
+
+import static org.opensearch.action.DocWriteResponse.Result.UPDATED;
+import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX;
+import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_STATE_INDEX;
+import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
+import static org.opensearch.flowframework.common.WorkflowResources.CREATE_CONNECTOR;
+import static org.opensearch.flowframework.common.WorkflowResources.DEPLOY_MODEL;
+import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
+import static org.opensearch.flowframework.common.WorkflowResources.REGISTER_REMOTE_MODEL;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 public class CreateWorkflowTransportActionTests extends OpenSearchTestCase {
 
     private CreateWorkflowTransportAction createWorkflowTransportAction;