Fixes #30573: search index=mlModelService no longer returns mlModel assets

openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/MlModelServiceSearchAliasIT.java

@@ -0,0 +1,134 @@
+/*
+ *  Copyright 2024 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.openmetadata.it.tests;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.time.Duration;
+import org.awaitility.Awaitility;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.parallel.Execution;
+import org.junit.jupiter.api.parallel.ExecutionMode;
+import org.openmetadata.it.factories.MlModelServiceTestFactory;
+import org.openmetadata.it.util.SdkClients;
+import org.openmetadata.it.util.TestNamespace;
+import org.openmetadata.it.util.TestNamespaceExtension;
+import org.openmetadata.schema.api.data.CreateMlModel;
+import org.openmetadata.schema.entity.data.MlModel;
+import org.openmetadata.schema.entity.services.MlModelService;
+import org.openmetadata.sdk.client.OpenMetadataClient;
+import org.openmetadata.service.Entity;
+
+/**
+ * Regression test for the ML model service search alias bug: querying {@code index=mlModelService}
+ * must return only ML model <em>services</em>, never ML model <em>assets</em>.
+ *
+ * <p>The mlModel asset lists {@code mlModelService} as a {@code parentAlias}, so the ES alias
+ * {@code mlModelService} is attached to both {@code mlmodel_service_search_index} and
+ * {@code mlmodel_search_index}. For every other {@code *Service} the search resolver collapses the
+ * alias token to the single concrete service index, but the ML service's {@code entityIndexMap} key
+ * ({@code mlmodelService}) differs in casing from its alias ({@code mlModelService}) that clients
+ * send, so the by-key lookup missed and the token fell through to ES-native alias expansion —
+ * pulling in assets. Resolving the token by the entity's alias (not just its key) fixes it.
+ */
+@Execution(ExecutionMode.CONCURRENT)
+@ExtendWith(TestNamespaceExtension.class)
+public class MlModelServiceSearchAliasIT {
+
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+  private static final Duration POLL_AT_MOST = Duration.ofSeconds(120);
+  private static final Duration POLL_INTERVAL = Duration.ofMillis(500);
+  private static final String MLMODEL_SERVICE_INDEX = "mlModelService";
+  private static final String MLMODEL_INDEX = "mlmodel";
+
+  @Test
+  @DisplayName("index=mlModelService returns only services, not mlModel assets")
+  void mlModelServiceAliasDoesNotLeakAssets(TestNamespace ns) throws Exception {
+    OpenMetadataClient client = SdkClients.adminClient();
+
+    MlModelService service = MlModelServiceTestFactory.createMlflow(ns);
+    MlModel asset =
+        ns.trackRoot(
+            Entity.MLMODEL,
+            client
+                .mlModels()
+                .create(
+                    new CreateMlModel()
+                        .withName(ns.prefix("mlmodel_alias_asset"))
+                        .withService(service.getFullyQualifiedName())
+                        .withAlgorithm("regression")));
+    String serviceId = service.getId().toString();
+    String assetId = asset.getId().toString();
+
+    JsonNode assetHit = awaitHitById(client, MLMODEL_INDEX, assetId);
+    assertEquals(
+        Entity.MLMODEL,
+        assetHit.path("entityType").asText(),
+        "asset must be reachable via its own index with entityType 'mlmodel'");
+
+    JsonNode serviceHit = awaitHitById(client, MLMODEL_SERVICE_INDEX, serviceId);
+    assertEquals(
+        Entity.MLMODEL_SERVICE,
+        serviceHit.path("entityType").asText(),
+        "service must be reachable via the mlModelService alias");
+
+    JsonNode leakedAsset = findHitById(client, MLMODEL_SERVICE_INDEX, assetId);
+    assertNull(
+        leakedAsset,
+        "index=mlModelService must not return the mlModel asset " + assetId + " (alias leak)");
+  }
+
+  /**
+   * Poll {@code index} for the document with id {@code id} until it appears, then return its
+   * {@code _source}. Matching by id (a keyword field) is analyzer-independent, so it behaves
+   * identically on Elasticsearch and OpenSearch. Indexing is async after the write API returns, so a
+   * fixed sleep flakes; this waits for convergence.
+   */
+  private JsonNode awaitHitById(OpenMetadataClient client, String index, String id) {
+    JsonNode[] match = new JsonNode[1];
+    Awaitility.await(id + " indexed in " + index)
+        .pollInterval(POLL_INTERVAL)
+        .atMost(POLL_AT_MOST)
+        .ignoreExceptions()
+        .untilAsserted(
+            () -> {
+              JsonNode source = findHitById(client, index, id);
+              assertNotNull(source, id + " not yet indexed in " + index);
+              match[0] = source;
+            });
+    return match[0];
+  }
+
+  private JsonNode findHitById(OpenMetadataClient client, String index, String id)
+      throws Exception {
+    String response =
+        client.search().query("id:" + id).index(index).size(5).deleted(false).execute();
+    JsonNode hits = OBJECT_MAPPER.readTree(response).path("hits").path("hits");
+    JsonNode result = null;
+    for (JsonNode hit : hits) {
+      JsonNode source = hit.path("_source");
+      if (id.equals(source.path("id").asText(""))) {
+        result = source;
+        break;
+      }
+    }
+    return result;
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java

@@ -409,6 +409,7 @@ public static void deferOrRunSearchWrite(
   }
 
   @Getter private Map<String, IndexMapping> entityIndexMap;
+  private Map<String, IndexMapping> aliasIndexMap;
 
   /**
    * Staged index names being populated by an in-flight reindex, keyed by the canonical index name
@@ -538,6 +539,21 @@ public SearchClient getSearchClient() {
   private void loadIndexMappings() {
     IndexMappingLoader mappingLoader = IndexMappingLoader.getInstance();
     entityIndexMap = mappingLoader.getIndexMapping();
+    aliasIndexMap = buildAliasIndexMap(entityIndexMap);
+  }
+
+  private static Map<String, IndexMapping> buildAliasIndexMap(
+      Map<String, IndexMapping> mappingsByKey) {
+    Map<String, IndexMapping> mappingsByAlias = new HashMap<>();
+    if (mappingsByKey != null) {
+      for (IndexMapping mapping : mappingsByKey.values()) {
+        String alias = mapping.getAlias(null);
+        if (alias != null && !mappingsByKey.containsKey(alias)) {
+          mappingsByAlias.putIfAbsent(alias, mapping);
+        }
+      }
+    }
+    return mappingsByAlias;
   }
 
   public SearchClient buildSearchClient(ElasticSearchConfiguration config) {
@@ -999,6 +1015,9 @@ private String resolveSingleAliasToken(String token, String clusterPrefix) {
       return token;
     }
     IndexMapping mapping = entityIndexMap == null ? null : entityIndexMap.get(token);
+    if (mapping == null && aliasIndexMap != null) {
+      mapping = aliasIndexMap.get(token);
+    }
     if (mapping != null) {
       return mapping.getIndexName(clusterAlias);
     }

openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRepositoryBehaviorTest.java

@@ -138,6 +138,14 @@ class SearchRepositoryBehaviorTest {
           .indexMappingFile("/elasticsearch/%s/database_service_index_mapping.json")
           .build();
 
+  private static final IndexMapping MLMODEL_SERVICE_MAPPING =
+      IndexMapping.builder()
+          .indexName("mlmodel_service_search_index")
+          .alias("mlModelService")
+          .childAliases(List.of("mlmodel"))
+          .indexMappingFile("/elasticsearch/%s/mlmodel_service_index_mapping.json")
+          .build();
+
   private static final IndexMapping DATABASE_MAPPING =
       IndexMapping.builder()
           .indexName("database_search_index")
@@ -223,6 +231,7 @@ void setUp() {
                 Map.entry(Entity.DOMAIN, DOMAIN_MAPPING),
                 Map.entry(Entity.DATA_PRODUCT, DATA_PRODUCT_MAPPING),
                 Map.entry(Entity.DATABASE_SERVICE, DATABASE_SERVICE_MAPPING),
+                Map.entry(Entity.MLMODEL_SERVICE, MLMODEL_SERVICE_MAPPING),
                 Map.entry(Entity.TAG, TABLE_MAPPING),
                 Map.entry(Entity.GLOSSARY_TERM, TABLE_MAPPING),
                 Map.entry(Entity.GLOSSARY, TABLE_MAPPING),
@@ -387,6 +396,20 @@ void getIndexOrAliasNameResolvesEntitySpecificAliasToCanonicalIndex() {
     assertEquals("cluster_domain_search_index", repository.getIndexOrAliasName("domain"));
   }
 
+  /**
+   * When an entity's {@code entityIndexMap} key differs from its {@code alias} (the mlModel service
+   * key is {@code mlmodelService} but its alias is {@code mlModelService}), a query for the alias
+   * must still resolve to the single canonical index. Without alias resolution the token misses the
+   * by-key lookup, passes through as a raw ES alias, and fans out to every index that carries it —
+   * for {@code mlModelService} that includes {@code mlmodel_search_index}, so the response leaks
+   * mlModel assets alongside the services.
+   */
+  @Test
+  void getIndexOrAliasNameResolvesEntityAliasWhenKeyCasingDiffers() {
+    assertEquals(
+        "cluster_mlmodel_service_search_index", repository.getIndexOrAliasName("mlModelService"));
+  }
+
   /**
    * Compound aliases like {@code "all"} and {@code "dataAsset"} have no entry in
    * {@code entityIndexMap} (they're meta-aliases registered against many entities at index
@@ -3470,6 +3493,7 @@ void repositoryMetadataHelpersExposeUnderlyingState() throws Exception {
             Entity.DOMAIN,
             Entity.DATA_PRODUCT,
             Entity.DATABASE_SERVICE,
+            Entity.MLMODEL_SERVICE,
             Entity.TAG,
             Entity.GLOSSARY_TERM,
             Entity.GLOSSARY,