Convert HostCreateFlow and HostCheckFlow to tm() (#910)

core/src/main/java/google/registry/flows/host/HostCreateFlow.java

@@ -21,10 +21,8 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.EppResourceUtils.createRepoId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
-import static google.registry.util.CollectionUtils.union;
 
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
@@ -137,13 +135,11 @@ public final class HostCreateFlow implements TransactionalFlow {
     ImmutableSet<ImmutableObject> entitiesToSave =
         ImmutableSet.of(
             newHost,
-            historyBuilder.build(),
+            historyBuilder.build().toChildHistoryEntity(),
             ForeignKeyIndex.create(newHost, newHost.getDeletionTime()),
             EppResourceIndex.create(Key.create(newHost)));
     if (superordinateDomain.isPresent()) {
-      entitiesToSave =
+      tm().update(
-          union(
-              entitiesToSave,
               superordinateDomain
                   .get()
                   .asBuilder()
@@ -153,7 +149,7 @@ public final class HostCreateFlow implements TransactionalFlow {
       // they are only written as NS records from the referencing domain.
       dnsQueue.addHostRefreshTask(targetId);
     }
-    ofy().save().entities(entitiesToSave);
+    tm().insertAll(entitiesToSave);
     return responseBuilder.setResData(HostCreateData.create(targetId, now)).build();
   }
 

core/src/main/java/google/registry/model/EppResourceUtils.java

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 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.isAtOrAfter;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.latestOf;
@@ -135,7 +136,7 @@ public final class EppResourceUtils {
         useCache
             ? ForeignKeyIndex.loadCached(clazz, ImmutableList.of(foreignKey), now)
                 .getOrDefault(foreignKey, null)
-            : ofy().load().type(ForeignKeyIndex.mapToFkiClass(clazz)).id(foreignKey).now();
+            : ForeignKeyIndex.load(clazz, foreignKey, now);
     // The value of fki.getResourceKey() might be null for hard-deleted prober data.
     if (fki == null || isAtOrAfter(now, fki.getDeletionTime()) || fki.getResourceKey() == null) {
       return Optional.empty();
@@ -143,7 +144,7 @@ public final class EppResourceUtils {
     T resource =
         useCache
             ? EppResource.loadCached(fki.getResourceKey())
-            : tm().maybeLoad(fki.getResourceKey()).orElse(null);
+            : transactIfJpaTm(() -> tm().maybeLoad(fki.getResourceKey()).orElse(null));
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return Optional.empty();
     }

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -15,9 +15,11 @@
 package google.registry.model.index;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.TypeUtils.instantiate;
 
@@ -29,6 +31,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Multimaps;
 import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
@@ -44,6 +47,8 @@ import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.util.NonFinalForTesting;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -76,13 +81,21 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource>
       implements DatastoreOnlyEntity {}
 
-  static final ImmutableMap<Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
+  private static final ImmutableMap<
+          Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
       RESOURCE_CLASS_TO_FKI_CLASS =
           ImmutableMap.of(
               ContactResource.class, ForeignKeyContactIndex.class,
               DomainBase.class, ForeignKeyDomainIndex.class,
               HostResource.class, ForeignKeyHostIndex.class);
 
+  private static final ImmutableMap<Class<? extends EppResource>, String>
+      RESOURCE_CLASS_TO_FKI_PROPERTY =
+          ImmutableMap.of(
+              ContactResource.class, "contactId",
+              DomainBase.class, "fullyQualifiedDomainName",
+              HostResource.class, "fullyQualifiedHostName");
+
   @Id String foreignKey;
 
   /**
@@ -179,9 +192,42 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> load(
       Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
-    return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
+    if (tm().isOfy()) {
-        .filter(e -> now.isBefore(e.getValue().deletionTime))
+      return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
-        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+          .filter(e -> now.isBefore(e.getValue().deletionTime))
+          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+    } else {
+      String property = RESOURCE_CLASS_TO_FKI_PROPERTY.get(clazz);
+      List<E> entities =
+          tm().transact(
+                  () -> {
+                    String entityName =
+                        jpaTm().getEntityManager().getMetamodel().entity(clazz).getName();
+                    return jpaTm()
+                        .getEntityManager()
+                        .createQuery(
+                            String.format(
+                                "FROM %s WHERE %s IN :propertyValue and deletionTime > :now ",
+                                entityName, property),
+                            clazz)
+                        .setParameter("propertyValue", foreignKeys)
+                        .setParameter("now", now)
+                        .getResultList();
+                  });
+      // We need to find and return the entities with the maximum deletionTime for each foreign key.
+      return Multimaps.index(entities, EppResource::getForeignKey).asMap().entrySet().stream()
+          .map(
+              entry ->
+                  Maps.immutableEntry(
+                      entry.getKey(),
+                      entry.getValue().stream()
+                          .max(Comparator.comparing(EppResource::getDeletionTime))
+                          .get()))
+          .collect(
+              toImmutableMap(
+                  Map.Entry::getKey,
+                  entry -> create(entry.getValue(), entry.getValue().getDeletionTime())));
+    }
   }
 
   static final CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
@@ -266,7 +312,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
               .filter(entry -> entry.getValue().isPresent())
               .filter(entry -> now.isBefore(entry.getValue().get().getDeletionTime()))
               .collect(
-                  ImmutableMap.toImmutableMap(
+                  toImmutableMap(
                       entry -> entry.getKey().getName(),
                       entry -> (ForeignKeyIndex<E>) entry.getValue().get()));
       return fkisFromCache;

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -29,6 +29,11 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
+import google.registry.model.ImmutableObject;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
@@ -56,6 +61,18 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final Retrier retrier = new Retrier(new SystemSleeper(), 3);
 
+  // The entity of classes in this set will be simply ignored when passed to modification
+  // operations, i.e. insert, put, update and delete. This is to help maintain a single code path
+  // when we switch from ofy() to tm() for the database migration as we don't need have a condition
+  // to exclude the Datastore specific entities when the underlying tm() is jpaTm().
+  // TODO(b/176108270): Remove this property after database migration.
+  private static final ImmutableSet<Class<? extends ImmutableObject>> IGNORED_ENTITY_CLASSES =
+      ImmutableSet.of(
+          EppResourceIndex.class,
+          ForeignKeyContactIndex.class,
+          ForeignKeyDomainIndex.class,
+          ForeignKeyHostIndex.class);
+
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
   private final Clock clock;
@@ -228,6 +245,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void insert(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().persist(entity);
     transactionInfo.get().addUpdate(entity);
@@ -253,6 +273,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void put(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().merge(entity);
     transactionInfo.get().addUpdate(entity);
@@ -278,6 +301,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void update(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     checkArgument(exists(entity), "Given entity does not exist");
     getEntityManager().merge(entity);
@@ -414,6 +440,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void delete(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
@@ -464,6 +493,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  private static boolean isEntityOfIgnoredClass(Object entity) {
+    return IGNORED_ENTITY_CLASSES.contains(entity.getClass());
+  }
+
   private static ImmutableSet<EntityId> getEntityIdsFromEntity(
       EntityType<?> entityType, Object entity) {
     if (entityType.hasSingleIdAttribute()) {

core/src/test/java/google/registry/flows/FlowTestCase.java

@@ -100,7 +100,11 @@ public abstract class FlowTestCase<F extends Flow> {
 
   @RegisterExtension
   final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+      AppEngineExtension.builder()
+          .withClock(clock)
+          .withDatastoreAndCloudSql()
+          .withTaskQueue()
+          .build();
 
   @BeforeEach
   public void beforeEachFlowTestCase() {
@@ -288,7 +292,7 @@ public abstract class FlowTestCase<F extends Flow> {
           e);
     }
     // Clear the cache so that we don't see stale results in tests.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     return output;
   }
 

core/src/test/java/google/registry/flows/ResourceFlowTestCase.java

@@ -19,6 +19,8 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.tmch.ClaimsListShardTest.createTestClaimsListShard;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
 import static google.registry.testing.LogsSubject.assertAboutLogs;
 import static google.registry.testing.TaskQueueHelper.assertTasksEnqueued;
@@ -72,7 +74,7 @@ public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource
   protected R reloadResourceByForeignKey(DateTime now) throws Exception {
     // Force the session to be cleared so that when we read it back, we read from Datastore and not
     // from the transaction's session cache.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     return loadByForeignKey(getResourceClass(), getUniqueIdFromCommand(), now).orElse(null);
   }
 
@@ -83,9 +85,9 @@ public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource
 
   protected <T extends EppResource> T reloadResourceAndCloneAtTime(T resource, DateTime now) {
     // Force the session to be cleared.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     @SuppressWarnings("unchecked")
-    T refreshedResource = (T) ofy().load().entity(resource).now().cloneProjectedAtTime(now);
+    T refreshedResource = (T) transactIfJpaTm(() -> tm().load(resource)).cloneProjectedAtTime(now);
     return refreshedResource;
   }
 

core/src/test/java/google/registry/flows/host/HostCheckFlowTest.java

@@ -24,16 +24,18 @@ import google.registry.flows.EppException;
 import google.registry.flows.ResourceCheckFlowTestCase;
 import google.registry.flows.exceptions.TooManyResourceChecksException;
 import google.registry.model.host.HostResource;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 
 /** Unit tests for {@link HostCheckFlow}. */
+@DualDatabaseTest
 class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostResource> {
 
   HostCheckFlowTest() {
     setEppInput("host_check.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testNothingExists() throws Exception {
     // These ids come from the check xml.
     doCheckTest(
@@ -42,7 +44,7 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExists() throws Exception {
     persistActiveHost("ns1.example.tld");
     // These ids come from the check xml.
@@ -52,7 +54,7 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExistsButWasDeleted() throws Exception {
     persistDeletedHost("ns1.example.tld", clock.nowUtc().minusDays(1));
     // These ids come from the check xml.
@@ -62,27 +64,27 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testXmlMatches() throws Exception {
     persistActiveHost("ns2.example.tld");
     runFlowAssertResponse(loadFile("host_check_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test50IdsAllowed() throws Exception {
     // Make sure we don't have a regression that reduces the number of allowed checks.
     setEppInput("host_check_50.xml");
     runFlow();
   }
 
-  @Test
+  @TestOfyAndSql
   void testTooManyIds() {
     setEppInput("host_check_51.xml");
     EppException thrown = assertThrows(TooManyResourceChecksException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-host-check");

core/src/test/java/google/registry/flows/host/HostCreateFlowTest.java

@@ -16,6 +16,7 @@ package google.registry.flows.host;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.assertNoBillingEvents;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.createTlds;
@@ -53,10 +54,12 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link HostCreateFlow}. */
+@DualDatabaseTest
 class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResource> {
 
   private void setEppHostCreateInput(String hostName, String hostAddrs) {
@@ -90,7 +93,9 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
         .hasOnlyOneHistoryEntryWhich()
         .hasType(HistoryEntry.Type.HOST_CREATE);
     assertNoBillingEvents();
-    assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    if (tm().isOfy()) {
+      assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    }
   }
 
   private void doSuccessfulInternalTest(String tld) throws Exception {
@@ -100,19 +105,19 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     dryRunFlowAssertResponse(loadFile("host_create_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_externalNeverExisted() throws Exception {
     doSuccessfulTest();
     assertAboutHosts().that(reloadResourceByForeignKey()).hasSuperordinateDomain(null);
     assertNoDnsTasksEnqueued();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_internalNeverExisted() throws Exception {
     doSuccessfulInternalTest("tld");
     HostResource host = reloadResourceByForeignKey();
@@ -123,7 +128,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertDnsTasksEnqueued("ns1.example.tld");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_multipartTLDsAndInvalidHost() {
     createTlds("bar.tld", "tld");
 
@@ -132,7 +137,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_externalExistedButWasDeleted() throws Exception {
     persistDeletedHost(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     doSuccessfulTest();
@@ -140,7 +145,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertNoDnsTasksEnqueued();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_internalExistedButWasDeleted() throws Exception {
     persistDeletedHost(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     doSuccessfulInternalTest("tld");
@@ -152,7 +157,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertDnsTasksEnqueued("ns1.example.tld");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_subordinateNeedsIps() {
     setEppHostCreateInput("ns1.example.tld", null);
     createTld("tld");
@@ -161,7 +166,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_externalMustNotHaveIps() {
     setEppHostCreateInputWithIps("ns1.example.external");
     createTld("tld");
@@ -170,7 +175,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_superordinateMissing() {
     setEppHostCreateInput("ns1.example.tld", null);
     createTld("tld");
@@ -179,7 +184,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertThat(thrown).hasMessageThat().contains("(example.tld)");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_superordinateInPendingDelete() {
     setEppHostCreateInputWithIps("ns1.example.tld");
     createTld("tld");
@@ -197,7 +202,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
         .contains("Superordinate domain for this hostname is in pending delete");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_alreadyExists() throws Exception {
     setEppHostCreateInput("ns1.example.tld", null);
     persistActiveHost(getUniqueIdFromCommand());
@@ -209,7 +214,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
             String.format("Object with given ID (%s) already exists", getUniqueIdFromCommand()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_resourceContention() throws Exception {
     setEppHostCreateInput("ns1.example.tld", null);
     String targetId = getUniqueIdFromCommand();
@@ -226,14 +231,14 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonLowerCaseHostname() {
     setEppHostCreateInput("ns1.EXAMPLE.tld", null);
     EppException thrown = assertThrows(HostNameNotLowerCaseException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonPunyCodedHostname() {
     setEppHostCreateInput("ns1.çauçalito.みんな", null);
     HostNameNotPunyCodedException thrown =
@@ -241,21 +246,21 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertThat(thrown).hasMessageThat().contains("expected ns1.xn--aualito-txac.xn--q9jyb4c");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonCanonicalHostname() {
     setEppHostCreateInput("ns1.example.tld.", null);
     EppException thrown = assertThrows(HostNameNotNormalizedException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_longHostName() {
     setEppHostCreateInputWithIps("a" + Strings.repeat(".labelpart", 25) + ".tld");
     EppException thrown = assertThrows(HostNameTooLongException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_ip4AddressWithIp6Declaration() {
     setEppHostCreateInput(
         "ns1.example.tld",
@@ -272,37 +277,37 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badCharacter() {
     doFailingHostNameTest("foo bar", InvalidHostNameException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowPublicSuffix() {
     doFailingHostNameTest("example.tld", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowCcTld() {
     doFailingHostNameTest("foo.co.uk", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_barePublicSuffix() {
     doFailingHostNameTest("com", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_bareCcTld() {
     doFailingHostNameTest("co.uk", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowNewTld() {
     doFailingHostNameTest("example.lol", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_ccTldInBailiwick() {
     createTld("co.uk");
     setEppHostCreateInputWithIps("foo.co.uk");
@@ -310,7 +315,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-host-create");

core/src/test/java/google/registry/flows/host/HostFlowUtilsTest.java

@@ -26,66 +26,68 @@ import google.registry.flows.host.HostFlowUtils.HostNameTooLongException;
 import google.registry.flows.host.HostFlowUtils.HostNameTooShallowException;
 import google.registry.flows.host.HostFlowUtils.InvalidHostNameException;
 import google.registry.testing.AppEngineExtension;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link HostFlowUtils}. */
+@DualDatabaseTest
 class HostFlowUtilsTest {
 
   @RegisterExtension
   final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostName_validates() throws Exception {
     assertThat(validateHostName("host.example.com").toString()).isEqualTo("host.example.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_validates() throws Exception {
     assertThat(validateHostName("host.blogspot.com").toString()).isEqualTo("host.blogspot.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_multipartTLD_validates() throws Exception {
     assertThat(validateHostName("ns1.host.co.uk").toString()).isEqualTo("ns1.host.co.uk");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_multipartTLD_tooShallow() {
     assertThrows(
         HostNameTooShallowException.class, () -> validateHostName("host.co.uk").toString());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameTooLong() {
     assertThrows(
         HostNameTooLongException.class,
         () -> validateHostName(Strings.repeat("na", 200) + ".wat.man"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotLowerCase() {
     assertThrows(HostNameNotLowerCaseException.class, () -> validateHostName("NA.CAPS.TLD"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotPunyCoded() {
     assertThrows(
         HostNameNotPunyCodedException.class, () -> validateHostName("motörhead.death.metal"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotNormalized() {
     assertThrows(HostNameNotNormalizedException.class, () -> validateHostName("root.node.yeah."));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameHasLeadingHyphen() {
     assertThrows(InvalidHostNameException.class, () -> validateHostName("-giga.mega.tld"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameTooShallow() {
     assertThrows(HostNameTooShallowException.class, () -> validateHostName("domain.tld"));
   }

core/src/test/java/google/registry/model/EppResourceUtilsTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadAtPointInTime;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newHostResource;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrars;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DatabaseHelper.persistResourceWithCommitLog;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -26,16 +27,19 @@ import static org.joda.time.DateTimeZone.UTC;
 import google.registry.model.host.HostResource;
 import google.registry.model.ofy.Ofy;
 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 google.registry.testing.TestOfyOnly;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Tests for {@link EppResourceUtils}. */
-public class EppResourceUtilsTest {
+@DualDatabaseTest
+class EppResourceUtilsTest {
 
   @RegisterExtension
   public final AppEngineExtension appEngine =
@@ -51,7 +55,7 @@ public class EppResourceUtilsTest {
     inject.setStaticField(Ofy.class, "clock", clock);
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadAtPointInTime_beforeCreated_returnsNull() {
     clock.advanceOneMilli();
     // Don't save a commit log, we shouldn't need one.
@@ -62,7 +66,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minus(Duration.millis(1))).now()).isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadAtPointInTime_atOrAfterLastAutoUpdateTime_returnsResource() {
     clock.advanceOneMilli();
     // Don't save a commit log, we shouldn't need one.
@@ -73,8 +77,9 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc()).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_usingIntactRevisionHistory_returnsMutationValue() {
+    persistNewRegistrars("OLD", "NEW");
     clock.advanceOneMilli();
     // Save resource with a commit log that we can read in later as a revisions map value.
     HostResource oldHost = persistResourceWithCommitLog(
@@ -94,7 +99,7 @@ public class EppResourceUtilsTest {
         .isEqualTo(oldHost);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_brokenRevisionHistory_returnsResourceAsIs() {
     // Don't save a commit log since we want to test the handling of a broken revisions key.
     HostResource oldHost = persistResource(
@@ -114,7 +119,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minusMillis(1)).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_fallback_returnsMutationValueForOldestRevision() {
     clock.advanceOneMilli();
     // Save a commit log that we can fall back to.
@@ -136,7 +141,7 @@ public class EppResourceUtilsTest {
         .isEqualTo(oldHost);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_ultimateFallback_onlyOneRevision_returnsCurrentResource() {
     clock.advanceOneMilli();
     // Don't save a commit log; we want to test that we load from the current resource.
@@ -151,7 +156,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minusMillis(1)).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_moreThanThirtyDaysInPast_historyIsPurged() {
     clock.advanceOneMilli();
     HostResource host =

core/src/test/java/google/registry/model/index/ForeignKeyIndexTest.java

@@ -20,6 +20,7 @@ import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.deleteResource;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistActiveHost;
 import static google.registry.testing.DatabaseHelper.persistDeletedHost;
@@ -29,16 +30,21 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import com.google.common.collect.ImmutableList;
 import google.registry.model.EntityTestCase;
 import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestCacheExtension;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link ForeignKeyIndex}. */
-public class ForeignKeyIndexTest extends EntityTestCase {
+@DualDatabaseTest
+class ForeignKeyIndexTest extends EntityTestCase {
 
   @RegisterExtension
   public final TestCacheExtension testCacheExtension =
@@ -49,7 +55,17 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     createTld("com");
   }
 
-  @Test
+  @TestSqlOnly
+  void testModifyForeignKeyIndex_notThrowExceptionInSql() {
+    DomainBase domainBase = newDomainBase("test.com");
+    ForeignKeyIndex<DomainBase> fki = ForeignKeyIndex.create(domainBase, fakeClock.nowUtc());
+    tm().transact(() -> tm().insert(fki));
+    tm().transact(() -> tm().put(fki));
+    tm().transact(() -> tm().delete(fki));
+    tm().transact(() -> tm().update(fki));
+  }
+
+  @TestOfyOnly
   void testPersistence() {
     // Persist a host and implicitly persist a ForeignKeyIndex for it.
     HostResource host = persistActiveHost("ns1.example.com");
@@ -59,7 +75,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     assertThat(fki.getDeletionTime()).isEqualTo(END_OF_TIME);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
     // Persist a host and implicitly persist a ForeignKeyIndex for it.
     persistActiveHost("ns1.example.com");
@@ -68,38 +84,50 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         "deletionTime");
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadForNonexistentForeignKey_returnsNull() {
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadForDeletedForeignKey_returnsNull() {
     HostResource host = persistActiveHost("ns1.example.com");
-    persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    if (tm().isOfy()) {
+      persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    } else {
+      persistResource(host.asBuilder().setDeletionTime(fakeClock.nowUtc().minusDays(1)).build());
+    }
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoad_newerKeyHasBeenSoftDeleted() {
     HostResource host1 = persistActiveHost("ns1.example.com");
     fakeClock.advanceOneMilli();
-    ForeignKeyHostIndex fki = new ForeignKeyHostIndex();
+    if (tm().isOfy()) {
-    fki.foreignKey = "ns1.example.com";
+      ForeignKeyHostIndex fki = new ForeignKeyHostIndex();
-    fki.topReference = host1.createVKey();
+      fki.foreignKey = "ns1.example.com";
-    fki.deletionTime = fakeClock.nowUtc();
+      fki.topReference = host1.createVKey();
-    persistResource(fki);
+      fki.deletionTime = fakeClock.nowUtc();
+      persistResource(fki);
+    } else {
+      persistResource(host1.asBuilder().setDeletionTime(fakeClock.nowUtc()).build());
+    }
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testBatchLoad_skipsDeletedAndNonexistent() {
     persistActiveHost("ns1.example.com");
     HostResource host = persistActiveHost("ns2.example.com");
-    persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    if (tm().isOfy()) {
+      persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    } else {
+      persistResource(host.asBuilder().setDeletionTime(fakeClock.nowUtc().minusDays(1)).build());
+    }
     assertThat(
             ForeignKeyIndex.load(
                     HostResource.class,
@@ -109,7 +137,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDeadCodeThatDeletedScrapCommandsReference() {
     persistActiveHost("omg");
     assertThat(ForeignKeyIndex.load(HostResource.class, "omg", fakeClock.nowUtc()).getForeignKey())
@@ -124,7 +152,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     return ForeignKeyIndex.load(ContactResource.class, contactId, fakeClock.nowUtc());
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesNonexistenceOfHosts() {
     assertThat(
             ForeignKeyIndex.loadCached(
@@ -144,7 +172,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns4.example.com", loadHostFki("ns4.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesExistenceOfHosts() {
     HostResource host1 = persistActiveHost("ns1.example.com");
     HostResource host2 = persistActiveHost("ns2.example.com");
@@ -172,7 +200,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
             "ns3.example.com", loadHostFki("ns3.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_doesntSeeHostChangesWhileCacheIsValid() {
     HostResource originalHost = persistActiveHost("ns1.example.com");
     ForeignKeyIndex<HostResource> originalFki = loadHostFki("ns1.example.com");
@@ -195,7 +223,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com", originalFki);
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_filtersOutSoftDeletedHosts() {
     persistActiveHost("ns1.example.com");
     persistDeletedHost("ns2.example.com", fakeClock.nowUtc().minusDays(1));
@@ -207,7 +235,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com", loadHostFki("ns1.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesContactFkis() {
     persistActiveContact("contactid1");
     ForeignKeyIndex<ContactResource> fki1 = loadContactFki("contactid1");

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

@@ -454,10 +454,6 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
           .setEnvIsAdmin(userInfo.isAdmin());
     }
 
-    if (clock != null) {
-      helper.setClock(() -> clock.nowUtc().getMillis());
-    }
-
     if (withLocalModules) {
       helper.setEnvInstance("0");
     }

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

@@ -18,6 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import google.registry.persistence.transaction.TransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import java.lang.reflect.Field;
@@ -160,6 +161,14 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
 
   private static boolean isDualDatabaseTest(ExtensionContext context) {
     Object testInstance = context.getTestInstance().orElseThrow(RuntimeException::new);
-    return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class);
+    // If the test method is declared in its parent class,
+    // e.g. google.registry.flows.ResourceFlowTestCase.testRequiresLogin,
+    // we don't consider it is a DualDatabaseTest. This is because there may exist some subclasses
+    // that have not been migrated to DualDatabaseTest.
+    boolean isDeclaredTestMethod =
+        ImmutableSet.copyOf(testInstance.getClass().getDeclaredMethods())
+            .contains(context.getTestMethod().orElseThrow(RuntimeException::new));
+    return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class)
+        && isDeclaredTestMethod;
   }
 }

core/src/test/java/google/registry/tools/CommandTestCase.java

@@ -17,7 +17,6 @@ package google.registry.tools;
 import static com.google.common.collect.Iterables.concat;
 import static com.google.common.collect.Iterables.toArray;
 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.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -118,7 +117,7 @@ public abstract class CommandTestCase<C extends Command> {
     } finally {
       // Clear the session cache so that subsequent reads for verification purposes hit Datastore.
       // This primarily matters for AutoTimestamp fields, which otherwise won't have updated values.
-      ofy().clearSessionCache();
+      tm().clearSessionCache();
       // Reset back to UNITTEST environment.
       RegistryToolEnvironment.UNITTEST.setup(systemPropertyExtension);
     }
@@ -192,7 +191,7 @@ public abstract class CommandTestCase<C extends Command> {
 
   /** Returns count of all poll messages in Datastore. */
   int getPollMessageCount() {
-    return ofy().load().type(PollMessage.class).count();
+    return transactIfJpaTm(() -> tm().loadAll(PollMessage.class).size());
   }
 
   /**