VitessTestDb: harden schema apply logic

GitOrigin-RevId: 85fe78af2e8a20c3ce940baa7e2539afa1edaf28

Author: aparajon

misk-vitess/src/testFixtures/kotlin/misk/vitess/testing/internal/VitessQueryExecutor.kt

@@ -5,6 +5,8 @@ import misk.vitess.testing.VitessTableType
 import misk.vitess.testing.VitessTestDbException
 import java.sql.Connection
 import java.sql.DriverManager
+import java.sql.SQLRecoverableException
+import java.sql.SQLTransientException
 import java.sql.Statement
 
 /**
@@ -36,6 +38,49 @@ internal class VitessQueryExecutor(
     return connection.use { conn -> executeUpdate(conn, query, target) }
   }
 
+  /**
+   * Execute update with retry logic to handle retryable failures.
+   */
+  fun executeUpdateWithRetries(query: String, target: String = "@primary"): Int {
+    var lastException: Exception? = null
+    val maxRetries = 3
+    val retryDelayMs = 500L
+
+    val retryableErrorMessages = listOf(
+      "Keyspace does not have exactly one shard",
+    )
+    
+    for (attempt in 1..maxRetries) {
+      try {
+        return executeUpdate(query, target)
+      } catch (e: Exception) {
+        lastException = e
+
+        // Check if this is a retryable exception type or message
+        val isRetryableException = e.cause is SQLRecoverableException || 
+                                   e.cause is SQLTransientException ||
+                                   e is SQLRecoverableException ||
+                                   e is SQLTransientException
+        
+        val isRetryableMessage = retryableErrorMessages.any { errorMessage ->
+          e.message?.contains(errorMessage, ignoreCase = true) == true
+        }
+        
+        if ((isRetryableException || isRetryableMessage) && attempt < maxRetries) {
+          Thread.sleep(retryDelayMs * attempt)
+        } else {
+          throw e // Re-throw if not a retry-able error or max retries reached
+        }
+      }
+    }
+    
+    // If we get here, all retries failed
+    throw VitessQueryExecutorException(
+      "Failed to executeUpdate after `$maxRetries` attempts. Query: `$query`, Last error: `${lastException?.message}`",
+      lastException
+    )
+  }
+
   fun executeTransaction(query: String, target: String = "@primary"): Boolean {
     val connection = getVtgateConnection()
     return connection.use { conn ->

misk-vitess/src/testFixtures/kotlin/misk/vitess/testing/internal/VitessSchemaApplier.kt

@@ -136,15 +136,28 @@ internal class VitessSchemaApplier(
     val ddlUpdates = ConcurrentLinkedQueue<DdlUpdate>()
     val vschemaUpdates = ConcurrentLinkedQueue<VSchemaUpdate>()
 
-    // Run all keyspace operations in parallel to speed up schema change processing.
-    val executorService = Executors.newFixedThreadPool(keyspaces.size)
+    // Apply VSchemas in parallel but wait for all to complete before proceeding to DDL's.
+    val vschemaExecutorService = Executors.newFixedThreadPool(keyspaces.size)
+    val vschemaFutures = mutableListOf<Future<VSchemaUpdate>>()
+    
+    keyspaces.forEach { keyspace ->
+      vschemaFutures.add(
+        vschemaExecutorService.submit<VSchemaUpdate> {
+          processVschema(keyspace)
+        }
+      )
+    }
+
+    vschemaFutures.forEach { future ->
+      vschemaUpdates.add(future.get())
+    }
+    vschemaExecutorService.shutdown()
+
+    val ddlExecutorService = Executors.newFixedThreadPool(keyspaces.size)
     val futures = mutableListOf<Future<Unit>>()
     keyspaces.forEach { keyspace ->
       futures.add(
-        executorService.submit<Unit> {
-          val vschemaUpdate = processVschema(keyspace)
-          vschemaUpdates.add(vschemaUpdate)
-
+        ddlExecutorService.submit<Unit> {
           var schemaChangesProcessedForKeyspace: Boolean
           if (enableDeclarativeSchemaChanges) {
             schemaChangesProcessedForKeyspace = applyDeclarativeSchemaChanges(keyspace, vitessQueryExecutor, ddlUpdates)
@@ -161,7 +174,7 @@ internal class VitessSchemaApplier(
     }
 
     futures.forEach { it.get() }
-    executorService.shutdown()
+    ddlExecutorService.shutdown()
 
     initializeSequenceTables(vitessQueryExecutor)
     println("🔧 Schema is applied.")
@@ -236,9 +249,9 @@ internal class VitessSchemaApplier(
     keyspace.ddlCommands
       .sortedBy { it.first }
       .map { it.second }
-      .forEach {
-        vitessQueryExecutor.executeUpdate(query = it, target = keyspace.name)
-        ddlUpdates.add(DdlUpdate(ddl = it, keyspace = keyspace.name))
+      .forEach { ddl ->
+        vitessQueryExecutor.executeUpdateWithRetries(ddl, keyspace.name)
+        ddlUpdates.add(DdlUpdate(ddl = ddl, keyspace = keyspace.name))
       }
 
     return true
@@ -286,7 +299,7 @@ internal class VitessSchemaApplier(
 
   private fun applyDdlCommands(ddl: String, keyspace: VitessKeyspace, vitessQueryExecutor: VitessQueryExecutor) {
     printDebug("Applying schema changes:\n$ddl")
-    vitessQueryExecutor.executeUpdate(ddl, keyspace.name)
+    vitessQueryExecutor.executeUpdateWithRetries(ddl, keyspace.name)
   }
 
   private fun executeDockerCommand(command: List<String>, keyspace: String): String {