add sdk implementation to the connector search (#3704)

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

Author: dhrubo-os

plugin/src/main/java/org/opensearch/ml/action/connector/SearchConnectorTransportAction.java

@@ -5,7 +5,6 @@
 
 package org.opensearch.ml.action.connector;
 
-import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
 import static org.opensearch.ml.utils.RestActionUtils.wrapListenerToHandleSearchIndexNotFound;
 
 import java.util.ArrayList;
@@ -15,6 +14,7 @@
 import java.util.stream.Collectors;
 
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.ShardSearchFailure;
@@ -24,9 +24,8 @@
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.index.IndexNotFoundException;
-import org.opensearch.index.query.BoolQueryBuilder;
-import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.connector.HttpConnector;
 import org.opensearch.ml.common.transport.connector.MLConnectorSearchAction;
@@ -36,6 +35,8 @@
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.SdkClient;
+import org.opensearch.remote.metadata.client.SearchDataObjectRequest;
+import org.opensearch.remote.metadata.common.SdkClientUtils;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.fetch.subphase.FetchSourceContext;
 import org.opensearch.search.internal.InternalSearchResponse;
@@ -108,22 +109,34 @@ private void search(SearchRequest request, String tenantId, ActionListener<Searc
             final ActionListener<SearchResponse> doubleWrappedListener = ActionListener
                 .wrap(wrappedListener::onResponse, e -> wrapListenerToHandleSearchIndexNotFound(e, wrappedListener));
 
-            if (tenantId != null) {
-                BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
-                if (request.source().query() != null) {
-                    queryBuilder.must(request.source().query());
-                }
-                queryBuilder.filter(QueryBuilders.termQuery(TENANT_ID_FIELD, tenantId)); // Replace with your tenant_id field
-                request.source().query(queryBuilder);
-            }
-
-            if (connectorAccessControlHelper.skipConnectorAccessControl(user)) {
-                client.search(request, doubleWrappedListener);
-            } else {
+            if (!connectorAccessControlHelper.skipConnectorAccessControl(user)) {
                 SearchSourceBuilder sourceBuilder = connectorAccessControlHelper.addUserBackendRolesFilter(user, request.source());
                 request.source(sourceBuilder);
-                client.search(request, doubleWrappedListener);
             }
+
+            SearchDataObjectRequest searchDataObjectRequest = SearchDataObjectRequest
+                .builder()
+                .indices(request.indices())
+                .searchSourceBuilder(request.source())
+                .tenantId(tenantId)
+                .build();
+            sdkClient.searchDataObjectAsync(searchDataObjectRequest).whenComplete((r, throwable) -> {
+                if (throwable != null) {
+                    Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
+                    log.error("Failed to search connector", cause);
+                    doubleWrappedListener.onFailure(cause);
+                } else {
+                    try {
+                        SearchResponse searchResponse = SearchResponse.fromXContent(r.parser());
+                        log.info("Connector search complete: {}", searchResponse.getHits().getTotalHits());
+                        doubleWrappedListener.onResponse(searchResponse);
+                    } catch (Exception e) {
+                        log.error("Failed to parse search response", e);
+                        doubleWrappedListener
+                            .onFailure(new OpenSearchStatusException("Failed to parse search response", RestStatus.INTERNAL_SERVER_ERROR));
+                    }
+                }
+            });
         } catch (Exception e) {
             log.error(e.getMessage(), e);
             actionListener.onFailure(e);