Refactor database initialization logic

airbyte-db/lib/src/main/java/io/airbyte/db/check/DatabaseAvailabilityCheck.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check;
+
+import static org.jooq.impl.DSL.select;
+
+import io.airbyte.commons.lang.Exceptions;
+import io.airbyte.db.Database;
+import java.util.Optional;
+import java.util.function.Function;
+import org.jooq.DSLContext;
+import org.slf4j.Logger;
+
+/**
+ * Performs a check to verify that the configured database is available.
+ */
+public interface DatabaseAvailabilityCheck extends DatabaseCheck {
+
+  /**
+   * The number of times to check if the database is available. TODO replace with a default value in a
+   * value injection annotation
+   */
+  int NUM_POLL_TIMES = 10;
+
+  /**
+   * Checks whether the configured database is available.
+   *
+   * @throws InterruptedException if unable to perform the check.
+   */
+  default void check() throws InterruptedException {
+    var initialized = false;
+    var totalTime = 0;
+    final var sleepTime = getTimeoutMs() / NUM_POLL_TIMES;
+
+    while (!initialized) {
+      getLogger().warn("Waiting for database to become available...");
+      if (totalTime >= getTimeoutMs()) {
+        throw new InterruptedException("Unable to connect to the database.");
+      }
+
+      final Optional<DSLContext> dslContext = getDslContext();
+
+      if(dslContext.isPresent()) {
+        final Database database = new Database(dslContext.get());
+        initialized = isDatabaseConnected(getDatabaseName()).apply(database);
+        if (!initialized) {
+          getLogger().info("Database is not ready yet. Please wait a moment, it might still be initializing...");
+          Exceptions.toRuntime(() -> Thread.sleep(sleepTime));
+          totalTime += sleepTime;
+        }
+      } else {
+        throw new InterruptedException("Database configuration not present.");
+      }
+    }
+  }
+
+  /**
+   * Generates a {@link Function} that is used to test if a connection can be made to the database by
+   * verifying that the {@code information_schema.tables} tables has been populated.
+   *
+   * @param databaseName The name of the database to test.
+   * @return A {@link Function} that can be invoked to test if the database is available.
+   */
+  default Function<Database, Boolean> isDatabaseConnected(final String databaseName) {
+    return database -> {
+      try {
+        getLogger().info("Testing {} database connection...", databaseName);
+        return database.query(ctx -> ctx.fetchExists(select().from("information_schema.tables")));
+      } catch (final Exception e) {
+        getLogger().error("Failed to verify database connection.", e);
+        return false;
+      }
+    };
+  }
+
+  /**
+   * Retrieves the configured database name to be tested.
+   *
+   * @return The name of the database to test.
+   */
+  String getDatabaseName();
+
+  /**
+   * Retrieves the configured {@link DSLContext} to be used to test the database availability.
+   *
+   * @return The configured {@link DSLContext} object.
+   */
+  Optional<DSLContext> getDslContext();
+
+  /**
+   * Retrieves the configured {@link Logger} object to be used to record progress of the migration
+   * check.
+   *
+   * @return The configured {@link Logger} object.
+   */
+  Logger getLogger();
+
+  /**
+   * Retrieves the timeout in milliseconds for the check. Once this timeout is exceeded, the check
+   * will fail with an {@link InterruptedException}.
+   *
+   * @return The timeout in milliseconds for the check.
+   */
+  long getTimeoutMs();
+
+}

airbyte-db/lib/src/main/java/io/airbyte/db/check/DatabaseCheck.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check;
+
+/**
+ * Defines the interface for performing checks against a database.
+ */
+public interface DatabaseCheck {
+
+  /**
+   * Checks whether the configured database is available.
+   *
+   * @throws InterruptedException if unable to perform the check.
+   */
+  void check() throws InterruptedException;
+
+}

airbyte-db/lib/src/main/java/io/airbyte/db/check/DatabaseMigrationCheck.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check;
+
+import java.util.Optional;
+import org.flywaydb.core.Flyway;
+import org.slf4j.Logger;
+
+/**
+ * Performs a check to verify that the configured database has been migrated to the appropriate
+ * version.
+ */
+public interface DatabaseMigrationCheck {
+
+  /**
+   * The number of times to check if the database has been migrated to the required schema version.
+   * TODO replace with a default value in a value injection annotation
+   */
+  int NUM_POLL_TIMES = 10;
+
+  /**
+   * Checks whether the configured database has been migrated to the required minimum schema version.
+   *
+   * @throws InterruptedException if unable to perform the check.
+   */
+  default void check() throws InterruptedException {
+    final var startTime = System.currentTimeMillis();
+    final var sleepTime = getTimeoutMs() / NUM_POLL_TIMES;
+    final Optional<Flyway> flywayOptional = getFlyway();
+
+    if(flywayOptional.isPresent()) {
+      final var flyway = flywayOptional.get();
+      var currDatabaseMigrationVersion = flyway.info().current().getVersion().getVersion();
+      getLogger().info("Current database migration version {}.", currDatabaseMigrationVersion);
+      getLogger().info("Minimum Flyway version required {}.", getMinimumFlywayVersion());
+
+      while (currDatabaseMigrationVersion.compareTo(getMinimumFlywayVersion()) < 0) {
+        if (System.currentTimeMillis() - startTime >= getTimeoutMs()) {
+          throw new InterruptedException("Timeout while waiting for database to fulfill minimum flyway migration version..");
+        }
+
+        Thread.sleep(sleepTime);
+        currDatabaseMigrationVersion = flyway.info().current().getVersion().getVersion();
+      }
+      getLogger().info("Verified that database has been migrated to the required minimum version {}.", getTimeoutMs());
+    } else {
+      throw new InterruptedException("Flyway configuration not present.");
+    }
+  }
+
+  /**
+   * Retrieves the configured {@link Flyway} object to be used to check the migration status of the
+   * database.
+   *
+   * @return The configured {@link Flyway} object.
+   */
+  Optional<Flyway> getFlyway();
+
+  /**
+   * Retrieves the configured {@link Logger} object to be used to record progress of the migration
+   * check.
+   *
+   * @return The configured {@link Logger} object.
+   */
+  Logger getLogger();
+
+  /**
+   * Retrieves the required minimum migration version of the schema.
+   *
+   * @return The required minimum migration version of the schema.
+   */
+  String getMinimumFlywayVersion();
+
+  /**
+   * Retrieves the timeout in milliseconds for the check. Once this timeout is exceeded, the check
+   * will fail with an {@link InterruptedException}.
+   *
+   * @return The timeout in milliseconds for the check.
+   */
+  long getTimeoutMs();
+
+}

airbyte-db/lib/src/main/java/io/airbyte/db/check/impl/ConfigsDatabaseAvailabilityCheck.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check.impl;
+
+import io.airbyte.db.check.DatabaseAvailabilityCheck;
+import io.airbyte.db.instance.configs.ConfigsDatabaseInstance;
+import java.util.Optional;
+import org.jooq.DSLContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Implementation of the {@link DatabaseAvailabilityCheck} for the Configurations database.
+ */
+public class ConfigsDatabaseAvailabilityCheck implements DatabaseAvailabilityCheck {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(ConfigsDatabaseAvailabilityCheck.class);
+
+  // TODO inject via dependency injection framework
+  private final DSLContext dslContext;
+
+  // TODO inject via dependency injection framework
+  private final long timeoutMs;
+
+  public ConfigsDatabaseAvailabilityCheck(final DSLContext dslContext, final long timeoutMs) {
+    this.dslContext = dslContext;
+    this.timeoutMs = timeoutMs;
+  }
+
+  @Override
+  public String getDatabaseName() {
+    return ConfigsDatabaseInstance.DATABASE_LOGGING_NAME;
+  }
+
+  @Override
+  public Optional<DSLContext> getDslContext() {
+    return Optional.ofNullable(dslContext);
+  }
+
+  @Override
+  public Logger getLogger() {
+    return LOGGER;
+  }
+
+  @Override
+  public long getTimeoutMs() {
+    return timeoutMs;
+  }
+
+}

airbyte-db/lib/src/main/java/io/airbyte/db/check/impl/ConfigsDatabaseMigrationCheck.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check.impl;
+
+import io.airbyte.db.check.DatabaseMigrationCheck;
+import java.util.Optional;
+import org.flywaydb.core.Flyway;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Implementation of the {@link DatabaseMigrationCheck} for the Configurations database.
+ */
+public class ConfigsDatabaseMigrationCheck implements DatabaseMigrationCheck {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(ConfigsDatabaseMigrationCheck.class);
+
+  // TODO inject via dependency injection framework
+  private final Flyway flyway;
+
+  // TODO inject via dependency injection framework
+  private final String minimumFlywayVersion;
+
+  // TODO inject via dependency injection framework
+  private final long timeoutMs;
+
+  public ConfigsDatabaseMigrationCheck(final Flyway flyway, final String minimumFlywayVersion, final long timeoutMs) {
+    this.flyway = flyway;
+    this.minimumFlywayVersion = minimumFlywayVersion;
+    this.timeoutMs = timeoutMs;
+  }
+
+  @Override
+  public Optional<Flyway> getFlyway() {
+    return Optional.ofNullable(flyway);
+  }
+
+  @Override
+  public Logger getLogger() {
+    return LOGGER;
+  }
+
+  @Override
+  public String getMinimumFlywayVersion() {
+    return minimumFlywayVersion;
+  }
+
+  @Override
+  public long getTimeoutMs() {
+    return timeoutMs;
+  }
+
+}

airbyte-db/lib/src/main/java/io/airbyte/db/check/impl/JobsDatabaseAvailabilityCheck.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.db.check.impl;
+
+import io.airbyte.db.check.DatabaseAvailabilityCheck;
+import java.util.Optional;
+import org.jooq.DSLContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Implementation of the {@link DatabaseAvailabilityCheck} for the Jobs database.
+ */
+public class JobsDatabaseAvailabilityCheck implements DatabaseAvailabilityCheck {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(JobsDatabaseAvailabilityCheck.class);
+
+  private static final String DATABASE_NAME = "airbyte jobs";
+
+  // TODO inject via dependency injection framework
+  private final DSLContext dslContext;
+
+  // TODO inject via dependency injection framework
+  private final long timeoutMs;
+
+  public JobsDatabaseAvailabilityCheck(final DSLContext dslContext, final long timeoutMs) {
+    this.dslContext = dslContext;
+    this.timeoutMs = timeoutMs;
+  }
+
+  @Override
+  public String getDatabaseName() {
+    return DATABASE_NAME;
+  }
+
+  @Override
+  public Optional<DSLContext> getDslContext() {
+    return Optional.ofNullable(dslContext);
+  }
+
+  @Override
+  public Logger getLogger() {
+    return LOGGER;
+  }
+
+  @Override
+  public long getTimeoutMs() {
+    return timeoutMs;
+  }
+
+}