feat(search): unified entity index (#14966)

Author: david-leifker

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/mappers/MapperUtilsTest.java

@@ -1,19 +1,17 @@
 package com.linkedin.datahub.graphql.types.mappers;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertThrows;
 
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.TestUtils;
 import com.linkedin.datahub.graphql.generated.MatchedField;
 import com.linkedin.metadata.entity.validation.ValidationApiUtils;
 import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.snapshot.Snapshot;
-import io.datahubproject.test.metadata.context.TestOperationContexts;
 import java.net.URISyntaxException;
 import java.util.List;
 import org.testng.annotations.BeforeTest;
@@ -44,9 +42,7 @@ public void testMatchedFieldValidation() throws URISyntaxException {
         IllegalArgumentException.class,
         () -> ValidationApiUtils.validateUrn(entityRegistry, invalidUrn));
 
-    QueryContext mockContext = mock(QueryContext.class);
-    when(mockContext.getOperationContext())
-        .thenReturn(TestOperationContexts.systemContextNoSearchAuthorization(entityRegistry));
+    QueryContext mockContext = TestUtils.getMockAllowContext();
 
     List<MatchedField> actualMatched =
         MapperUtils.getMatchedFieldEntry(

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/loadindices/LoadIndicesIndexManager.java

@@ -59,21 +59,25 @@ public LoadIndicesIndexManager(
   public List<ReindexConfig> discoverDataHubIndexConfigs() throws IOException {
     List<ReindexConfig> configs = new ArrayList<>();
 
-    // Get entity indices using IndexConvention pattern
-    String entityPattern = indexConvention.getAllEntityIndicesPattern();
-    log.debug("Querying entity indices with pattern: {}", entityPattern);
-    GetIndexRequest entityRequest = new GetIndexRequest(entityPattern);
-    GetIndexResponse entityResponse = searchClient.getIndex(entityRequest, RequestOptions.DEFAULT);
-    String[] entityIndices = entityResponse.getIndices();
-
-    for (String indexName : entityIndices) {
-      try {
-        ReindexConfig config = indexBuilder.buildReindexState(indexName, Map.of(), Map.of());
-        configs.add(config);
-        log.debug("Added entity index config: {}", indexName);
-      } catch (IOException e) {
-        log.warn(
-            "Failed to build reindex config for entity index {}: {}", indexName, e.getMessage());
+    // Get entity indices using IndexConvention patterns
+    List<String> entityPatterns = indexConvention.getAllEntityIndicesPatterns();
+    log.debug("Querying entity indices with patterns: {}", entityPatterns);
+
+    for (String entityPattern : entityPatterns) {
+      GetIndexRequest entityRequest = new GetIndexRequest(entityPattern);
+      GetIndexResponse entityResponse =
+          searchClient.getIndex(entityRequest, RequestOptions.DEFAULT);
+      String[] entityIndices = entityResponse.getIndices();
+
+      for (String indexName : entityIndices) {
+        try {
+          ReindexConfig config = indexBuilder.buildReindexState(indexName, Map.of(), Map.of());
+          configs.add(config);
+          log.debug("Added entity index config: {}", indexName);
+        } catch (IOException e) {
+          log.warn(
+              "Failed to build reindex config for entity index {}: {}", indexName, e.getMessage());
+        }
       }
     }
 

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/cron/steps/TweakReplicasStep.java

@@ -57,7 +57,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
       CronArgs args = createArgs(context);
       try {
         for (ElasticSearchIndexed service : services) {
-          service.tweakReplicasAll(structuredProperties, args.dryRun);
+          service.tweakReplicasAll(context.opContext(), structuredProperties, args.dryRun);
         }
       } catch (Exception e) {
         log.error("TweakReplicasStep failed.", e);

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPostStep.java

@@ -50,7 +50,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
       try {
 
         List<ReindexConfig> indexConfigs =
-            getAllReindexConfigs(services, structuredProperties).stream()
+            getAllReindexConfigs(context.opContext(), services, structuredProperties).stream()
                 .filter(ReindexConfig::requiresReindex)
                 .collect(Collectors.toList());
 

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPreStep.java

@@ -53,7 +53,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
     return (context) -> {
       try {
         final List<ReindexConfig> reindexConfigs =
-            getAllReindexConfigs(services, structuredProperties);
+            getAllReindexConfigs(context.opContext(), services, structuredProperties);
 
         // Get indices to update
         List<ReindexConfig> indexConfigs =

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesStep.java

@@ -37,7 +37,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
     return (context) -> {
       try {
         for (ElasticSearchIndexed service : services) {
-          service.reindexAll(structuredProperties);
+          service.reindexAll(context.opContext(), structuredProperties);
         }
       } catch (Exception e) {
         log.error("BuildIndicesStep failed.", e);

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/CleanIndicesStep.java

@@ -51,7 +51,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
     return (context) -> {
       try {
         // Each service enumerates and cleans up their own indices
-        IndexUtils.getAllReindexConfigs(indexedServices, structuredProperties)
+        IndexUtils.getAllReindexConfigs(context.opContext(), indexedServices, structuredProperties)
             .forEach(
                 reindexConfig ->
                     ESIndexBuilder.cleanOrphanedIndices(searchClient, esConfig, reindexConfig));

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/CreateUsageEventIndicesStep.java

@@ -47,8 +47,8 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
             configurationProvider.getElasticSearch().getIndex().getFinalPrefix();
 
         boolean useOpenSearch = esComponents.getSearchClient().getEngineType().isOpenSearch();
-        int numShards = esComponents.getIndexBuilder().getNumShards();
-        int numReplicas = esComponents.getIndexBuilder().getNumReplicas();
+        int numShards = configurationProvider.getElasticSearch().getIndex().getNumShards();
+        int numReplicas = configurationProvider.getElasticSearch().getIndex().getNumReplicas();
 
         if (useOpenSearch) {
           setupOpenSearchUsageEvents(indexPrefix, numShards, numReplicas, context.opContext());

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/ReindexDebugStep.java

@@ -12,6 +12,7 @@
 import com.linkedin.structured.StructuredPropertyDefinition;
 import com.linkedin.upgrade.DataHubUpgradeState;
 import com.linkedin.util.Pair;
+import io.datahubproject.metadata.context.OperationContext;
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
@@ -80,7 +81,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
           log.error("ReindexDebugStep failed: No ElasticSearchService found");
           return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
         }
-        setConfig(args.index);
+        setConfig(context.opContext(), args.index);
         if (config == null) {
           log.error("ReindexDebugStep failed: No matching config found for index: {}", args.index);
           return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
@@ -99,10 +100,11 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
     };
   }
 
-  void setConfig(String targetIndex) throws IOException, IllegalAccessException {
+  void setConfig(OperationContext opContext, String targetIndex)
+      throws IOException, IllegalAccessException {
     // datahubpolicyindex_v2 has some docs upon starting quickdebug...
     //  String targetIndex = "datahubpolicyindex_v2";
-    List<ReindexConfig> configs = service.buildReindexConfigs(structuredProperties);
+    List<ReindexConfig> configs = service.buildReindexConfigs(opContext, structuredProperties);
     config = null; // Reset config to null
     for (ReindexConfig cfg : configs) {
       String cfgname = cfg.name();

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/util/IndexUtils.java

@@ -54,14 +54,16 @@ private IndexUtils() {}
   private static List<ReindexConfig> _reindexConfigs = new ArrayList<>();
 
   public static List<ReindexConfig> getAllReindexConfigs(
+      OperationContext opContext,
       List<ElasticSearchIndexed> elasticSearchIndexedList,
       Collection<Pair<Urn, StructuredPropertyDefinition>> structuredProperties)
       throws IOException {
     // Avoid locking & reprocessing
     List<ReindexConfig> reindexConfigs = new ArrayList<>(_reindexConfigs);
     if (reindexConfigs.isEmpty()) {
       for (ElasticSearchIndexed elasticSearchIndexed : elasticSearchIndexedList) {
-        reindexConfigs.addAll(elasticSearchIndexed.buildReindexConfigs(structuredProperties));
+        reindexConfigs.addAll(
+            elasticSearchIndexed.buildReindexConfigs(opContext, structuredProperties));
       }
       _reindexConfigs = new ArrayList<>(reindexConfigs);
     }