Adds Ability to Extend search functionalty (#17508)

* Adds Ability to Extend search functionalty

* Fix Tests

openmetadata-service/src/main/java/org/openmetadata/service/Entity.java

@@ -596,4 +596,8 @@ public final class Entity {
   public static <T> T getDao() {
     return (T) collectionDAO;
   }
+
+  public static <T> T getSearchRepo() {
+    return (T) searchRepository;
+  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -65,6 +65,7 @@ import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 import org.openmetadata.schema.api.security.ClientType;
 import org.openmetadata.schema.configuration.LimitsConfiguration;
+import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
 import org.openmetadata.schema.services.connections.metadata.AuthProvider;
 import org.openmetadata.service.apps.ApplicationHandler;
 import org.openmetadata.service.apps.scheduler.AppScheduler;
@@ -164,9 +165,7 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
     jdbi = createAndSetupJDBI(environment, catalogConfig.getDataSourceFactory());
     Entity.setCollectionDAO(getDao(jdbi));
 
-    // initialize Search Repository, all repositories use SearchRepository this line should always
+    installSearchRepository(catalogConfig.getElasticSearchConfiguration());
-    // before initializing repository
-    new SearchRepository(catalogConfig.getElasticSearchConfiguration());
     // Initialize the MigrationValidationClient, used in the Settings Repository
     MigrationValidationClient.initialize(jdbi.onDemand(MigrationDAO.class), catalogConfig);
     // as first step register all the repositories
@@ -302,6 +301,13 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
     }
   }
 
+  protected void installSearchRepository(ElasticSearchConfiguration esConfig) {
+    // initialize Search Repository, all repositories use SearchRepository this line should always
+    // before initializing repository
+    SearchRepository searchRepository = new SearchRepository(esConfig);
+    Entity.setSearchRepository(searchRepository);
+  }
+
   private void registerHealthCheck(Environment environment) {
     environment
         .healthChecks()

openmetadata-service/src/main/java/org/openmetadata/service/resources/apps/AppResource.java

@@ -115,7 +115,7 @@ public class AppResource extends EntityResource<App, AppRepository> {
 
       // Create an On Demand DAO
       CollectionDAO dao = Entity.getCollectionDAO();
-      searchRepository = new SearchRepository(config.getElasticSearchConfiguration());
+      searchRepository = Entity.getSearchRepository();
       AppScheduler.initialize(config, dao, searchRepository);
 
       // Initialize Default Apps

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

@@ -135,6 +135,11 @@ public interface SearchClient {
       String entityType)
       throws IOException;
 
+  default Response listPageHierarchy(String parent) {
+    throw new CustomExceptionMessage(
+        Response.Status.NOT_IMPLEMENTED, NOT_IMPLEMENTED_ERROR_TYPE, NOT_IMPLEMENTED_METHOD);
+  }
+
   Map<String, Object> searchLineageInternal(
       String fqn,
       int upstreamDepth,

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

@@ -109,30 +109,14 @@ public class SearchRepository {
 
   public SearchRepository(ElasticSearchConfiguration config) {
     elasticSearchConfiguration = config;
-    if (config != null
+    searchClient = buildSearchClient(config);
-        && config.getSearchType() == ElasticSearchConfiguration.SearchType.OPENSEARCH) {
+    searchIndexFactory = buildIndexFactory();
-      searchClient = new OpenSearchClient(config);
-    } else {
-      searchClient = new ElasticSearchClient(config);
-    }
-    try {
-      if (config != null && (!nullOrEmpty(config.getSearchIndexFactoryClassName()))) {
-        this.searchIndexFactory =
-            Class.forName(config.getSearchIndexFactoryClassName())
-                .asSubclass(SearchIndexFactory.class)
-                .getDeclaredConstructor()
-                .newInstance();
-      }
-    } catch (Exception e) {
-      LOG.warn("Failed to initialize search index factory using default one", e);
-    }
     language =
         config != null && config.getSearchIndexMappingLanguage() != null
             ? config.getSearchIndexMappingLanguage().value()
             : "en";
     clusterAlias = config != null ? config.getClusterAlias() : "";
     loadIndexMappings();
-    Entity.setSearchRepository(this);
   }
 
   private void loadIndexMappings() {
@@ -164,6 +148,21 @@ public class SearchRepository {
     }
   }
 
+  public SearchClient buildSearchClient(ElasticSearchConfiguration config) {
+    SearchClient sc;
+    if (config != null
+        && config.getSearchType() == ElasticSearchConfiguration.SearchType.OPENSEARCH) {
+      sc = new OpenSearchClient(config);
+    } else {
+      sc = new ElasticSearchClient(config);
+    }
+    return sc;
+  }
+
+  public SearchIndexFactory buildIndexFactory() {
+    return new SearchIndexFactory();
+  }
+
   public ElasticSearchConfiguration.SearchType getSearchType() {
     return searchClient.getSearchType();
   }
@@ -881,4 +880,8 @@ public class SearchRepository {
     }
     return new ArrayList<>();
   }
+
+  public <T> T getRestHighLevelClient() {
+    return (T) searchClient;
+  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java

@@ -208,7 +208,7 @@ import org.openmetadata.service.workflows.searchIndex.ReindexingUtil;
 public class ElasticSearchClient implements SearchClient {
 
   @SuppressWarnings("deprecated")
-  private final RestHighLevelClient client;
+  protected final RestHighLevelClient client;
 
   private final boolean isClientAvailable;
   public static final NamedXContentRegistry xContentRegistry;

openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java

@@ -205,7 +205,7 @@ import os.org.opensearch.search.suggest.completion.context.CategoryQueryContext;
 @Slf4j
 // Not tagged with Repository annotation as it is programmatically initialized
 public class OpenSearchClient implements SearchClient {
-  private final RestHighLevelClient client;
+  protected final RestHighLevelClient client;
   public static final NamedXContentRegistry X_CONTENT_REGISTRY;
   private final boolean isClientAvailable;
 

openmetadata-service/src/test/java/org/openmetadata/service/OpenMetadataApplicationTest.java

@@ -232,7 +232,9 @@ public abstract class OpenMetadataApplicationTest {
             pipelineServiceClientConfiguration,
             forceMigrations);
     // Initialize search repository
-    new SearchRepository(config.getElasticSearchConfiguration());
+    SearchRepository searchRepository =
+        new SearchRepository(config.getElasticSearchConfiguration());
+    Entity.setSearchRepository(searchRepository);
     Entity.setCollectionDAO(jdbi.onDemand(CollectionDAO.class));
     Entity.initializeRepositories(config, jdbi);
     workflow.loadMigrations();