From 71352191513f4fee083009ea13cd0685939e8458 Mon Sep 17 00:00:00 2001
From: Weimin Yu <weiminyu@google.com>
Date: Fri, 25 Mar 2022 15:55:00 -0400
Subject: [PATCH] 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).
---
 .../main/java/google/registry/beam/common/RegistryJpaIO.java | 5 +++++
 .../google/registry/beam/common/DatabaseSnapshotTest.java    | 5 +++--
 2 files changed, 8 insertions(+), 2 deletions(-)

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 =