Fix Connector Update, Search, and Delete tenant awareness

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Sep 1, 2024 · ef820ac · ef820ac
1 parent ab045f6
commit ef820ac
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 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
@@ -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
@@ -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);
                     }
                 }