Fix sporadic SQL Snapshot failure (#1571)

* Fix sporadic SQL Snapshot failure

The Postgresql set-snapshot statement (called in
JpaTransactionManager.setDatabaseSnapshot() method) must be the first
statement in the SQL transaction.

Currenty the JpaTransaction.transact() method may insert a query for
DatabaseMigrationStateSchedule before the user query when the cache is
empty or the cached value expires.

This PR proactively preloads the cache in RegistryJpaIO to prevent cache
loading  inside the transaction.

This PR also changes some DatabaseSnapshotTest tests to be retrying, in
case they run just after the cache expires. (This has happened before in
CI).

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(
                 () -> {