From afaa06df3f7f2be18f94b65b086906ffb38981ff Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Sun, 7 Jan 2024 10:44:09 -0800 Subject: [PATCH] Use timeout in other search requests Signed-off-by: Daniel Widdis --- .../rest/AbstractSearchWorkflowAction.java | 11 ++-- .../CreateWorkflowTransportActionTests.java | 55 ++++++++++--------- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java index e5ad65c43..0e62de62f 100644 --- a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java @@ -48,7 +48,7 @@ public abstract class AbstractSearchWorkflowAction e /** Search action type*/ protected final ActionType actionType; /** Settings to enable FlowFramework API*/ - protected final FlowFrameworkSettings flowFrameworkFeatureEnabledSetting; + protected final FlowFrameworkSettings flowFrameworkSettings; /** * Instantiates a new AbstractSearchWorkflowAction @@ -56,25 +56,25 @@ public abstract class AbstractSearchWorkflowAction e * @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 urlPaths, String index, Class clazz, ActionType 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 index e3a28a12f..9aac97b14 100644 --- 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;