Change search config api request (opensearch-project#1089)

Signed-off-by: Joanne Wang <jowg@amazon.com>
sbcd90 · Jun 20, 2024 · 1927009 · 1927009
1 parent c28f84f
commit 1927009
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 35 deletions.
diff --git a/.../org/opensearch/securityanalytics/threatIntel/action/SASearchTIFSourceConfigsRequest.java b/.../org/opensearch/securityanalytics/threatIntel/action/SASearchTIFSourceConfigsRequest.java
@@ -7,9 +7,9 @@
 
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
-import org.opensearch.action.search.SearchRequest;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.search.builder.SearchSourceBuilder;
 
 import java.io.IOException;
 
@@ -19,29 +19,29 @@
 public class SASearchTIFSourceConfigsRequest extends ActionRequest {
 
     // TODO: add pagination parameters
-    private SearchRequest searchRequest;
+    private final SearchSourceBuilder searchSourceBuilder;
 
-    public SASearchTIFSourceConfigsRequest(SearchRequest searchRequest) {
+    public SASearchTIFSourceConfigsRequest(SearchSourceBuilder searchSourceBuilder) {
         super();
-        this.searchRequest = searchRequest;
+        this.searchSourceBuilder = searchSourceBuilder;
     }
 
     public SASearchTIFSourceConfigsRequest(StreamInput sin) throws IOException {
-        searchRequest = new SearchRequest(sin);
+        searchSourceBuilder = new SearchSourceBuilder(sin);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        searchRequest.writeTo(out);
+        searchSourceBuilder.writeTo(out);
     }
 
     @Override
     public ActionRequestValidationException validate() {
         return null;
     }
 
-    public SearchRequest getSearchRequest() {
-        return searchRequest;
+    public SearchSourceBuilder getSearchSourceBuilder() {
+        return searchSourceBuilder;
     }
 
 }
diff --git a/...pensearch/securityanalytics/threatIntel/resthandler/RestSearchTIFSourceConfigsAction.java b/...pensearch/securityanalytics/threatIntel/resthandler/RestSearchTIFSourceConfigsAction.java
@@ -2,33 +2,20 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.node.NodeClient;
-import org.opensearch.cluster.routing.Preference;
-import org.opensearch.common.xcontent.LoggingDeprecationHandler;
-import org.opensearch.common.xcontent.XContentFactory;
-import org.opensearch.common.xcontent.XContentType;
-import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.ToXContent;
-import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.index.query.BoolQueryBuilder;
-import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestChannel;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
-import org.opensearch.search.SearchHit;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.fetch.subphase.FetchSourceContext;
 import org.opensearch.securityanalytics.SecurityAnalyticsPlugin;
-import org.opensearch.securityanalytics.action.CorrelatedFindingResponse;
 import org.opensearch.securityanalytics.threatIntel.action.SASearchTIFSourceConfigsAction;
 import org.opensearch.securityanalytics.threatIntel.action.SASearchTIFSourceConfigsRequest;
-import org.opensearch.securityanalytics.threatIntel.model.SATIFSourceConfigDto;
 
 import java.io.IOException;
 import java.util.List;
@@ -56,14 +43,11 @@ public List<Route> routes() {
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         log.debug(String.format(Locale.getDefault(), "%s %s", request.method(), SecurityAnalyticsPlugin.THREAT_INTEL_SOURCE_URI + "/" + "_search"));
 
-        // TODO: Change request to take in a BoolQueryBuilder
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         searchSourceBuilder.parseXContent(request.contentOrSourceParamParser());
         searchSourceBuilder.fetchSource(FetchSourceContext.parseFromRestRequest(request));
-        searchSourceBuilder.seqNoAndPrimaryTerm(true);
-        searchSourceBuilder.version(true);
 
-        SASearchTIFSourceConfigsRequest req = new SASearchTIFSourceConfigsRequest(new SearchRequest().source(searchSourceBuilder));
+        SASearchTIFSourceConfigsRequest req = new SASearchTIFSourceConfigsRequest(searchSourceBuilder);
 
         return channel -> client.execute(
                 SASearchTIFSourceConfigsAction.INSTANCE,

diff --git a/.../opensearch/securityanalytics/threatIntel/service/SATIFSourceConfigManagementService.java b/.../opensearch/securityanalytics/threatIntel/service/SATIFSourceConfigManagementService.java
@@ -3,8 +3,6 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchException;
-import org.opensearch.OpenSearchStatusException;
-import org.opensearch.ResourceNotFoundException;
 import org.opensearch.action.delete.DeleteResponse;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -15,7 +13,6 @@
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
-import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -26,6 +23,7 @@
 import org.opensearch.jobscheduler.spi.LockModel;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.search.SearchHit;
+import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.securityanalytics.SecurityAnalyticsPlugin;
 import org.opensearch.securityanalytics.services.STIX2IOCFetchService;
 import org.opensearch.securityanalytics.threatIntel.common.TIFJobState;
@@ -35,7 +33,6 @@
 import org.opensearch.securityanalytics.threatIntel.model.SATIFSourceConfigDto;
 
 import java.time.Instant;
-import java.util.Locale;
 
 /**
  * Service class for threat intel feed source config object
@@ -168,14 +165,14 @@ public void getTIFSourceConfig(
     }
 
     public void searchTIFSourceConfigs(
-            final SearchRequest searchRequest,
+            final SearchSourceBuilder searchSourceBuilder,
             final ActionListener<SearchResponse> listener
     ) {
         try {
-            SearchRequest newSearchRequest = getSearchRequest(searchRequest);
+            SearchRequest searchRequest = getSearchRequest(searchSourceBuilder);
 
             // convert search response to threat intel source config dtos
-            saTifSourceConfigService.searchTIFSourceConfigs(newSearchRequest, ActionListener.wrap(
+            saTifSourceConfigService.searchTIFSourceConfigs(searchRequest, ActionListener.wrap(
                     searchResponse -> {
                         for (SearchHit hit: searchResponse.getHits()) {
                             XContentParser xcp = XContentType.JSON.xContent().createParser(
@@ -188,7 +185,7 @@ public void searchTIFSourceConfigs(
                         }
                         listener.onResponse(searchResponse);
                     }, e -> {
-                        log.error("Failed to fetch all threat intel source configs for search request [{}]", searchRequest, e);
+                        log.error("Failed to fetch all threat intel source configs for search request [{}]", searchSourceBuilder, e);
                         listener.onFailure(e);
                     }
             ));
@@ -198,7 +195,14 @@ public void searchTIFSourceConfigs(
         }
     }
 
-    private static SearchRequest getSearchRequest(SearchRequest searchRequest) {
+    private static SearchRequest getSearchRequest(SearchSourceBuilder searchSourceBuilder) {
+
+        // update search source builder
+        searchSourceBuilder.seqNoAndPrimaryTerm(true);
+        searchSourceBuilder.version(true);
+
+        // construct search request
+        SearchRequest searchRequest = new SearchRequest().source(searchSourceBuilder);
         searchRequest.indices(SecurityAnalyticsPlugin.JOB_INDEX_NAME);
         searchRequest.preference(Preference.PRIMARY_FIRST.type());
 

diff --git a/...search/securityanalytics/threatIntel/transport/TransportSearchTIFSourceConfigsAction.java b/...search/securityanalytics/threatIntel/transport/TransportSearchTIFSourceConfigsAction.java
@@ -64,7 +64,7 @@ protected void doExecute(Task task, SASearchTIFSourceConfigsRequest request, Act
 
         this.threadPool.getThreadContext().stashContext(); // TODO: sync up with @deysubho about thread context
 
-        saTifConfigService.searchTIFSourceConfigs(request.getSearchRequest(), ActionListener.wrap(
+        saTifConfigService.searchTIFSourceConfigs(request.getSearchSourceBuilder(), ActionListener.wrap(
                 r -> {
                     log.debug("Successfully listed all threat intel source configs");
                     actionListener.onResponse(r);