Fix migration dev center schema dump by run db-specific initializatio…

…n script (#18984)

* Run db-specific initialization script

* Refactor code

* Use constants

* Add unit test for dev center

airbyte-bootloader/src/main/java/io/airbyte/bootloader/BootloaderApp.java

@@ -148,10 +148,10 @@ public BootloaderApp(final Configs configs,
   public void load() throws Exception {
     LOGGER.info("Initializing databases...");
     DatabaseCheckFactory.createConfigsDatabaseInitializer(configsDslContext,
-        configs.getConfigsDatabaseInitializationTimeoutMs(), MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH)).initialize();
+        configs.getConfigsDatabaseInitializationTimeoutMs(), MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH)).initialize();
 
     DatabaseCheckFactory.createJobsDatabaseInitializer(jobsDslContext,
-        configs.getJobsDatabaseInitializationTimeoutMs(), MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH)).initialize();
+        configs.getJobsDatabaseInitializationTimeoutMs(), MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH)).initialize();
     LOGGER.info("Databases initialized.");
 
     LOGGER.info("Setting up config database and default workspace...");

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/DatabaseConstants.java

@@ -26,7 +26,9 @@ public final class DatabaseConstants {
   /**
    * Path to the script that contains the initial schema definition for the Configurations database.
    */
-  public static final String CONFIGS_SCHEMA_PATH = "configs_database/schema.sql";
+  public static final String CONFIGS_INITIAL_SCHEMA_PATH = "configs_database/schema.sql";
+
+  public static final String CONFIGS_SCHEMA_DUMP_PATH = "src/main/resources/configs_database/schema_dump.txt";
 
   /**
    * Logical name of the Jobs database.
@@ -41,7 +43,9 @@ public final class DatabaseConstants {
   /**
    * Path to the script that contains the initial schema definition for the Jobs database.
    */
-  public static final String JOBS_SCHEMA_PATH = "jobs_database/schema.sql";
+  public static final String JOBS_INITIAL_SCHEMA_PATH = "jobs_database/schema.sql";
+
+  public static final String JOBS_SCHEMA_DUMP_PATH = "src/main/resources/jobs_database/schema_dump.txt";
 
   /**
    * Default database connection timeout in milliseconds.

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseMigrationDevCenter.java

@@ -6,32 +6,26 @@
 
 import io.airbyte.db.Database;
 import io.airbyte.db.factory.FlywayFactory;
+import io.airbyte.db.instance.DatabaseConstants;
 import io.airbyte.db.instance.FlywayDatabaseMigrator;
 import io.airbyte.db.instance.development.MigrationDevCenter;
-import java.io.IOException;
 import javax.sql.DataSource;
 import org.flywaydb.core.Flyway;
-import org.jooq.DSLContext;
 
 /**
  * Helper class for migration development. See README for details.
  */
 public class ConfigsDatabaseMigrationDevCenter extends MigrationDevCenter {
 
   public ConfigsDatabaseMigrationDevCenter() {
-    super("configs", "src/main/resources/configs_database/schema_dump.txt");
+    super("configs", DatabaseConstants.CONFIGS_SCHEMA_DUMP_PATH, DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
   }
 
   @Override
   protected FlywayDatabaseMigrator getMigrator(final Database database, final Flyway flyway) {
     return new ConfigsDatabaseMigrator(database, flyway);
   }
 
-  @Override
-  protected Database getDatabase(final DSLContext dslContext) throws IOException {
-    return new Database(dslContext);
-  }
-
   @Override
   protected Flyway getFlyway(final DataSource dataSource) {
     return FlywayFactory.create(dataSource, getClass().getSimpleName(), ConfigsDatabaseMigrator.DB_IDENTIFIER,

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java

@@ -35,7 +35,7 @@ public ConfigsDatabaseTestProvider(final DSLContext dslContext, final Flyway fly
 
   @Override
   public Database create(final boolean runMigration) throws IOException, DatabaseInitializationException {
-    final String initalSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH);
+    final String initalSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
     DatabaseCheckFactory.createConfigsDatabaseInitializer(dslContext, DatabaseConstants.DEFAULT_CONNECTION_TIMEOUT_MS, initalSchema).initialize();
 
     final Database database = new Database(dslContext);

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/configs/ConfigsFlywayMigrationDatabase.java

@@ -47,7 +47,7 @@ protected String[] getMigrationFileLocations() {
 
   @Override
   protected void initializeDatabase(final DSLContext dslContext) throws DatabaseInitializationException, IOException {
-    final String initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH);
+    final String initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
     DatabaseCheckFactory.createConfigsDatabaseInitializer(dslContext, DatabaseConstants.DEFAULT_CONNECTION_TIMEOUT_MS, initialSchema).initialize();
   }
 

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/development/MigrationDevCenter.java

@@ -4,6 +4,7 @@
 
 package io.airbyte.db.instance.development;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.db.Database;
 import io.airbyte.db.factory.DSLContextFactory;
 import io.airbyte.db.factory.DataSourceFactory;
@@ -37,30 +38,33 @@ private enum Command {
 
   private final String dbIdentifier;
   private final String schemaDumpFile;
+  private final String initialScript;
 
-  protected MigrationDevCenter(final String dbIdentifier, final String schemaDumpFile) {
+  protected MigrationDevCenter(final String dbIdentifier, final String schemaDumpFile, final String initialScript) {
     this.dbIdentifier = dbIdentifier;
     this.schemaDumpFile = schemaDumpFile;
+    this.initialScript = initialScript;
   }
 
-  private static PostgreSQLContainer<?> createContainer() {
+  private PostgreSQLContainer<?> createContainer() {
     final PostgreSQLContainer<?> container = new PostgreSQLContainer<>("postgres:13-alpine")
         .withDatabaseName("airbyte")
         .withUsername("docker")
         .withPassword("docker");
     container.start();
-
-    initializeDatabase(container);
-
+    final var containerDelegate = new JdbcDatabaseDelegate(container, "");
+    ScriptUtils.runInitScript(containerDelegate, initialScript);
     return container;
   }
 
   protected abstract FlywayDatabaseMigrator getMigrator(Database database, Flyway flyway);
 
-  protected abstract Database getDatabase(DSLContext dslContext) throws IOException;
-
   protected abstract Flyway getFlyway(DataSource dataSource);
 
+  private Database getDatabase(final DSLContext dslContext) throws IOException {
+    return new Database(dslContext);
+  }
+
   private void createMigration() {
     try (final PostgreSQLContainer<?> container = createContainer()) {
       final DataSource dataSource =
@@ -94,7 +98,8 @@ private void runLastMigration() {
     }
   }
 
-  private void dumpSchema() {
+  @VisibleForTesting
+  public String dumpSchema(final boolean persistToFile) {
     try (final PostgreSQLContainer<?> container = createContainer()) {
       final DataSource dataSource =
           DataSourceFactory.create(container.getUsername(), container.getPassword(), container.getDriverClassName(), container.getJdbcUrl());
@@ -104,19 +109,16 @@ private void dumpSchema() {
         final FlywayDatabaseMigrator migrator = getMigrator(database, flyway);
         migrator.migrate();
         final String schema = migrator.dumpSchema();
-        MigrationDevHelper.dumpSchema(schema, schemaDumpFile, true);
+        if (persistToFile) {
+          MigrationDevHelper.dumpSchema(schema, schemaDumpFile, true);
+        }
+        return schema;
       }
     } catch (final Exception e) {
       throw new RuntimeException(e);
     }
   }
 
-  private static void initializeDatabase(final PostgreSQLContainer container) {
-    final var containerDelegate = new JdbcDatabaseDelegate(container, "");
-    ScriptUtils.runInitScript(containerDelegate, "configs_database/schema.sql");
-    ScriptUtils.runInitScript(containerDelegate, "jobs_database/schema.sql");
-  }
-
   public static void main(final String[] args) {
     final MigrationDevCenter devCenter;
 
@@ -131,7 +133,7 @@ public static void main(final String[] args) {
     switch (command) {
       case CREATE -> devCenter.createMigration();
       case MIGRATE -> devCenter.runLastMigration();
-      case DUMP_SCHEMA -> devCenter.dumpSchema();
+      case DUMP_SCHEMA -> devCenter.dumpSchema(true);
       default -> throw new IllegalArgumentException("Unexpected command: " + args[1]);
     }
   }

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseMigrationDevCenter.java

@@ -6,32 +6,26 @@
 
 import io.airbyte.db.Database;
 import io.airbyte.db.factory.FlywayFactory;
+import io.airbyte.db.instance.DatabaseConstants;
 import io.airbyte.db.instance.FlywayDatabaseMigrator;
 import io.airbyte.db.instance.development.MigrationDevCenter;
-import java.io.IOException;
 import javax.sql.DataSource;
 import org.flywaydb.core.Flyway;
-import org.jooq.DSLContext;
 
 /**
  * Helper class for migration development. See README for details.
  */
 public class JobsDatabaseMigrationDevCenter extends MigrationDevCenter {
 
   public JobsDatabaseMigrationDevCenter() {
-    super("jobs", "src/main/resources/jobs_database/schema_dump.txt");
+    super("jobs", DatabaseConstants.JOBS_SCHEMA_DUMP_PATH, DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
   }
 
   @Override
   protected FlywayDatabaseMigrator getMigrator(final Database database, final Flyway flyway) {
     return new JobsDatabaseMigrator(database, flyway);
   }
 
-  @Override
-  protected Database getDatabase(final DSLContext dslContext) throws IOException {
-    return new Database(dslContext);
-  }
-
   @Override
   protected Flyway getFlyway(final DataSource dataSource) {
     return FlywayFactory.create(dataSource, getClass().getSimpleName(), JobsDatabaseMigrator.DB_IDENTIFIER,

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseTestProvider.java

@@ -27,7 +27,7 @@ public JobsDatabaseTestProvider(final DSLContext dslContext, final Flyway flyway
 
   @Override
   public Database create(final boolean runMigration) throws IOException, DatabaseInitializationException {
-    final String initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final String initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     DatabaseCheckFactory.createJobsDatabaseInitializer(dslContext, DatabaseConstants.DEFAULT_CONNECTION_TIMEOUT_MS, initialSchema).initialize();
 
     final Database jobsDatabase = new Database(dslContext);

airbyte-db/db-lib/src/main/java/io/airbyte/db/instance/jobs/JobsFlywayMigrationDatabase.java

@@ -47,7 +47,7 @@ protected String[] getMigrationFileLocations() {
 
   @Override
   protected void initializeDatabase(final DSLContext dslContext) throws DatabaseInitializationException, IOException {
-    final String initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final String initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     DatabaseCheckFactory.createJobsDatabaseInitializer(dslContext, DatabaseConstants.DEFAULT_CONNECTION_TIMEOUT_MS, initialSchema).initialize();
   }
 

airbyte-db/db-lib/src/test/java/io/airbyte/db/init/impl/ConfigsDatabaseInitializerTest.java

@@ -25,7 +25,7 @@ class ConfigsDatabaseInitializerTest extends CommonDatabaseInitializerTest {
   @Test
   void testInitializingSchema() throws IOException {
     final var databaseAvailabilityCheck = mock(ConfigsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
     final var initializer = new ConfigsDatabaseInitializer(databaseAvailabilityCheck, dslContext, initialSchema);
 
     Assertions.assertDoesNotThrow(() -> initializer.initialize());
@@ -35,7 +35,7 @@ void testInitializingSchema() throws IOException {
   @Test
   void testInitializingSchemaAlreadyExists() throws IOException {
     final var databaseAvailabilityCheck = mock(ConfigsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
     dslContext.execute(initialSchema);
     final var initializer = new ConfigsDatabaseInitializer(databaseAvailabilityCheck, dslContext, initialSchema);
 
@@ -46,7 +46,7 @@ void testInitializingSchemaAlreadyExists() throws IOException {
   @Test
   void testInitializationException() throws IOException, DatabaseCheckException {
     final var databaseAvailabilityCheck = mock(ConfigsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.CONFIGS_INITIAL_SCHEMA_PATH);
 
     doThrow(new DatabaseCheckException("test")).when(databaseAvailabilityCheck).check();
 

airbyte-db/db-lib/src/test/java/io/airbyte/db/init/impl/JobsDatabaseInitializerTest.java

@@ -25,7 +25,7 @@ class JobsDatabaseInitializerTest extends CommonDatabaseInitializerTest {
   @Test
   void testInitializingSchema() throws IOException {
     final var databaseAvailabilityCheck = mock(JobsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     final var initializer = new JobsDatabaseInitializer(databaseAvailabilityCheck, dslContext, initialSchema);
 
     Assertions.assertDoesNotThrow(() -> initializer.initialize());
@@ -35,7 +35,7 @@ void testInitializingSchema() throws IOException {
   @Test
   void testInitializingSchemaAlreadyExists() throws IOException {
     final var databaseAvailabilityCheck = mock(JobsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     dslContext.execute(initialSchema);
     final var initializer = new JobsDatabaseInitializer(databaseAvailabilityCheck, dslContext, initialSchema);
 
@@ -46,7 +46,7 @@ void testInitializingSchemaAlreadyExists() throws IOException {
   @Test
   void testInitializationException() throws IOException, DatabaseCheckException {
     final var databaseAvailabilityCheck = mock(JobsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
 
     doThrow(new DatabaseCheckException("test")).when(databaseAvailabilityCheck).check();
 
@@ -56,15 +56,15 @@ void testInitializationException() throws IOException, DatabaseCheckException {
 
   @Test
   void testInitializationNullAvailabilityCheck() throws IOException {
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     final var initializer = new JobsDatabaseInitializer(null, dslContext, initialSchema);
     Assertions.assertThrows(DatabaseInitializationException.class, () -> initializer.initialize());
   }
 
   @Test
   void testInitializationNullDslContext() throws IOException {
     final var databaseAvailabilityCheck = mock(JobsDatabaseAvailabilityCheck.class);
-    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_SCHEMA_PATH);
+    final var initialSchema = MoreResources.readResource(DatabaseConstants.JOBS_INITIAL_SCHEMA_PATH);
     final var initializer = new JobsDatabaseInitializer(databaseAvailabilityCheck, null, initialSchema);
     Assertions.assertThrows(DatabaseInitializationException.class, () -> initializer.initialize());
   }

airbyte-db/db-lib/src/test/java/io/airbyte/db/instance/configs/ConfigsDatabaseMigrationDevCenterTest.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.instance.configs;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import io.airbyte.commons.io.IOs;
+import io.airbyte.db.instance.DatabaseConstants;
+import io.airbyte.db.instance.development.MigrationDevCenter;
+import java.nio.file.Path;
+import org.junit.jupiter.api.Test;
+
+class ConfigsDatabaseMigrationDevCenterTest {
+
+  /**
+   * This test ensures that the dev center is working correctly end-to-end. If it fails, it means
+   * either the migration is not run properly, or the database initialization is incorrect in the dev
+   * center implementation.
+   */
+  @Test
+  void testSchemaDump() {
+    final MigrationDevCenter devCenter = new ConfigsDatabaseMigrationDevCenter();
+    final String schemaDump = IOs.readFile(Path.of(DatabaseConstants.CONFIGS_SCHEMA_DUMP_PATH));
+    assertEquals(schemaDump.trim(), devCenter.dumpSchema(false));
+  }
+
+}

airbyte-db/db-lib/src/test/java/io/airbyte/db/instance/configs/ConfigsDatabaseMigratorTest.java

@@ -5,6 +5,7 @@
 package io.airbyte.db.instance.configs;
 
 import io.airbyte.db.factory.FlywayFactory;
+import io.airbyte.db.instance.DatabaseConstants;
 import io.airbyte.db.instance.DatabaseMigrator;
 import io.airbyte.db.instance.development.MigrationDevHelper;
 import java.io.IOException;
@@ -14,16 +15,14 @@
 @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert")
 class ConfigsDatabaseMigratorTest extends AbstractConfigsDatabaseTest {
 
-  private static final String SCHEMA_DUMP_FILE = "src/main/resources/configs_database/schema_dump.txt";
-
   @Test
   void dumpSchema() throws IOException {
     final Flyway flyway = FlywayFactory.create(getDataSource(), getClass().getSimpleName(), ConfigsDatabaseMigrator.DB_IDENTIFIER,
         ConfigsDatabaseMigrator.MIGRATION_FILE_LOCATION);
     final DatabaseMigrator migrator = new ConfigsDatabaseMigrator(database, flyway);
     migrator.migrate();
     final String schema = migrator.dumpSchema();
-    MigrationDevHelper.dumpSchema(schema, SCHEMA_DUMP_FILE, false);
+    MigrationDevHelper.dumpSchema(schema, DatabaseConstants.CONFIGS_SCHEMA_DUMP_PATH, false);
   }
 
 }

airbyte-db/db-lib/src/test/java/io/airbyte/db/instance/jobs/JobsDatabaseMigrationDevCenterTest.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.instance.jobs;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.airbyte.commons.io.IOs;
+import io.airbyte.db.instance.DatabaseConstants;
+import io.airbyte.db.instance.development.MigrationDevCenter;
+import java.nio.file.Path;
+import org.junit.jupiter.api.Test;
+
+class JobsDatabaseMigrationDevCenterTest {
+
+  /**
+   * This test ensures that the dev center is working correctly end-to-end. If it fails, it means
+   * either the migration is not run properly, or the database initialization is incorrect.
+   */
+  @Test
+  void testSchemaDump() {
+    final MigrationDevCenter devCenter = new JobsDatabaseMigrationDevCenter();
+    final String schemaDump = IOs.readFile(Path.of(DatabaseConstants.JOBS_SCHEMA_DUMP_PATH));
+    assertEquals(schemaDump.trim(), devCenter.dumpSchema(false));
+  }
+
+}