[ML] Check Search Inference ID Usage When Deleting An Inference Endpoint (#113895) (#114012)

Author: Mikep86

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/utils/SemanticTextInfoExtractor.java

@@ -33,15 +33,11 @@ public static Set<String> extractIndexesReferencingInferenceEndpoints(Metadata m
         Map<String, IndexMetadata> indices = metadata.indices();
 
         indices.forEach((indexName, indexMetadata) -> {
-            if (indexMetadata.getInferenceFields() != null) {
-                Map<String, InferenceFieldMetadata> inferenceFields = indexMetadata.getInferenceFields();
-                if (inferenceFields.entrySet()
-                    .stream()
-                    .anyMatch(
-                        entry -> entry.getValue().getInferenceId() != null && endpointIds.contains(entry.getValue().getInferenceId())
-                    )) {
-                    referenceIndices.add(indexName);
-                }
+            Map<String, InferenceFieldMetadata> inferenceFields = indexMetadata.getInferenceFields();
+            if (inferenceFields.values()
+                .stream()
+                .anyMatch(im -> endpointIds.contains(im.getInferenceId()) || endpointIds.contains(im.getSearchInferenceId()))) {
+                referenceIndices.add(indexName);
             }
         });
 

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -168,6 +168,26 @@ protected void putSemanticText(String endpointId, String indexName) throws IOExc
         assertOkOrCreated(response);
     }
 
+    protected void putSemanticText(String endpointId, String searchEndpointId, String indexName) throws IOException {
+        var request = new Request("PUT", Strings.format("%s", indexName));
+        String body = Strings.format("""
+            {
+                "mappings": {
+                "properties": {
+                    "inference_field": {
+                        "type": "semantic_text",
+                            "inference_id": "%s",
+                            "search_inference_id": "%s"
+                    }
+                }
+                }
+            }
+            """, endpointId, searchEndpointId);
+        request.setJsonEntity(body);
+        var response = client().performRequest(request);
+        assertOkOrCreated(response);
+    }
+
     protected Map<String, Object> putModel(String modelId, String modelConfig, TaskType taskType) throws IOException {
         String endpoint = Strings.format("_inference/%s/%s", taskType, modelId);
         return putRequest(endpoint, modelConfig);

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java

@@ -17,6 +17,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
@@ -154,34 +155,55 @@ public void testDeleteEndpointWhileReferencedByPipeline() throws IOException {
     }
 
     public void testDeleteEndpointWhileReferencedBySemanticText() throws IOException {
-        String endpointId = "endpoint_referenced_by_semantic_text";
+        final String endpointId = "endpoint_referenced_by_semantic_text";
+        final String searchEndpointId = "search_endpoint_referenced_by_semantic_text";
+        final String indexName = randomAlphaOfLength(10).toLowerCase();
+        final Function<String, String> buildErrorString = endpointName -> " Inference endpoint "
+            + endpointName
+            + " is being used in the mapping for indexes: "
+            + Set.of(indexName)
+            + ". Ensure that no index mappings are using this inference endpoint, or use force to ignore this warning and delete the"
+            + " inference endpoint.";
+
         putModel(endpointId, mockSparseServiceModelConfig(), TaskType.SPARSE_EMBEDDING);
-        String indexName = randomAlphaOfLength(10).toLowerCase();
         putSemanticText(endpointId, indexName);
         {
-
-            var errorString = new StringBuilder().append(" Inference endpoint ")
-                .append(endpointId)
-                .append(" is being used in the mapping for indexes: ")
-                .append(Set.of(indexName))
-                .append(". ")
-                .append("Ensure that no index mappings are using this inference endpoint, ")
-                .append("or use force to ignore this warning and delete the inference endpoint.");
             var e = expectThrows(ResponseException.class, () -> deleteModel(endpointId));
-            assertThat(e.getMessage(), containsString(errorString.toString()));
+            assertThat(e.getMessage(), containsString(buildErrorString.apply(endpointId)));
         }
         {
             var response = deleteModel(endpointId, "dry_run=true");
             var entityString = EntityUtils.toString(response.getEntity());
             assertThat(entityString, containsString("\"acknowledged\":false"));
             assertThat(entityString, containsString(indexName));
+            assertThat(entityString, containsString(endpointId));
         }
         {
             var response = deleteModel(endpointId, "force=true");
             var entityString = EntityUtils.toString(response.getEntity());
             assertThat(entityString, containsString("\"acknowledged\":true"));
         }
         deleteIndex(indexName);
+
+        putModel(searchEndpointId, mockSparseServiceModelConfig(), TaskType.SPARSE_EMBEDDING);
+        putSemanticText(endpointId, searchEndpointId, indexName);
+        {
+            var e = expectThrows(ResponseException.class, () -> deleteModel(searchEndpointId));
+            assertThat(e.getMessage(), containsString(buildErrorString.apply(searchEndpointId)));
+        }
+        {
+            var response = deleteModel(searchEndpointId, "dry_run=true");
+            var entityString = EntityUtils.toString(response.getEntity());
+            assertThat(entityString, containsString("\"acknowledged\":false"));
+            assertThat(entityString, containsString(indexName));
+            assertThat(entityString, containsString(searchEndpointId));
+        }
+        {
+            var response = deleteModel(searchEndpointId, "force=true");
+            var entityString = EntityUtils.toString(response.getEntity());
+            assertThat(entityString, containsString("\"acknowledged\":true"));
+        }
+        deleteIndex(indexName);
     }
 
     public void testDeleteEndpointWhileReferencedBySemanticTextAndPipeline() throws IOException {
@@ -217,6 +239,7 @@ public void testDeleteEndpointWhileReferencedBySemanticTextAndPipeline() throws
             assertThat(entityString, containsString("\"acknowledged\":false"));
             assertThat(entityString, containsString(indexName));
             assertThat(entityString, containsString(pipelineId));
+            assertThat(entityString, containsString(endpointId));
         }
         {
             var response = deleteModel(endpointId, "force=true");