Fix Failing Maven CI Test (#12151)

open-metadata · Jun 26, 2023 · 80b273a · 80b273a
1 parent 73f5859
commit 80b273a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 31 deletions.
diff --git a/...adata-service/src/main/java/org/openmetadata/service/migration/api/MigrationWorkflow.java b/...adata-service/src/main/java/org/openmetadata/service/migration/api/MigrationWorkflow.java
@@ -7,44 +7,25 @@
 import org.jdbi.v3.core.Handle;
 import org.jdbi.v3.core.Jdbi;
 import org.openmetadata.service.jdbi3.MigrationDAO;
-import org.openmetadata.service.jdbi3.locator.ConnectionType;
 
 @Slf4j
 public class MigrationWorkflow {
   private final List<MigrationStep> migrations;
   private final MigrationDAO migrationDAO;
   private final Jdbi jdbi;
-  private final ConnectionType workflowDatabaseConnectionType;
-
   private boolean ignoreFileChecksum = false;
 
-  public MigrationWorkflow(
-      Jdbi jdbi, ConnectionType type, List<MigrationStep> migrationSteps, boolean ignoreFileChecksum) {
+  public MigrationWorkflow(Jdbi jdbi, List<MigrationStep> migrationSteps, boolean ignoreFileChecksum) {
     this.jdbi = jdbi;
-    this.workflowDatabaseConnectionType = type;
     this.migrationDAO = jdbi.onDemand(MigrationDAO.class);
     this.ignoreFileChecksum = ignoreFileChecksum;
-
-    // Validate Migration
-    validateMigrations(migrationSteps);
-
     // Sort Migration on the basis of version
     migrationSteps.sort(Comparator.comparing(MigrationStep::getMigrationVersion));
 
     // Filter Migrations to Be Run
     this.migrations = filterAndGetMigrationsToRun(migrationSteps);
   }
 
-  private void validateMigrations(List<MigrationStep> migrationSteps) {
-    for (MigrationStep step : migrationSteps) {
-      if (!step.getDatabaseConnectionType().equals(this.workflowDatabaseConnectionType)) {
-        throw new IllegalArgumentException(
-            String.format(
-                "Provided Migration File is for Database Connection %s", step.getDatabaseConnectionType().toString()));
-      }
-    }
-  }
-
   private List<MigrationStep> filterAndGetMigrationsToRun(List<MigrationStep> migrations) {
     LOG.debug("Filtering Server Migrations");
     String maxMigration = migrations.get(migrations.size() - 1).getMigrationVersion();

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/TablesInitializer.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/TablesInitializer.java
@@ -50,10 +50,10 @@
 import org.openmetadata.service.fernet.Fernet;
 import org.openmetadata.service.jdbi3.CollectionDAO;
 import org.openmetadata.service.jdbi3.locator.ConnectionAwareAnnotationSqlLocator;
+import org.openmetadata.service.jdbi3.locator.ConnectionType;
 import org.openmetadata.service.migration.MigrationFile;
 import org.openmetadata.service.migration.api.MigrationStep;
 import org.openmetadata.service.migration.api.MigrationWorkflow;
-import org.openmetadata.service.resources.databases.DatasourceConfig;
 import org.openmetadata.service.search.IndexUtil;
 import org.openmetadata.service.search.SearchClient;
 import org.openmetadata.service.secrets.SecretsManagerFactory;
@@ -278,7 +278,8 @@ private static void execute(
       case MIGRATE:
         flyway.migrate();
         // Validate and Run System Data Migrations
-        validateAndRunSystemDataMigrations(jdbi, config, ignoreServerFileChecksum);
+        validateAndRunSystemDataMigrations(
+            jdbi, ConnectionType.from(config.getDataSourceFactory().getDriverClass()), ignoreServerFileChecksum);
         break;
       case INFO:
         printToConsoleMandatory(dumpToAsciiTable(flyway.info().all()));
@@ -328,21 +329,18 @@ private static void printToConsoleInDebug(String message) {
     }
   }
 
-  private static void validateAndRunSystemDataMigrations(
-      Jdbi jdbi, OpenMetadataApplicationConfig config, boolean ignoreFileChecksum) {
-    DatasourceConfig.initialize(config);
-    List<MigrationStep> loadedMigrationFiles = getServerMigrationFiles();
-    MigrationWorkflow workflow =
-        new MigrationWorkflow(
-            jdbi, DatasourceConfig.getInstance().getDatabaseConnectionType(), loadedMigrationFiles, ignoreFileChecksum);
+  public static void validateAndRunSystemDataMigrations(
+      Jdbi jdbi, ConnectionType connType, boolean ignoreFileChecksum) {
+    List<MigrationStep> loadedMigrationFiles = getServerMigrationFiles(connType);
+    MigrationWorkflow workflow = new MigrationWorkflow(jdbi, loadedMigrationFiles, ignoreFileChecksum);
     workflow.runMigrationWorkflows();
   }
 
-  private static List<MigrationStep> getServerMigrationFiles() {
+  private static List<MigrationStep> getServerMigrationFiles(ConnectionType connType) {
     List<MigrationStep> migrations = new ArrayList<>();
     try {
       String prefix =
-          Boolean.TRUE.equals(DatasourceConfig.getInstance().isMySQL())
+          connType.equals(ConnectionType.MYSQL)
               ? "org.openmetadata.service.migration.versions.mysql"
               : "org.openmetadata.service.migration.versions.postgres";
       Reflections reflections = new Reflections(prefix);

diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/OpenMetadataApplicationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/OpenMetadataApplicationTest.java
@@ -14,6 +14,7 @@
 package org.openmetadata.service;
 
 import static java.lang.String.format;
+import static org.openmetadata.service.util.TablesInitializer.validateAndRunSystemDataMigrations;
 
 import io.dropwizard.jersey.jackson.JacksonFeature;
 import io.dropwizard.testing.ConfigOverride;
@@ -30,11 +31,16 @@
 import org.glassfish.jersey.client.ClientProperties;
 import org.glassfish.jersey.client.HttpUrlConnectorProvider;
 import org.glassfish.jersey.client.JerseyClientBuilder;
+import org.jdbi.v3.core.Jdbi;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+import org.jdbi.v3.sqlobject.SqlObjects;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.TestInstance;
 import org.openmetadata.common.utils.CommonUtil;
 import org.openmetadata.service.fernet.Fernet;
+import org.openmetadata.service.jdbi3.locator.ConnectionAwareAnnotationSqlLocator;
+import org.openmetadata.service.jdbi3.locator.ConnectionType;
 import org.openmetadata.service.resources.CollectionRegistry;
 import org.openmetadata.service.resources.events.WebhookCallbackResource;
 import org.openmetadata.service.resources.tags.TagLabelCache;
@@ -140,6 +146,13 @@ public static void createApplication() throws Exception {
     ConfigOverride[] configOverridesArray = configOverrides.toArray(new ConfigOverride[configOverrides.size()]);
     APP = new DropwizardAppExtension<>(OpenMetadataApplication.class, CONFIG_PATH, configOverridesArray);
 
+    // Run System Migrations
+    final Jdbi jdbi = Jdbi.create(sqlContainer.getJdbcUrl(), sqlContainer.getUsername(), sqlContainer.getPassword());
+    jdbi.installPlugin(new SqlObjectPlugin());
+    jdbi.getConfig(SqlObjects.class)
+        .setSqlLocator(new ConnectionAwareAnnotationSqlLocator(sqlContainer.getDriverClassName()));
+    validateAndRunSystemDataMigrations(jdbi, ConnectionType.from(sqlContainer.getDriverClassName()), false);
+
     APP.before();
   }