From ef820ac15b0af90ea25b3c5e621724b49758f891 Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Thu, 29 Aug 2024 17:37:06 -0700
Subject: [PATCH] Fix Connector Update, Search, and Delete tenant awareness

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../connector/DeleteConnectorTransportAction.java |  1 +
 .../connector/UpdateConnectorTransportAction.java | 15 ++++++---------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java
index 22b8e8be5b..ddef704bfa 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/connector/DeleteConnectorTransportAction.java
@@ -129,6 +129,7 @@ private void checkForModelsUsingConnector(String connectorId, String tenantId, A
             SearchDataObjectRequest searchDataObjectRequest = SearchDataObjectRequest
                 .builder()
                 .indices(ML_MODEL_INDEX)
+                .tenantId(tenantId)
                 .searchSourceBuilder(sourceBuilder)
                 .build();
             sdkClient
diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/UpdateConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/UpdateConnectorTransportAction.java
index 362761393a..0b8b83ca98 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/connector/UpdateConnectorTransportAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/connector/UpdateConnectorTransportAction.java
@@ -104,25 +104,20 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
             return;
         }
         String connectorId = mlUpdateConnectorAction.getConnectorId();
+        String tenantId = mlCreateConnectorInput.getTenantId();
         FetchSourceContext fetchSourceContext = new FetchSourceContext(true, Strings.EMPTY_ARRAY, Strings.EMPTY_ARRAY);
         GetDataObjectRequest getDataObjectRequest = GetDataObjectRequest
             .builder()
             .index(ML_CONNECTOR_INDEX)
             .id(connectorId)
+            .tenantId(tenantId)
             .fetchSourceContext(fetchSourceContext)
             .build();
-
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             connectorAccessControlHelper
                 .getConnector(sdkClient, client, context, getDataObjectRequest, connectorId, ActionListener.wrap(connector -> {
                     // context is already restored here
-                    if (TenantAwareHelper
-                        .validateTenantResource(
-                            mlFeatureEnabledSetting,
-                            mlCreateConnectorInput.getTenantId(),
-                            connector.getTenantId(),
-                            listener
-                        )) {
+                    if (TenantAwareHelper.validateTenantResource(mlFeatureEnabledSetting, tenantId, connector.getTenantId(), listener)) {
                         boolean hasPermission = connectorAccessControlHelper.validateConnectorAccess(client, connector);
                         if (hasPermission) {
                             connector.update(mlUpdateConnectorAction.getUpdateContent(), mlEngine::encrypt);
@@ -131,6 +126,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
                                 .builder()
                                 .index(ML_CONNECTOR_INDEX)
                                 .id(connectorId)
+                                .tenantId(tenantId)
                                 .dataObject(connector)
                                 .build();
                             try (ThreadContext.StoredContext innerContext = client.threadPool().getThreadContext().stashContext()) {
@@ -173,6 +169,7 @@ private void updateUndeployedConnector(
         SearchDataObjectRequest searchDataObjectRequest = SearchDataObjectRequest
             .builder()
             .indices(ML_MODEL_INDEX)
+            .tenantId(updateDataObjectRequest.tenantId())
             .searchSourceBuilder(sourceBuilder)
             .build();
         sdkClient
@@ -211,7 +208,6 @@ private void updateUndeployedConnector(
                     }
                 } else {
                     Exception cause = SdkClientUtils.unwrapAndConvertToException(st);
-                    log.error("Failed to update ML connector: " + connectorId, cause);
                     if (cause instanceof IndexNotFoundException) {
                         sdkClient
                             .updateDataObjectAsync(updateDataObjectRequest, client.threadPool().executor(GENERAL_THREAD_POOL))
@@ -220,6 +216,7 @@ private void updateUndeployedConnector(
                             });
                         return;
                     } else {
+                        log.error("Failed to update ML connector: " + connectorId, cause);
                         listener.onFailure(cause);
                     }
                 }