Convert RDE classes to use tm() (#1044)

This is mostly just using the generic Cursor load methods with the
slight difference that before we relied on ofy() returning null on
absent entities.

core/src/main/java/google/registry/model/rde/RdeRevision.java

@@ -17,6 +17,7 @@ package google.registry.model.rde;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.rde.RdeNamingUtils.makePartialName;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.base.VerifyException;
 import com.googlecode.objectify.Key;
@@ -97,7 +98,8 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
     Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, id);
     Optional<RdeRevision> revisionOptional =
-        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey)));
     return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
   }
 

core/src/main/java/google/registry/rde/EscrowTaskRunner.java

@@ -14,8 +14,8 @@
 
 package google.registry.rde;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.flogger.FluentLogger;
 import google.registry.model.common.Cursor;
@@ -90,8 +90,13 @@ class EscrowTaskRunner {
         () -> {
           logger.atInfo().log("TLD: %s", registry.getTld());
           DateTime startOfToday = clock.nowUtc().withTimeAtStartOfDay();
-          Cursor cursor = ofy().load().key(Cursor.createKey(cursorType, registry)).now();
-          final DateTime nextRequiredRun = (cursor == null ? startOfToday : cursor.getCursorTime());
+          DateTime nextRequiredRun =
+              transactIfJpaTm(
+                      () ->
+                          tm().loadByKeyIfPresent(
+                                  Cursor.createVKey(cursorType, registry.getTldStr())))
+                  .map(Cursor::getCursorTime)
+                  .orElse(startOfToday);
           if (nextRequiredRun.isAfter(startOfToday)) {
             throw new NoContentException("Already completed");
           }

core/src/main/java/google/registry/rde/PendingDepositChecker.java

@@ -15,8 +15,8 @@
 package google.registry.rde;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.google.common.collect.ImmutableSetMultimap;
@@ -28,6 +28,7 @@ import google.registry.model.registry.Registries;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.util.Clock;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
@@ -89,13 +90,15 @@ public final class PendingDepositChecker {
         continue;
       }
       // Avoid creating a transaction unless absolutely necessary.
-      Cursor cursor = ofy().load().key(Cursor.createKey(cursorType, registry)).now();
-      DateTime cursorValue = (cursor != null ? cursor.getCursorTime() : startingPoint);
+      Optional<Cursor> maybeCursor =
+          transactIfJpaTm(
+              () -> tm().loadByKeyIfPresent(Cursor.createVKey(cursorType, registry.getTldStr())));
+      DateTime cursorValue = maybeCursor.map(Cursor::getCursorTime).orElse(startingPoint);
       if (isBeforeOrAt(cursorValue, now)) {
         DateTime watermark =
-            (cursor != null
-                ? cursor.getCursorTime()
-                : transactionallyInitializeCursor(registry, cursorType, startingPoint));
+            maybeCursor
+                .map(Cursor::getCursorTime)
+                .orElse(transactionallyInitializeCursor(registry, cursorType, startingPoint));
         if (isBeforeOrAt(watermark, now)) {
           builder.put(tld, PendingDeposit.create(tld, watermark, mode, cursorType, interval));
         }
@@ -108,9 +111,10 @@ public final class PendingDepositChecker {
       final Registry registry, final CursorType cursorType, final DateTime initialValue) {
     return tm().transact(
             () -> {
-              Cursor cursor = ofy().load().key(Cursor.createKey(cursorType, registry)).now();
-              if (cursor != null) {
-                return cursor.getCursorTime();
+              Optional<Cursor> maybeCursor =
+                  tm().loadByKeyIfPresent(Cursor.createVKey(cursorType, registry.getTldStr()));
+              if (maybeCursor.isPresent()) {
+                return maybeCursor.get().getCursorTime();
               }
               tm().put(Cursor.create(cursorType, initialValue, registry));
               return initialValue;

core/src/main/java/google/registry/rde/RdeReportAction.java

@@ -17,8 +17,9 @@ package google.registry.rde;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.rde.RdeMode.FULL;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
@@ -77,7 +78,7 @@ public final class RdeReportAction implements Runnable, EscrowTask {
   @Override
   public void runWithLock(DateTime watermark) throws Exception {
     Cursor cursor =
-        ofy().load().key(Cursor.createKey(CursorType.RDE_UPLOAD, Registry.get(tld))).now();
+        transactIfJpaTm(() -> tm().loadByKey(Cursor.createVKey(CursorType.RDE_UPLOAD, tld)));
     DateTime cursorTime = getCursorTimeOrStartOfTime(cursor);
     if (isBeforeOrAt(cursorTime, watermark)) {
       throw new NoContentException(

core/src/main/java/google/registry/rde/RdeUploadAction.java

@@ -18,12 +18,14 @@ import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static com.jcraft.jsch.ChannelSftp.OVERWRITE;
+import static google.registry.model.common.Cursor.CursorType.RDE_STAGING;
 import static google.registry.model.common.Cursor.CursorType.RDE_UPLOAD_SFTP;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.rde.RdeMode.FULL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.request.Action.Method.POST;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static java.util.Arrays.asList;
 
@@ -131,8 +133,7 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
   @Override
   public void runWithLock(final DateTime watermark) throws Exception {
     logger.atInfo().log("Verifying readiness to upload the RDE deposit.");
-    Cursor cursor =
-        ofy().load().key(Cursor.createKey(CursorType.RDE_STAGING, Registry.get(tld))).now();
+    Cursor cursor = transactIfJpaTm(() -> tm().loadByKey(Cursor.createVKey(RDE_STAGING, tld)));
     DateTime stagingCursorTime = getCursorTimeOrStartOfTime(cursor);
     if (isBeforeOrAt(stagingCursorTime, watermark)) {
       throw new NoContentException(
@@ -141,9 +142,10 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
                   + "last RDE staging completion was at %s",
               tld, watermark, stagingCursorTime));
     }
-    Cursor sftpCursor =
-        ofy().load().key(Cursor.createKey(RDE_UPLOAD_SFTP, Registry.get(tld))).now();
-    DateTime sftpCursorTime = getCursorTimeOrStartOfTime(sftpCursor);
+    DateTime sftpCursorTime =
+        transactIfJpaTm(() -> tm().loadByKeyIfPresent(Cursor.createVKey(RDE_UPLOAD_SFTP, tld)))
+            .map(Cursor::getCursorTime)
+            .orElse(START_OF_TIME);
     Duration timeSinceLastSftp = new Duration(sftpCursorTime, clock.nowUtc());
     if (timeSinceLastSftp.isShorterThan(sftpCooldown)) {
       throw new NoContentException(

core/src/test/java/google/registry/rde/EscrowTaskRunnerTest.java

@@ -15,8 +15,9 @@
 package google.registry.rde;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.joda.time.Duration.standardDays;
 import static org.joda.time.Duration.standardSeconds;
@@ -31,16 +32,18 @@ import google.registry.rde.EscrowTaskRunner.EscrowTask;
 import google.registry.request.HttpException.NoContentException;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeLockHandler;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link EscrowTaskRunner}. */
+@DualDatabaseTest
 public class EscrowTaskRunnerTest {
 
   @RegisterExtension
@@ -62,7 +65,7 @@ public class EscrowTaskRunnerTest {
     runner.clock = clock;
     runner.lockHandler = new FakeLockHandler(true);
     previousDateTimeZone = DateTimeZone.getDefault();
-    DateTimeZone.setDefault(DateTimeZone.forID("America/New_York"));  // Make sure UTC stuff works.
+    DateTimeZone.setDefault(DateTimeZone.forID("America/New_York")); // Make sure UTC stuff works.
   }
 
   @AfterEach
@@ -70,7 +73,7 @@ public class EscrowTaskRunnerTest {
     DateTimeZone.setDefault(previousDateTimeZone);
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_cursorIsToday_advancesCursorToTomorrow() throws Exception {
     clock.setTo(DateTime.parse("2006-06-06T00:30:00Z"));
     persistResource(
@@ -78,23 +81,22 @@ public class EscrowTaskRunnerTest {
     runner.lockRunAndRollForward(
         task, registry, standardSeconds(30), CursorType.RDE_STAGING, standardDays(1));
     verify(task).runWithLock(DateTime.parse("2006-06-06TZ"));
-    ofy().clearSessionCache();
-    Cursor cursor = ofy().load().key(Cursor.createKey(CursorType.RDE_STAGING, registry)).now();
+    tm().clearSessionCache();
+    Cursor cursor = loadByKey(Cursor.createVKey(CursorType.RDE_STAGING, registry.getTldStr()));
     assertThat(cursor.getCursorTime()).isEqualTo(DateTime.parse("2006-06-07TZ"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_cursorMissing_assumesTodayAndAdvancesCursorToTomorrow() throws Exception {
     clock.setTo(DateTime.parse("2006-06-06T00:30:00Z"));
     runner.lockRunAndRollForward(
         task, registry, standardSeconds(30), CursorType.RDE_STAGING, standardDays(1));
     verify(task).runWithLock(DateTime.parse("2006-06-06TZ"));
-    Cursor cursor =
-        ofy().load().key(Cursor.createKey(CursorType.RDE_STAGING, Registry.get("lol"))).now();
+    Cursor cursor = loadByKey(Cursor.createVKey(CursorType.RDE_STAGING, "lol"));
     assertThat(cursor.getCursorTime()).isEqualTo(DateTime.parse("2006-06-07TZ"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_cursorInTheFuture_doesNothing() {
     clock.setTo(DateTime.parse("2006-06-06T00:30:00Z"));
     persistResource(
@@ -108,7 +110,7 @@ public class EscrowTaskRunnerTest {
     assertThat(thrown).hasMessageThat().contains("Already completed");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_lockIsntAvailable_throws503() {
     String lockName = "EscrowTaskRunner " + task.getClass().getSimpleName();
     clock.setTo(DateTime.parse("2006-06-06T00:30:00Z"));

core/src/test/java/google/registry/rde/PendingDepositCheckerTest.java

@@ -17,29 +17,33 @@ package google.registry.rde;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.common.Cursor.CursorType.BRDA;
 import static google.registry.model.common.Cursor.CursorType.RDE_STAGING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.rde.RdeMode.FULL;
 import static google.registry.model.rde.RdeMode.THIN;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadByKey;
+import static google.registry.testing.DatabaseHelper.loadByKeyIfPresent;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.joda.time.DateTimeConstants.TUESDAY;
 import static org.joda.time.Duration.standardDays;
 
 import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.truth.Truth8;
 import google.registry.model.common.Cursor;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.ofy.Ofy;
 import google.registry.model.registry.Registry;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link PendingDepositChecker}. */
+@DualDatabaseTest
 public class PendingDepositCheckerTest {
 
   @RegisterExtension public final InjectExtension inject = new InjectExtension();
@@ -60,14 +64,14 @@ public class PendingDepositCheckerTest {
     checker.rdeInterval = standardDays(1);
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_noTldsWithEscrowEnabled_returnsEmpty() {
     createTld("pal");
     createTld("fun");
     assertThat(checker.getTldsAndWatermarksPendingDepositForRdeAndBrda()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_firstDeposit_depositsRdeTodayAtMidnight() {
     clock.setTo(DateTime.parse("2000-01-01T08:00Z"));  // Saturday
     createTldWithEscrowEnabled("lol");
@@ -78,7 +82,7 @@ public class PendingDepositCheckerTest {
                 "lol", DateTime.parse("2000-01-01TZ"), FULL, RDE_STAGING, standardDays(1))));
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_firstDepositOnBrdaDay_depositsBothRdeAndBrda() {
     clock.setTo(DateTime.parse("2000-01-04T08:00Z"));  // Tuesday
     createTldWithEscrowEnabled("lol");
@@ -91,19 +95,20 @@ public class PendingDepositCheckerTest {
                 "lol", DateTime.parse("2000-01-04TZ"), THIN, BRDA, standardDays(7))));
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_firstRdeDeposit_initializesCursorToMidnightToday() {
     clock.setTo(DateTime.parse("2000-01-01TZ"));  // Saturday
     createTldWithEscrowEnabled("lol");
     clock.advanceOneMilli();
     Registry registry = Registry.get("lol");
-    assertThat(ofy().load().key(Cursor.createKey(RDE_STAGING, registry)).now()).isNull();
+    Truth8.assertThat(loadByKeyIfPresent(Cursor.createVKey(RDE_STAGING, registry.getTldStr())))
+        .isEmpty();
     checker.getTldsAndWatermarksPendingDepositForRdeAndBrda();
-    assertThat(ofy().load().key(Cursor.createKey(RDE_STAGING, registry)).now().getCursorTime())
+    assertThat(loadByKey(Cursor.createVKey(RDE_STAGING, registry.getTldStr())).getCursorTime())
         .isEqualTo(DateTime.parse("2000-01-01TZ"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_subsequentRdeDeposit_doesntMutateCursor() {
     clock.setTo(DateTime.parse("2000-01-01TZ"));  // Saturday
     createTldWithEscrowEnabled("lol");
@@ -112,11 +117,11 @@ public class PendingDepositCheckerTest {
     setCursor(Registry.get("lol"), RDE_STAGING, yesterday);
     clock.advanceOneMilli();
     checker.getTldsAndWatermarksPendingDepositForRdeAndBrda();
-    Cursor cursor = ofy().load().key(Cursor.createKey(RDE_STAGING, Registry.get("lol"))).now();
+    Cursor cursor = loadByKey(Cursor.createVKey(RDE_STAGING, "lol"));
     assertThat(cursor.getCursorTime()).isEqualTo(yesterday);
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_firstBrdaDepositButNotOnBrdaDay_doesntInitializeCursor() {
     clock.setTo(DateTime.parse("2000-01-01TZ"));  // Saturday
     createTldWithEscrowEnabled("lol");
@@ -124,12 +129,12 @@ public class PendingDepositCheckerTest {
     clock.advanceOneMilli();
     setCursor(registry, RDE_STAGING, DateTime.parse("2000-01-02TZ")); // assume rde is already done
     clock.advanceOneMilli();
-    assertThat(ofy().load().key(Cursor.createKey(BRDA, registry)).now()).isNull();
+    Truth8.assertThat(loadByKeyIfPresent(Cursor.createVKey(BRDA, registry.getTldStr()))).isEmpty();
     assertThat(checker.getTldsAndWatermarksPendingDepositForRdeAndBrda()).isEmpty();
-    assertThat(ofy().load().key(Cursor.createKey(BRDA, registry)).now()).isNull();
+    Truth8.assertThat(loadByKeyIfPresent(Cursor.createVKey(BRDA, registry.getTldStr()))).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_backloggedTwoDays_onlyWantsLeastRecentDay() {
     clock.setTo(DateTime.parse("2000-01-01TZ"));
     createTldWithEscrowEnabled("lol");
@@ -142,7 +147,7 @@ public class PendingDepositCheckerTest {
                 "lol", DateTime.parse("1999-12-30TZ"), FULL, RDE_STAGING, standardDays(1))));
   }
 
-  @Test
+  @TestOfyAndSql
   void testMethod_multipleTldsWithEscrowEnabled_depositsBoth() {
     clock.setTo(DateTime.parse("2000-01-01TZ"));  // Saturday
     createTldWithEscrowEnabled("pal");

core/src/test/java/google/registry/rde/RdeReportActionTest.java

@@ -19,8 +19,8 @@ import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.common.Cursor.CursorType.RDE_REPORT;
 import static google.registry.model.common.Cursor.CursorType.RDE_UPLOAD;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.GcsTestingUtils.writeGcsFile;
 import static javax.servlet.http.HttpServletResponse.SC_BAD_REQUEST;
@@ -51,10 +51,12 @@ import google.registry.request.HttpException.InternalServerErrorException;
 import google.registry.request.HttpException.NoContentException;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeKeyringModule;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.FakeSleeper;
+import google.registry.testing.TestOfyAndSql;
 import google.registry.util.Retrier;
 import google.registry.xjc.XjcXmlTransformer;
 import google.registry.xjc.rdereport.XjcRdeReportReport;
@@ -65,11 +67,11 @@ import java.util.Map;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.ArgumentCaptor;
 
 /** Unit tests for {@link RdeReportAction}. */
+@DualDatabaseTest
 public class RdeReportActionTest {
 
   private static final ByteSource REPORT_XML = RdeTestData.loadBytes("report.xml");
@@ -123,7 +125,7 @@ public class RdeReportActionTest {
     writeGcsFile(gcsService, reportFile, Ghostryde.encode(REPORT_XML.read(), encryptKey));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun() {
     createTld("lol");
     RdeReportAction action = createAction();
@@ -134,7 +136,7 @@ public class RdeReportActionTest {
     verifyNoMoreInteractions(runner);
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock() throws Exception {
     when(httpResponse.getResponseCode()).thenReturn(SC_OK);
     when(httpResponse.getContent()).thenReturn(IIRDEA_GOOD_XML.read());
@@ -160,7 +162,7 @@ public class RdeReportActionTest {
     assertThat(report.getWatermark()).isEqualTo(DateTime.parse("2010-10-17T00:00:00Z"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_uploadNotFinished_throws204() {
     persistResource(
         Cursor.create(RDE_UPLOAD, DateTime.parse("2006-06-06TZ"), Registry.get("test")));
@@ -174,7 +176,7 @@ public class RdeReportActionTest {
                 + "last upload completion was at 2006-06-06T00:00:00.000Z");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_badRequest_throws500WithErrorInfo() throws Exception {
     when(httpResponse.getResponseCode()).thenReturn(SC_BAD_REQUEST);
     when(httpResponse.getContent()).thenReturn(IIRDEA_BAD_XML.read());
@@ -186,7 +188,7 @@ public class RdeReportActionTest {
     assertThat(thrown).hasMessageThat().contains("The structure of the report is invalid.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_fetchFailed_throwsRuntimeException() throws Exception {
     class ExpectedThrownException extends RuntimeException {}
     when(urlFetchService.fetch(any(HTTPRequest.class))).thenThrow(new ExpectedThrownException());
@@ -194,7 +196,7 @@ public class RdeReportActionTest {
         ExpectedThrownException.class, () -> createAction().runWithLock(loadRdeReportCursor()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_socketTimeout_doesRetry() throws Exception {
     when(httpResponse.getResponseCode()).thenReturn(SC_OK);
     when(httpResponse.getContent()).thenReturn(IIRDEA_GOOD_XML.read());
@@ -208,11 +210,7 @@ public class RdeReportActionTest {
   }
 
   private DateTime loadRdeReportCursor() {
-    return ofy()
-        .load()
-        .key(Cursor.createKey(RDE_REPORT, Registry.get("test")))
-        .now()
-        .getCursorTime();
+    return loadByKey(Cursor.createVKey(RDE_REPORT, "test")).getCursorTime();
   }
 
   private static ImmutableMap<String, String> mapifyHeaders(Iterable<HTTPHeader> headers) {

core/src/test/java/google/registry/rde/RdeUploadActionTest.java

@@ -65,6 +65,7 @@ import google.registry.request.HttpException.NoContentException;
 import google.registry.request.RequestParameters;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeKeyringModule;
 import google.registry.testing.FakeResponse;
@@ -72,6 +73,7 @@ import google.registry.testing.FakeSleeper;
 import google.registry.testing.GpgSystemCommandExtension;
 import google.registry.testing.Lazies;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
+import google.registry.testing.TestOfyAndSql;
 import google.registry.testing.sftp.SftpServerExtension;
 import google.registry.util.Retrier;
 import google.registry.util.TaskQueueUtils;
@@ -84,12 +86,12 @@ import java.net.URI;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.api.io.TempDir;
 import org.mockito.stubbing.OngoingStubbing;
 
 /** Unit tests for {@link RdeUploadAction}. */
+@DualDatabaseTest
 public class RdeUploadActionTest {
 
   private static final int BUFFER_SIZE = 64 * 1024;
@@ -198,7 +200,7 @@ public class RdeUploadActionTest {
             });
   }
 
-  @Test
+  @TestOfyAndSql
   void testSocketConnection() throws Exception {
     int port = sftpd.serve("user", "password", folder);
     try (Socket socket = new Socket("localhost", port)) {
@@ -206,7 +208,7 @@ public class RdeUploadActionTest {
     }
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun() {
     createTld("lol");
     RdeUploadAction action = createAction(null);
@@ -220,7 +222,7 @@ public class RdeUploadActionTest {
     verifyNoMoreInteractions(runner);
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_succeedsOnThirdTry() throws Exception {
     int port = sftpd.serve("user", "password", folder);
     URI uploadUrl = URI.create(String.format("sftp://user:password@localhost:%d/", port));
@@ -239,7 +241,7 @@ public class RdeUploadActionTest {
         .containsExactly("tld_2010-10-17_full_S1_R0.ryde", "tld_2010-10-17_full_S1_R0.sig");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_failsAfterThreeAttempts() throws Exception {
     int port = sftpd.serve("user", "password", folder);
     URI uploadUrl = URI.create(String.format("sftp://user:password@localhost:%d/", port));
@@ -253,7 +255,7 @@ public class RdeUploadActionTest {
     assertThat(thrown).hasMessageThat().contains("The crow flies in square circles.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_copiesOnGcs() throws Exception {
     int port = sftpd.serve("user", "password", folder);
     URI uploadUrl = URI.create(String.format("sftp://user:password@localhost:%d/", port));
@@ -275,7 +277,7 @@ public class RdeUploadActionTest {
         .isEqualTo(Files.toByteArray(new File(folder, sigFilename)));
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_resend() throws Exception {
     tm().transact(() -> RdeRevision.saveRevision("tld", DateTime.parse("2010-10-17TZ"), FULL, 1));
     int port = sftpd.serve("user", "password", folder);
@@ -293,7 +295,7 @@ public class RdeUploadActionTest {
         .containsExactly("tld_2010-10-17_full_S1_R1.ryde", "tld_2010-10-17_full_S1_R1.sig");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_producesValidSignature() throws Exception {
     assumeTrue(hasCommand("gpg --version"));
     int port = sftpd.serve("user", "password", folder);
@@ -316,7 +318,7 @@ public class RdeUploadActionTest {
     assertThat(stderr).contains("rde-unittest@registry.test");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_stagingNotFinished_throws204() {
     URI url = URI.create("sftp://user:password@localhost:32323/");
     DateTime stagingCursor = DateTime.parse("2010-10-17TZ");
@@ -331,7 +333,7 @@ public class RdeUploadActionTest {
                 + "last RDE staging completion was at 2010-10-17T00:00:00.000Z");
   }
 
-  @Test
+  @TestOfyAndSql
   void testRunWithLock_sftpCooldownNotPassed_throws204() {
     RdeUploadAction action = createAction(URI.create("sftp://user:password@localhost:32323/"));
     action.sftpCooldown = standardHours(2);

core/src/test/java/google/registry/testing/DatabaseHelper.java

@@ -117,6 +117,7 @@ import google.registry.tmch.LordnTaskUtils;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import javax.annotation.Nullable;
@@ -1265,6 +1266,16 @@ public class DatabaseHelper {
     return transactIfJpaTm(() -> tm().loadByKey(key));
   }
 
+  /**
+   * Loads the specified entity by its key from the DB or empty if it doesn't exist.
+   *
+   * <p>If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for
+   * convenience, so you don't need to wrap it in a transaction at the callsite.
+   */
+  public static <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
+    return transactIfJpaTm(() -> tm().loadByKeyIfPresent(key));
+  }
+
   /**
    * Loads the specified entities by their keys from the DB.
    *