Add some robustness to local development with Vitess

cashapp · Apr 11, 2019 · 79ad8cf · 79ad8cf
1 parent ee7501d
commit 79ad8cf
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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,15 +414,16 @@ class StartVitessService(
      * MyAppVitessDaemon.kt:
      *
      *  fun main() {
-     *    val config = MiskConfig.load<MyAppConfig>("myapp")
+     *    val config = MiskConfig.load<MyAppConfig>("myapp", Environment.TESTING)
      *    startVitessDaemon(MyAppDb::class, config.data_source_clusters.values.first().writer)
      *  }
      *
      */
     fun startVitessDaemon(
       /** The same qualifier passed into [HibernateModule], used to uniquely name the container */
       qualifier: KClass<out Annotation>,
-      /** 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
@@ -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()
   }
 }