diff --git a/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java b/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java
index 93c141915..d170bf543 100644
--- a/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java
+++ b/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java
@@ -23,6 +23,7 @@ import com.google.common.collect.Streams;
 import google.registry.beam.common.RegistryQuery.CriteriaQuerySupplier;
 import google.registry.model.UpdateAutoTimestamp;
 import google.registry.model.UpdateAutoTimestamp.DisableAutoUpdateResource;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.replay.SqlEntity;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -234,6 +235,10 @@ public final class RegistryJpaIO {
 
       @ProcessElement
       public void processElement(OutputReceiver<T> outputReceiver) {
+        // Preload the migration schedule into cache, otherwise the cache loading query may happen
+        // before the setDatabaseSnapshot call in the transaction below, causing it to fail.
+        DatabaseMigrationStateSchedule.get();
+
         jpaTm()
             .transactNoRetry(
                 () -> {
diff --git a/core/src/test/java/google/registry/beam/common/DatabaseSnapshotTest.java b/core/src/test/java/google/registry/beam/common/DatabaseSnapshotTest.java
index d71db5fee..7ad60e2de 100644
--- a/core/src/test/java/google/registry/beam/common/DatabaseSnapshotTest.java
+++ b/core/src/test/java/google/registry/beam/common/DatabaseSnapshotTest.java
@@ -37,6 +37,7 @@ import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junitpioneer.jupiter.RetryingTest;
 import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
@@ -106,7 +107,7 @@ public class DatabaseSnapshotTest {
     try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {}
   }
 
-  @Test
+  @RetryingTest(2) // TODO(b/226945065) revert to @Test when DatbaseMigrationSchedule is removed.
   void readSnapshot() {
     try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {
       Registry snapshotRegistry =
@@ -120,7 +121,7 @@ public class DatabaseSnapshotTest {
     }
   }
 
-  @Test
+  @RetryingTest(2) // TODO(b/226945065) revert to @Test when DatbaseMigrationSchedule is removed.
   void readSnapshot_withSubsequentChange() {
     try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {
       Registry updated =