Refactor Repositories: Remove passing around CollectionDAO (#13427)

* Refactor repository registration

* Refactor repository registration

* Refactor repository registration

* Refactor repository registration

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

@@ -31,7 +31,6 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
 import javax.ws.rs.core.UriInfo;
@@ -66,7 +65,7 @@
 @Slf4j
 public final class Entity {
   private static volatile boolean initializedRepositories = false;
-  @Getter private static volatile CollectionDAO collectionDAO;
+  @Getter @Setter private static CollectionDAO collectionDAO;
   public static final String SEPARATOR = "."; // Fully qualified name separator
 
   // Canonical entity name to corresponding EntityRepository map
@@ -242,19 +241,16 @@ public final class Entity {
 
   private Entity() {}
 
-  public static void initializeRepositories(Jdbi jdbi, CollectionDAO collectionDAO) {
+  public static void initializeRepositories(Jdbi jdbi) {
     if (!initializedRepositories) {
-      Entity.collectionDAO = collectionDAO;
-      tokenRepository = new TokenRepository(collectionDAO);
-      // Check Collection DAO
-      Objects.requireNonNull(collectionDAO, "CollectionDAO must not be null");
+      tokenRepository = new TokenRepository();
       List<Class<?>> repositories = getRepositories();
       for (Class<?> clz : repositories) {
         if (Modifier.isAbstract(clz.getModifiers())) {
           continue; // Don't instantiate abstract classes
         }
         try {
-          clz.getDeclaredConstructor(CollectionDAO.class).newInstance(collectionDAO);
+          clz.getDeclaredConstructor().newInstance();
         } catch (Exception e) {
           try {
             clz.getDeclaredConstructor(Jdbi.class).newInstance(jdbi);
@@ -270,6 +266,8 @@ public static void initializeRepositories(Jdbi jdbi, CollectionDAO collectionDAO
 
   public static void cleanup() {
     initializedRepositories = false;
+    collectionDAO = null;
+    searchRepository = null;
     ENTITY_REPOSITORY_MAP.clear();
   }
 
@@ -385,7 +383,6 @@ public static <T> T getEntity(String entityType, UUID id, String fields, Include
     return getEntity(entityType, id, fields, include, true);
   }
 
-  // TODO remove throwing IOException
   /** Retrieve the entity using id from given entity reference and fields */
   public static <T> T getEntityByName(
       String entityType, String fqn, String fields, Include include, boolean fromCache) {

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

@@ -126,7 +126,6 @@
 public class OpenMetadataApplication extends Application<OpenMetadataApplicationConfig> {
   private Authorizer authorizer;
   private AuthenticatorHandler authenticatorHandler;
-  private static CollectionDAO collectionDAO;
 
   private static SearchRepository searchRepository;
 
@@ -143,16 +142,16 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ
     ChangeEventConfig.initialize(catalogConfig);
     final Jdbi jdbi = createAndSetupJDBI(environment, catalogConfig.getDataSourceFactory());
     JdbiUnitOfWorkProvider jdbiUnitOfWorkProvider = JdbiUnitOfWorkProvider.withDefault(jdbi);
-    collectionDAO = (CollectionDAO) getWrappedInstanceForDaoClass(CollectionDAO.class);
     JdbiTransactionManager.initialize(jdbiUnitOfWorkProvider.getHandleManager());
+    CollectionDAO collectionDAO = (CollectionDAO) getWrappedInstanceForDaoClass(CollectionDAO.class);
+    Entity.setCollectionDAO(collectionDAO);
     environment.jersey().register(new JdbiUnitOfWorkApplicationEventListener(new HashSet<>()));
 
     // initialize Search Repository, all repositories use SearchRepository this line should always before initializing
     // repository
-    searchRepository = new SearchRepository(catalogConfig.getElasticSearchConfiguration(), collectionDAO);
-
+    searchRepository = new SearchRepository(catalogConfig.getElasticSearchConfiguration());
     // as first step register all the repositories
-    Entity.initializeRepositories(jdbi, collectionDAO);
+    Entity.initializeRepositories(jdbi);
 
     // Init Settings Cache after repositories
     SettingsCache.initialize(catalogConfig);
@@ -202,10 +201,10 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ
     // start event hub before registering publishers
     EventPubSub.start();
 
-    registerResources(catalogConfig, environment, jdbi, collectionDAO);
+    registerResources(catalogConfig, environment, jdbi);
 
     // Register Event Handler
-    registerEventFilter(catalogConfig, environment, jdbiUnitOfWorkProvider);
+    registerEventFilter(catalogConfig, environment);
     environment.lifecycle().manage(new ManagedShutdown());
     // Register Event publishers
     registerEventPublisher(catalogConfig);
@@ -215,10 +214,10 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ
 
     // start authorizer after event publishers
     // authorizer creates admin/bot users, ES publisher should start before to index users created by authorizer
-    authorizer.init(catalogConfig, collectionDAO);
+    authorizer.init(catalogConfig);
 
     // authenticationHandler Handles auth related activities
-    authenticatorHandler.init(catalogConfig, collectionDAO);
+    authenticatorHandler.init(catalogConfig);
 
     webAnalyticEvents = MicrometerBundleSingleton.latencyTimer(catalogConfig.getEventMonitorConfiguration());
     FilterRegistration.Dynamic micrometerFilter =
@@ -405,10 +404,9 @@ private void registerAuthenticator(OpenMetadataApplicationConfig catalogConfig)
     }
   }
 
-  private void registerEventFilter(
-      OpenMetadataApplicationConfig catalogConfig, Environment environment, JdbiUnitOfWorkProvider provider) {
+  private void registerEventFilter(OpenMetadataApplicationConfig catalogConfig, Environment environment) {
     if (catalogConfig.getEventHandlerConfiguration() != null) {
-      ContainerResponseFilter eventFilter = new EventFilter(catalogConfig, provider);
+      ContainerResponseFilter eventFilter = new EventFilter(catalogConfig);
       environment.jersey().register(eventFilter);
       ContainerResponseFilter reindexingJobs = new SearchIndexEvent();
       environment.jersey().register(reindexingJobs);
@@ -428,11 +426,9 @@ private void registerEventPublisher(OpenMetadataApplicationConfig openMetadataAp
     }
   }
 
-  private void registerResources(
-      OpenMetadataApplicationConfig config, Environment environment, Jdbi jdbi, CollectionDAO daoObject) {
+  private void registerResources(OpenMetadataApplicationConfig config, Environment environment, Jdbi jdbi) {
     CollectionRegistry.initialize();
-    CollectionRegistry.getInstance()
-        .registerResources(jdbi, environment, config, daoObject, authorizer, authenticatorHandler);
+    CollectionRegistry.getInstance().registerResources(jdbi, environment, config, authorizer, authenticatorHandler);
     environment.jersey().register(new JsonPatchProvider());
     OMErrorPageHandler eph = new OMErrorPageHandler(config.getWebConfiguration());
     eph.addErrorPage(Response.Status.NOT_FOUND.getStatusCode(), "/");

openmetadata-service/src/main/java/org/openmetadata/service/events/ChangeEventHandler.java

@@ -58,7 +58,7 @@ public Void process(ContainerRequestContext requestContext, ContainerResponseCon
     try {
       CollectionDAO collectionDAO = (CollectionDAO) getWrappedInstanceForDaoClass(CollectionDAO.class);
       CollectionDAO.ChangeEventDAO changeEventDAO = collectionDAO.changeEventDAO();
-      FeedRepository feedRepository = new FeedRepository(collectionDAO);
+      FeedRepository feedRepository = new FeedRepository();
       if (responseContext.getEntity() != null && responseContext.getEntity().getClass().equals(Thread.class)) {
         // we should move this to Email Application notifications instead of processing it here.
         notificationHandler.processNotifications(responseContext);

openmetadata-service/src/main/java/org/openmetadata/service/events/EventFilter.java

@@ -26,7 +26,6 @@
 import javax.ws.rs.ext.Provider;
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.service.OpenMetadataApplicationConfig;
-import org.openmetadata.service.jdbi3.unitofwork.JdbiUnitOfWorkProvider;
 import org.openmetadata.service.security.JwtFilter;
 import org.openmetadata.service.util.ParallelStreamUtil;
 
@@ -38,13 +37,13 @@ public class EventFilter implements ContainerResponseFilter {
   private final ForkJoinPool forkJoinPool;
   private final List<EventHandler> eventHandlers;
 
-  public EventFilter(OpenMetadataApplicationConfig config, JdbiUnitOfWorkProvider provider) {
+  public EventFilter(OpenMetadataApplicationConfig config) {
     this.forkJoinPool = new ForkJoinPool(FORK_JOIN_POOL_PARALLELISM);
     this.eventHandlers = new ArrayList<>();
-    registerEventHandlers(config, provider);
+    registerEventHandlers(config);
   }
 
-  private void registerEventHandlers(OpenMetadataApplicationConfig config, JdbiUnitOfWorkProvider provider) {
+  private void registerEventHandlers(OpenMetadataApplicationConfig config) {
     try {
       Set<String> eventHandlerClassNames =
           new HashSet<>(config.getEventHandlerConfiguration().getEventHandlerClassNames());

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/BotRepository.java

@@ -27,8 +27,9 @@
 public class BotRepository extends EntityRepository<Bot> {
   static final String BOT_UPDATE_FIELDS = "botUser";
 
-  public BotRepository(CollectionDAO dao) {
-    super(BotResource.COLLECTION_PATH, Entity.BOT, Bot.class, dao.botDAO(), dao, "", BOT_UPDATE_FIELDS);
+  public BotRepository() {
+    super(
+        BotResource.COLLECTION_PATH, Entity.BOT, Bot.class, Entity.getCollectionDAO().botDAO(), "", BOT_UPDATE_FIELDS);
     quoteFqn = true;
   }
 

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ChangeEventRepository.java

@@ -26,8 +26,8 @@
 public class ChangeEventRepository {
   private final CollectionDAO.ChangeEventDAO dao;
 
-  public ChangeEventRepository(CollectionDAO dao) {
-    this.dao = dao.changeEventDAO();
+  public ChangeEventRepository() {
+    this.dao = Entity.getCollectionDAO().changeEventDAO();
     Entity.setChangeEventRepository(this);
   }
 

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ChartRepository.java

@@ -30,8 +30,8 @@
 
 @Slf4j
 public class ChartRepository extends EntityRepository<Chart> {
-  public ChartRepository(CollectionDAO dao) {
-    super(ChartResource.COLLECTION_PATH, Entity.CHART, Chart.class, dao.chartDAO(), dao, "", "");
+  public ChartRepository() {
+    super(ChartResource.COLLECTION_PATH, Entity.CHART, Chart.class, Entity.getCollectionDAO().chartDAO(), "", "");
     supportsSearch = true;
   }
 

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ClassificationRepository.java

@@ -37,13 +37,12 @@
 
 @Slf4j
 public class ClassificationRepository extends EntityRepository<Classification> {
-  public ClassificationRepository(CollectionDAO dao) {
+  public ClassificationRepository() {
     super(
         ClassificationResource.TAG_COLLECTION_PATH,
         Entity.CLASSIFICATION,
         Classification.class,
-        dao.classificationDAO(),
-        dao,
+        Entity.getCollectionDAO().classificationDAO(),
         "",
         "");
     quoteFqn = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContainerRepository.java

@@ -37,13 +37,12 @@ public class ContainerRepository extends EntityRepository<Container> {
   private static final String CONTAINER_UPDATE_FIELDS = "dataModel";
   private static final String CONTAINER_PATCH_FIELDS = "dataModel";
 
-  public ContainerRepository(CollectionDAO dao) {
+  public ContainerRepository() {
     super(
         ContainerResource.COLLECTION_PATH,
         Entity.CONTAINER,
         Container.class,
-        dao.containerDAO(),
-        dao,
+        Entity.getCollectionDAO().containerDAO(),
         CONTAINER_PATCH_FIELDS,
         CONTAINER_UPDATE_FIELDS);
     supportsSearch = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardDataModelRepository.java

@@ -44,13 +44,12 @@
 
 @Slf4j
 public class DashboardDataModelRepository extends EntityRepository<DashboardDataModel> {
-  public DashboardDataModelRepository(CollectionDAO dao) {
+  public DashboardDataModelRepository() {
     super(
         DashboardDataModelResource.COLLECTION_PATH,
         Entity.DASHBOARD_DATA_MODEL,
         DashboardDataModel.class,
-        dao.dashboardDataModelDAO(),
-        dao,
+        Entity.getCollectionDAO().dashboardDataModelDAO(),
         "",
         "");
     supportsSearch = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardRepository.java

@@ -44,13 +44,12 @@ public class DashboardRepository extends EntityRepository<Dashboard> {
   private static final String DASHBOARD_PATCH_FIELDS = "charts,dataModels";
   private static final String DASHBOARD_URL = "sourceUrl";
 
-  public DashboardRepository(CollectionDAO dao) {
+  public DashboardRepository() {
     super(
         DashboardResource.COLLECTION_PATH,
         Entity.DASHBOARD,
         Dashboard.class,
-        dao.dashboardDAO(),
-        dao,
+        Entity.getCollectionDAO().dashboardDAO(),
         DASHBOARD_PATCH_FIELDS,
         DASHBOARD_UPDATE_FIELDS);
     supportsSearch = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardServiceRepository.java

@@ -23,12 +23,11 @@
 @Slf4j
 public class DashboardServiceRepository extends ServiceEntityRepository<DashboardService, DashboardConnection> {
 
-  public DashboardServiceRepository(CollectionDAO dao) {
+  public DashboardServiceRepository() {
     super(
         DashboardServiceResource.COLLECTION_PATH,
         Entity.DASHBOARD_SERVICE,
-        dao,
-        dao.dashboardServiceDAO(),
+        Entity.getCollectionDAO().dashboardServiceDAO(),
         DashboardConnection.class,
         "",
         ServiceType.DASHBOARD);

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataInsightChartRepository.java

@@ -6,6 +6,7 @@
 import java.util.List;
 import org.openmetadata.schema.dataInsight.DataInsightChart;
 import org.openmetadata.schema.dataInsight.DataInsightChartResult;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.util.EntityUtil;
 
 public class DataInsightChartRepository extends EntityRepository<DataInsightChart> {
@@ -63,8 +64,14 @@ public class DataInsightChartRepository extends EntityRepository<DataInsightChar
   public static final List<String> SUPPORTS_NULL_DATE_RANGE =
       Arrays.asList(DataInsightChartResult.DataInsightChartType.UNUSED_ASSETS.toString());
 
-  public DataInsightChartRepository(CollectionDAO dao) {
-    super(COLLECTION_PATH, DATA_INSIGHT_CHART, DataInsightChart.class, dao.dataInsightChartDAO(), dao, "", "");
+  public DataInsightChartRepository() {
+    super(
+        COLLECTION_PATH,
+        DATA_INSIGHT_CHART,
+        DataInsightChart.class,
+        Entity.getCollectionDAO().dataInsightChartDAO(),
+        "",
+        "");
   }
 
   @Override

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataProductRepository.java

@@ -34,13 +34,12 @@
 public class DataProductRepository extends EntityRepository<DataProduct> {
   private static final String UPDATE_FIELDS = "experts,assets"; // Domain field can't be updated
 
-  public DataProductRepository(CollectionDAO dao) {
+  public DataProductRepository() {
     super(
         DataProductResource.COLLECTION_PATH,
         Entity.DATA_PRODUCT,
         DataProduct.class,
-        dao.dataProductDAO(),
-        dao,
+        Entity.getCollectionDAO().dataProductDAO(),
         UPDATE_FIELDS,
         UPDATE_FIELDS);
     supportsSearch = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseRepository.java

@@ -32,8 +32,14 @@
 
 @Slf4j
 public class DatabaseRepository extends EntityRepository<Database> {
-  public DatabaseRepository(CollectionDAO dao) {
-    super(DatabaseResource.COLLECTION_PATH, Entity.DATABASE, Database.class, dao.databaseDAO(), dao, "", "");
+  public DatabaseRepository() {
+    super(
+        DatabaseResource.COLLECTION_PATH,
+        Entity.DATABASE,
+        Database.class,
+        Entity.getCollectionDAO().databaseDAO(),
+        "",
+        "");
     supportsSearch = true;
   }
 

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseSchemaRepository.java

@@ -32,13 +32,12 @@
 
 @Slf4j
 public class DatabaseSchemaRepository extends EntityRepository<DatabaseSchema> {
-  public DatabaseSchemaRepository(CollectionDAO dao) {
+  public DatabaseSchemaRepository() {
     super(
         DatabaseSchemaResource.COLLECTION_PATH,
         Entity.DATABASE_SCHEMA,
         DatabaseSchema.class,
-        dao.databaseSchemaDAO(),
-        dao,
+        Entity.getCollectionDAO().databaseSchemaDAO(),
         "",
         "");
     supportsSearch = true;

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseServiceRepository.java

@@ -22,12 +22,11 @@
 
 @Slf4j
 public class DatabaseServiceRepository extends ServiceEntityRepository<DatabaseService, DatabaseConnection> {
-  public DatabaseServiceRepository(CollectionDAO dao) {
+  public DatabaseServiceRepository() {
     super(
         DatabaseServiceResource.COLLECTION_PATH,
         Entity.DATABASE_SERVICE,
-        dao,
-        dao.dbServiceDAO(),
+        Entity.getCollectionDAO().dbServiceDAO(),
         DatabaseConnection.class,
         "",
         ServiceType.DATABASE);

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DocumentRepository.java

@@ -17,6 +17,7 @@
 
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.schema.entities.docStore.Document;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.resources.docstore.DocStoreResource;
 import org.openmetadata.service.util.EntityUtil.Fields;
 
@@ -25,13 +26,12 @@ public class DocumentRepository extends EntityRepository<Document> {
   static final String DOCUMENT_UPDATE_FIELDS = "data";
   static final String DOCUMENT_PATCH_FIELDS = "data";
 
-  public DocumentRepository(CollectionDAO dao) {
+  public DocumentRepository() {
     super(
         DocStoreResource.COLLECTION_PATH,
         DOCUMENT,
         Document.class,
-        dao.docStoreDAO(),
-        dao,
+        Entity.getCollectionDAO().docStoreDAO(),
         DOCUMENT_UPDATE_FIELDS,
         DOCUMENT_PATCH_FIELDS);
     supportsSearch = false;