diff --git a/misk-hibernate-testing/src/main/kotlin/misk/jdbc/VitessScaleSafetyChecks.kt b/misk-hibernate-testing/src/main/kotlin/misk/jdbc/VitessScaleSafetyChecks.kt index 27116438549..bf8c187866f 100644 --- a/misk-hibernate-testing/src/main/kotlin/misk/jdbc/VitessScaleSafetyChecks.kt +++ b/misk-hibernate-testing/src/main/kotlin/misk/jdbc/VitessScaleSafetyChecks.kt @@ -262,7 +262,7 @@ class VitessScaleSafetyChecks( * Digs into the MySQL log to find the last executed DML statement that passed through Vitess. */ private fun extractLastDmlQuery(): String? { - return connect().let { c -> + return connect()?.let { c -> c.createStatement().use { s -> s.executeQuery(""" SELECT argument @@ -288,11 +288,11 @@ class VitessScaleSafetyChecks( * closed. We shut down the Vitess docker container after the tests have completed running so * this doesn't need to be closed explicitly. */ - private fun connect(): Connection { + private fun connect(): Connection? { var connection = connection if (connection != null) return connection - val cluster = startVitessService.cluster() + val cluster = startVitessService.cluster() ?: return null connection = cluster.openMysqlConnection() this.connection = connection return connection diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/SchemaValidatorService.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/SchemaValidatorService.kt index 81fffa84977..f038da80b61 100644 --- a/misk-hibernate/src/main/kotlin/misk/hibernate/SchemaValidatorService.kt +++ b/misk-hibernate/src/main/kotlin/misk/hibernate/SchemaValidatorService.kt @@ -3,8 +3,12 @@ package misk.hibernate import com.google.common.util.concurrent.AbstractIdleService import com.google.inject.Key import misk.DependentService +import misk.backoff.Backoff +import misk.backoff.ExponentialBackoff +import misk.backoff.retry import misk.inject.toKey import misk.jdbc.DataSourceConfig +import java.time.Duration import javax.inject.Provider import javax.inject.Singleton import kotlin.reflect.KClass @@ -27,7 +31,10 @@ internal class SchemaValidatorService internal constructor( synchronized(this) { val validator = SchemaValidator() val sessionFactoryService = sessionFactoryServiceProvider.get() - validator.validate(transacterProvider.get(), sessionFactoryService.hibernateMetadata) + // Sometimes the schema hasn't been refreshed at this point in Vitess. So we retry a few times. + retry(5, ExponentialBackoff(Duration.ofMillis(10), Duration.ofMillis(100))) { + validator.validate(transacterProvider.get(), sessionFactoryService.hibernateMetadata) + } } } diff --git a/misk-hibernate/src/main/kotlin/misk/vitess/StartVitessService.kt b/misk-hibernate/src/main/kotlin/misk/vitess/StartVitessService.kt index b08c8bf3248..298abdd725f 100644 --- a/misk-hibernate/src/main/kotlin/misk/vitess/StartVitessService.kt +++ b/misk-hibernate/src/main/kotlin/misk/vitess/StartVitessService.kt @@ -342,9 +342,9 @@ class StartVitessService( val environment: Environment ) - fun cluster() = cluster!!.cluster + fun cluster() = cluster?.cluster - fun shouldRunVitess() = config.type == DataSourceType.VITESS && (environment == TESTING || environment == DEVELOPMENT) + fun shouldRunVitess() = config.type == DataSourceType.VITESS && (environment == TESTING || environment == DEVELOPMENT) override fun shutDown() { } @@ -414,7 +414,7 @@ class StartVitessService( * MyAppVitessDaemon.kt: * * fun main() { - * val config = MiskConfig.load("myapp") + * val config = MiskConfig.load("myapp", Environment.TESTING) * startVitessDaemon(MyAppDb::class, config.data_source_clusters.values.first().writer) * } * @@ -422,7 +422,8 @@ class StartVitessService( fun startVitessDaemon( /** The same qualifier passed into [HibernateModule], used to uniquely name the container */ qualifier: KClass, - /** Config for the Vitess cluster */ + /** Config for the Vitess clu + * ster */ config: DataSourceConfig ) { val docker: DockerClient = DockerClientBuilder.getInstance() diff --git a/misk-hibernate/src/test/kotlin/misk/hibernate/VitessSchemaMigratorTest.kt b/misk-hibernate/src/test/kotlin/misk/hibernate/VitessSchemaMigratorTest.kt index b3c1d16f5be..d805229382a 100644 --- a/misk-hibernate/src/test/kotlin/misk/hibernate/VitessSchemaMigratorTest.kt +++ b/misk-hibernate/src/test/kotlin/misk/hibernate/VitessSchemaMigratorTest.kt @@ -27,7 +27,7 @@ internal class VitessSchemaMigratorTest { } @AfterEach fun cleanUpMigrationTable() { - openDirectConnection().use { c -> + openDirectConnection()?.use { c -> val schemaVersion = c.prepareStatement(""" |DELETE FROM `vt_movies_-80`.schema_version |""".trimMargin()) @@ -62,7 +62,7 @@ internal class VitessSchemaMigratorTest { // The schema_version is unknown to Vitess which means we can query it with shard targetting // but we can't insert into it (unless we specify -queryserver-config-allowunsafe-dmls which // vttestserver currently does not). So we bypass Vitess to insert into it directly. - openDirectConnection().use { c -> + openDirectConnection()?.use { c -> val schemaVersion = c.prepareStatement(""" |INSERT INTO `vt_movies_-80`.schema_version (version, installed_by) VALUES (?, ?) |""".trimMargin()) @@ -73,8 +73,8 @@ internal class VitessSchemaMigratorTest { } /** Open a direct connection to the Vitess MySQL instance. */ - private fun openDirectConnection(): Connection { + private fun openDirectConnection(): Connection? { val cluster = vitessService.cluster() - return cluster.openMysqlConnection() + return cluster?.openMysqlConnection() } }