Use TransactionManager APIs in DatastoreHelper (#849)

* Make DatastoreHelper support Postgresql

* Rebase on HEAD

* Resolve comments

* Use put* inside insert* and update*

* Resolve comments

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -60,7 +60,9 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
   @Id
   @Access(AccessType.PROPERTY)
   public String getContactRepoId() {
-    return parent.getName();
+    // We need to handle null case here because Hibernate sometimes accesses this method before
+    // parent gets initialized
+    return parent == null ? null : parent.getName();
   }
 
   /** This method is private because it is only used by Hibernate. */

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -75,7 +75,9 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   @Id
   @Access(AccessType.PROPERTY)
   public String getDomainRepoId() {
-    return parent.getName();
+    // We need to handle null case here because Hibernate sometimes accesses this method before
+    // parent gets initialized
+    return parent == null ? null : parent.getName();
   }
 
   /** This method is private because it is only used by Hibernate. */
@@ -127,14 +129,24 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
    *
    * <p>This will be empty for any DomainHistory/HistoryEntry generated before this field was added,
    * mid-2017, as well as any action that does not generate billable events (e.g. updates).
+   *
+   * <p>This method is dedicated for Hibernate, external caller should use {@link
+   * #getDomainTransactionRecords()}.
    */
   @Access(AccessType.PROPERTY)
   @OneToMany(cascade = {CascadeType.ALL})
   @JoinColumn(name = "historyRevisionId", referencedColumnName = "historyRevisionId")
   @JoinColumn(name = "domainRepoId", referencedColumnName = "domainRepoId")
-  @Override
+  @SuppressWarnings("unused")
-  public Set<DomainTransactionRecord> getDomainTransactionRecords() {
+  private Set<DomainTransactionRecord> getInternalDomainTransactionRecords() {
-    return super.getDomainTransactionRecords();
+    return domainTransactionRecords;
+  }
+
+  /** Sets the domain transaction records. This method is dedicated for Hibernate. */
+  @SuppressWarnings("unused")
+  private void setInternalDomainTransactionRecords(
+      Set<DomainTransactionRecord> domainTransactionRecords) {
+    this.domainTransactionRecords = domainTransactionRecords;
   }
 
   @Id

core/src/main/java/google/registry/model/host/HostHistory.java

@@ -61,7 +61,9 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
   @Id
   @Access(AccessType.PROPERTY)
   public String getHostRepoId() {
-    return parent.getName();
+    // We need to handle null case here because Hibernate sometimes accesses this method before
+    // parent gets initialized
+    return parent == null ? null : parent.getName();
   }
 
   /** This method is private because it is only used by Hibernate. */

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -23,7 +23,9 @@ import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
+import com.googlecode.objectify.Result;
 import google.registry.model.contact.ContactHistory;
 import google.registry.model.host.HostHistory;
 import google.registry.model.reporting.HistoryEntry;
@@ -102,12 +104,22 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void insert(Object entity) {
-    saveEntity(entity);
+    put(entity);
   }
 
   @Override
   public void insertAll(ImmutableCollection<?> entities) {
-    getOfy().save().entities(entities);
+    putAll(entities);
+  }
+
+  @Override
+  public void insertWithoutBackup(Object entity) {
+    putWithoutBackup(entity);
+  }
+
+  @Override
+  public void insertAllWithoutBackup(ImmutableCollection<?> entities) {
+    putAllWithoutBackup(entities);
   }
 
   @Override
@@ -117,17 +129,37 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void putAll(ImmutableCollection<?> entities) {
-    getOfy().save().entities(entities);
+    syncIfTransactionless(getOfy().save().entities(entities));
+  }
+
+  @Override
+  public void putWithoutBackup(Object entity) {
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entity));
+  }
+
+  @Override
+  public void putAllWithoutBackup(ImmutableCollection<?> entities) {
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entities));
   }
 
   @Override
   public void update(Object entity) {
-    saveEntity(entity);
+    put(entity);
   }
 
   @Override
   public void updateAll(ImmutableCollection<?> entities) {
-    getOfy().save().entities(entities);
+    putAll(entities);
+  }
+
+  @Override
+  public void updateWithoutBackup(Object entity) {
+    putWithoutBackup(entity);
+  }
+
+  @Override
+  public void updateAllWithoutBackup(ImmutableCollection<?> entities) {
+    putAllWithoutBackup(entities);
   }
 
   @Override
@@ -158,6 +190,11 @@ public class DatastoreTransactionManager implements TransactionManager {
     return result;
   }
 
+  @Override
+  public <T> T load(T entity) {
+    return ofy().load().entity(entity).now();
+  }
+
   @Override
   public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
     // Keep track of the Key -> VKey mapping so we can translate them back.
@@ -175,13 +212,17 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAll(Class<T> clazz) {
-    // We can do a ofy().load().type(clazz), but this doesn't work in a transaction.
+    return ImmutableList.copyOf(getOfy().load().type(clazz));
-    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
+    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
   }
 
   @Override
   public void delete(VKey<?> key) {
-    getOfy().delete().key(key.getOfyKey()).now();
+    syncIfTransactionless(getOfy().delete().key(key.getOfyKey()));
   }
 
   @Override
@@ -192,7 +233,35 @@ public class DatastoreTransactionManager implements TransactionManager {
         StreamSupport.stream(vKeys.spliterator(), false)
             .map(VKey::getOfyKey)
             .collect(toImmutableList());
-    getOfy().delete().keys(list).now();
+    syncIfTransactionless(getOfy().delete().keys(list));
+  }
+
+  @Override
+  public void delete(Object entity) {
+    syncIfTransactionless(getOfy().delete().entity(entity));
+  }
+
+  @Override
+  public void deleteWithoutBackup(VKey<?> key) {
+    syncIfTransactionless(getOfy().deleteWithoutBackup().key(key.getOfyKey()));
+  }
+
+  @Override
+  public void deleteWithoutBackup(Iterable<? extends VKey<?>> keys) {
+    syncIfTransactionless(
+        getOfy()
+            .deleteWithoutBackup()
+            .keys(Streams.stream(keys).map(VKey::getOfyKey).collect(toImmutableList())));
+  }
+
+  @Override
+  public void deleteWithoutBackup(Object entity) {
+    syncIfTransactionless(getOfy().deleteWithoutBackup().entity(entity));
+  }
+
+  @Override
+  public void clearSessionCache() {
+    getOfy().clearSessionCache();
   }
 
   /**
@@ -209,7 +278,7 @@ public class DatastoreTransactionManager implements TransactionManager {
     if (entity instanceof HistoryEntry) {
       entity = ((HistoryEntry) entity).asHistoryEntry();
     }
-    getOfy().save().entity(entity);
+    syncIfTransactionless(getOfy().save().entity(entity));
   }
 
   @SuppressWarnings("unchecked")
@@ -227,4 +296,19 @@ public class DatastoreTransactionManager implements TransactionManager {
   private <T> T loadNullable(VKey<T> key) {
     return toChildHistoryEntryIfPossible(getOfy().load().key(key.getOfyKey()).now());
   }
+
+  /**
+   * Executes the given {@link Result} instance synchronously if not in a transaction.
+   *
+   * <p>The {@link Result} instance contains a task that will be executed by Objectify
+   * asynchronously. If it is in a transaction, we don't need to execute the task immediately
+   * because it is guaranteed to be done by the end of the transaction. However, if it is not in a
+   * transaction, we need to execute it in case the following code expects that happens before
+   * themselves.
+   */
+  private void syncIfTransactionless(Result<?> result) {
+    if (!inTransaction()) {
+      result.now();
+    }
+  }
 }

core/src/main/java/google/registry/model/registrar/Registrar.java

@@ -32,6 +32,7 @@ import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registries.assertTldsExist;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableSortedCopy;
 import static google.registry.util.PasswordUtils.SALT_SUPPLIER;
@@ -704,10 +705,18 @@ public class Registrar extends ImmutableObject
     return new Builder(clone(this));
   }
 
+  /** Creates a {@link VKey} for this instance. */
   public VKey<Registrar> createVKey() {
     return VKey.create(Registrar.class, clientIdentifier, Key.create(this));
   }
 
+  /** Creates a {@link VKey} for the given {@code registrarId}. */
+  public static VKey<Registrar> createVKey(String registrarId) {
+    checkArgumentNotNull(registrarId, "registrarId must be specified");
+    return VKey.create(
+        Registrar.class, registrarId, Key.create(getCrossTldKey(), Registrar.class, registrarId));
+  }
+
   /** A builder for constructing {@link Registrar}, since it is immutable. */
   public static class Builder extends Buildable.Builder<Registrar> {
     public Builder() {}
@@ -991,8 +1000,7 @@ public class Registrar extends ImmutableObject
   /** Loads and returns a registrar entity by its client id directly from Datastore. */
   public static Optional<Registrar> loadByClientId(String clientId) {
     checkArgument(!Strings.isNullOrEmpty(clientId), "clientId must be specified");
-    return Optional.ofNullable(
+    return transactIfJpaTm(() -> tm().maybeLoad(createVKey(clientId)));
-        ofy().load().type(Registrar.class).parent(getCrossTldKey()).id(clientId).now());
   }
 
   /**

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

@@ -25,6 +25,7 @@ import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
 import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -59,15 +60,21 @@ public final class Registries {
             tm().doTransactionless(
                     () -> {
                       ImmutableSet<String> tlds =
-                          ofy()
+                          ofyOrJpaTm(
-                              .load()
+                              () ->
-                              .type(Registry.class)
+                                  ofy()
-                              .ancestor(getCrossTldKey())
+                                      .load()
-                              .keys()
+                                      .type(Registry.class)
-                              .list()
+                                      .ancestor(getCrossTldKey())
-                              .stream()
+                                      .keys()
-                              .map(Key::getName)
+                                      .list()
-                              .collect(toImmutableSet());
+                                      .stream()
+                                      .map(Key::getName)
+                                      .collect(toImmutableSet()),
+                              () ->
+                                  tm().loadAll(Registry.class).stream()
+                                      .map(Registry::getTldStr)
+                                      .collect(toImmutableSet()));
                       return Registry.getAll(tlds).stream()
                           .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
                           .collect(entriesToImmutableMap());

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

@@ -22,7 +22,6 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Maps.toMap;
 import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -61,6 +60,7 @@ import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.ReservedList;
+import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import google.registry.util.Idn;
 import java.util.Map;
@@ -267,28 +267,24 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
                 public Optional<Registry> load(final String tld) {
                   // Enter a transaction-less context briefly; we don't want to enroll every TLD in
                   // a transaction that might be wrapping this call.
-                  return Optional.ofNullable(
+                  return tm().doTransactionless(() -> tm().maybeLoad(createVKey(tld)));
-                      tm().doTransactionless(
-                              () ->
-                                  ofy()
-                                      .load()
-                                      .key(Key.create(getCrossTldKey(), Registry.class, tld))
-                                      .now()));
                 }
 
                 @Override
                 public Map<String, Optional<Registry>> loadAll(Iterable<? extends String> tlds) {
-                  ImmutableMap<String, Key<Registry>> keysMap =
+                  ImmutableMap<String, VKey<Registry>> keysMap =
-                      toMap(
+                      toMap(ImmutableSet.copyOf(tlds), Registry::createVKey);
-                          ImmutableSet.copyOf(tlds),
+                  Map<VKey<? extends Registry>, Registry> entities =
-                          tld -> Key.create(getCrossTldKey(), Registry.class, tld));
+                      tm().doTransactionless(() -> tm().load(keysMap.values()));
-                  Map<Key<Registry>, Registry> entities =
-                      tm().doTransactionless(() -> ofy().load().keys(keysMap.values()));
                   return Maps.transformEntries(
                       keysMap, (k, v) -> Optional.ofNullable(entities.getOrDefault(v, null)));
                 }
               });
 
+  public static VKey<Registry> createVKey(String tld) {
+    return VKey.create(Registry.class, tld, Key.create(getCrossTldKey(), Registry.class, tld));
+  }
+
   /**
    * The name of the pricing engine that this TLD uses.
    *
@@ -386,7 +382,7 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
   CreateAutoTimestamp creationTime = CreateAutoTimestamp.create(null);
 
   /** The set of reserved lists that are applicable to this registry. */
-  @Column(name = "reserved_list_names", nullable = false)
+  @Column(name = "reserved_list_names")
   Set<Key<ReservedList>> reservedLists;
 
   /** Retrieves an ImmutableSet of all ReservedLists associated with this tld. */

core/src/main/java/google/registry/model/reporting/HistoryEntry.java

@@ -193,7 +193,7 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
    * transaction counts (such as contact or host mutations).
    */
   @Transient // domain-specific
-  Set<DomainTransactionRecord> domainTransactionRecords;
+  protected Set<DomainTransactionRecord> domainTransactionRecords;
 
   public long getId() {
     // For some reason, Hibernate throws NPE during some initialization phase if we don't deal with
@@ -273,7 +273,8 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
   /** This method exists solely to satisfy Hibernate. Use the {@link Builder} instead. */
   @SuppressWarnings("UnusedMethod")
   private void setDomainTransactionRecords(Set<DomainTransactionRecord> domainTransactionRecords) {
-    this.domainTransactionRecords = ImmutableSet.copyOf(domainTransactionRecords);
+    this.domainTransactionRecords =
+        domainTransactionRecords == null ? null : ImmutableSet.copyOf(domainTransactionRecords);
   }
 
   public static VKey<HistoryEntry> createVKey(Key<HistoryEntry> key) {

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

@@ -15,6 +15,7 @@
 package google.registry.persistence.transaction;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
@@ -25,6 +26,7 @@ import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 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.persistence.JpaRetries;
@@ -238,6 +240,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     entities.forEach(this::insert);
   }
 
+  @Override
+  public void insertWithoutBackup(Object entity) {
+    insert(entity);
+  }
+
+  @Override
+  public void insertAllWithoutBackup(ImmutableCollection<?> entities) {
+    insertAll(entities);
+  }
+
   @Override
   public void put(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
@@ -253,6 +265,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     entities.forEach(this::put);
   }
 
+  @Override
+  public void putWithoutBackup(Object entity) {
+    put(entity);
+  }
+
+  @Override
+  public void putAllWithoutBackup(ImmutableCollection<?> entities) {
+    putAll(entities);
+  }
+
   @Override
   public void update(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
@@ -269,6 +291,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     entities.forEach(this::update);
   }
 
+  @Override
+  public void updateWithoutBackup(Object entity) {
+    update(entity);
+  }
+
+  @Override
+  public void updateAllWithoutBackup(ImmutableCollection<?> entities) {
+    updateAll(entities);
+  }
+
   @Override
   public <T> boolean exists(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
@@ -315,6 +347,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     return result;
   }
 
+  @Override
+  public <T> T load(T entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    return (T)
+        load(VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
+  }
+
   @Override
   public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");
@@ -342,6 +382,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             .getResultList());
   }
 
+  @Override
+  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
+    return Streams.stream(entities).map(this::load).collect(toImmutableList());
+  }
+
   private int internalDelete(VKey<?> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
@@ -366,6 +411,37 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     vKeys.forEach(this::internalDelete);
   }
 
+  @Override
+  public void delete(Object entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    Object managedEntity = entity;
+    if (!getEntityManager().contains(entity)) {
+      managedEntity = getEntityManager().merge(entity);
+    }
+    getEntityManager().remove(managedEntity);
+  }
+
+  @Override
+  public void deleteWithoutBackup(VKey<?> key) {
+    delete(key);
+  }
+
+  @Override
+  public void deleteWithoutBackup(Iterable<? extends VKey<?>> keys) {
+    delete(keys);
+  }
+
+  @Override
+  public void deleteWithoutBackup(Object entity) {
+    delete(entity);
+  }
+
+  @Override
+  public void clearSessionCache() {
+    // This is an intended no-op method as there is no session cache in Postgresql.
+  }
+
   @Override
   public <T> void assertDelete(VKey<T> key) {
     if (internalDelete(key) != 1) {

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

@@ -91,18 +91,90 @@ public interface TransactionManager {
   /** Persists all new entities in the database, throws exception if any entity already exists. */
   void insertAll(ImmutableCollection<?> entities);
 
+  /**
+   * Persists a new entity in the database without writing any backup if the underlying database is
+   * Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void insertWithoutBackup(Object entity);
+
+  /**
+   * Persists all new entities in the database without writing any backup if the underlying database
+   * is Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void insertAllWithoutBackup(ImmutableCollection<?> entities);
+
   /** Persists a new entity or update the existing entity in the database. */
   void put(Object entity);
 
   /** Persists all new entities or update the existing entities in the database. */
   void putAll(ImmutableCollection<?> entities);
 
+  /**
+   * Persists a new entity or update the existing entity in the database without writing any backup
+   * if the underlying database is Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void putWithoutBackup(Object entity);
+
+  /**
+   * Persists all new entities or update the existing entities in the database without writing any
+   * backup if the underlying database is Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void putAllWithoutBackup(ImmutableCollection<?> entities);
+
   /** Updates an entity in the database, throws exception if the entity does not exist. */
   void update(Object entity);
 
   /** Updates all entities in the database, throws exception if any entity does not exist. */
   void updateAll(ImmutableCollection<?> entities);
 
+  /**
+   * Updates an entity in the database without writing any backup if the underlying database is
+   * Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void updateWithoutBackup(Object entity);
+
+  /**
+   * Updates all entities in the database without writing any backup if the underlying database is
+   * Datastore.
+   *
+   * <p>This method is for the sake of keeping a single code path when replacing ofy() with tm() in
+   * the application code. When the method is invoked with Datastore, it won't write the commit log
+   * backup; when invoked with Cloud SQL, it behaves the same as the method which doesn't have
+   * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
+   * SQL.
+   */
+  void updateAllWithoutBackup(ImmutableCollection<?> entities);
+
   /** Returns whether the given entity with same ID exists. */
   boolean exists(Object entity);
 
@@ -115,6 +187,11 @@ public interface TransactionManager {
   /** Loads the entity by its id, throws NoSuchElementException if it doesn't exist. */
   <T> T load(VKey<T> key);
 
+  /**
+   * Loads the given entity from the database, throws NoSuchElementException if it doesn't exist.
+   */
+  <T> T load(T entity);
+
   /**
    * Loads the set of entities by their key id.
    *
@@ -125,9 +202,38 @@ public interface TransactionManager {
   /** Loads all entities of the given type, returns empty if there is no such entity. */
   <T> ImmutableList<T> loadAll(Class<T> clazz);
 
+  /**
+   * Loads all given entities from the database, throws NoSuchElementException if it doesn't exist.
+   */
+  <T> ImmutableList<T> loadAll(Iterable<T> entities);
+
   /** Deletes the entity by its id. */
   void delete(VKey<?> key);
 
   /** Deletes the set of entities by their key id. */
   void delete(Iterable<? extends VKey<?>> keys);
+
+  /** Deletes the given entity from the database. */
+  void delete(Object entity);
+
+  /**
+   * Deletes the entity by its id without writing any backup if the underlying database is
+   * Datastore.
+   */
+  void deleteWithoutBackup(VKey<?> key);
+
+  /**
+   * Deletes the set of entities by their key id without writing any backup if the underlying
+   * database is Datastore.
+   */
+  void deleteWithoutBackup(Iterable<? extends VKey<?>> keys);
+
+  /**
+   * Deletes the given entity from the database without writing any backup if the underlying
+   * database is Datastore.
+   */
+  void deleteWithoutBackup(Object entity);
+
+  /** Clears the session cache if the underlying database is Datastore, otherwise it is a no-op. */
+  void clearSessionCache();
 }

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

@@ -0,0 +1,110 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.transaction;
+
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import google.registry.model.ofy.DatastoreTransactionManager;
+import java.util.function.Supplier;
+
+/** Utility class that provides supplementary methods for {@link TransactionManager}. */
+public class TransactionManagerUtil {
+
+  /**
+   * Returns the result of the given {@link Supplier}.
+   *
+   * <p>If {@link TransactionManagerFactory#tm()} returns a {@link JpaTransactionManager} instance,
+   * the {@link Supplier} is executed in a transaction.
+   */
+  public static <T> T transactIfJpaTm(Supplier<T> supplier) {
+    if (tm() instanceof JpaTransactionManager) {
+      return tm().transact(supplier);
+    } else {
+      return supplier.get();
+    }
+  }
+
+  /**
+   * Executes the given {@link Runnable}.
+   *
+   * <p>If {@link TransactionManagerFactory#tm()} returns a {@link JpaTransactionManager} instance,
+   * the {@link Runnable} is executed in a transaction.
+   */
+  public static void transactIfJpaTm(Runnable runnable) {
+    transactIfJpaTm(
+        () -> {
+          runnable.run();
+          return null;
+        });
+  }
+
+  /**
+   * Executes either {@code ofyRunnable} if {@link TransactionManagerFactory#tm()} returns a {@link
+   * JpaTransactionManager} instance, or {@code jpaRunnable} if {@link
+   * TransactionManagerFactory#tm()} returns a {@link DatastoreTransactionManager} instance.
+   */
+  public static void ofyOrJpaTm(Runnable ofyRunnable, Runnable jpaRunnable) {
+    ofyOrJpaTm(
+        () -> {
+          ofyRunnable.run();
+          return null;
+        },
+        () -> {
+          jpaRunnable.run();
+          return null;
+        });
+  }
+
+  /**
+   * Returns the result from either {@code ofySupplier} if {@link TransactionManagerFactory#tm()}
+   * returns a {@link JpaTransactionManager} instance, or {@code jpaSupplier} if {@link
+   * TransactionManagerFactory#tm()} returns a {@link DatastoreTransactionManager} instance.
+   */
+  public static <T> T ofyOrJpaTm(Supplier<T> ofySupplier, Supplier<T> jpaSupplier) {
+    if (tm() instanceof DatastoreTransactionManager) {
+      return ofySupplier.get();
+    } else if (tm() instanceof JpaTransactionManager) {
+      return jpaSupplier.get();
+    } else {
+      throw new IllegalStateException(
+          "Expected tm() to be DatastoreTransactionManager or JpaTransactionManager, but got "
+              + tm().getClass());
+    }
+  }
+
+  /**
+   * Executes the given {@link Runnable} if {@link TransactionManagerFactory#tm()} returns a {@link
+   * DatastoreTransactionManager} instance, otherwise does nothing.
+   */
+  public static void ofyTmOrDoNothing(Runnable ofyRunnable) {
+    if (tm() instanceof DatastoreTransactionManager) {
+      ofyRunnable.run();
+    }
+  }
+
+  /**
+   * Returns the result from the given {@link Supplier} if {@link TransactionManagerFactory#tm()}
+   * returns a {@link DatastoreTransactionManager} instance, otherwise returns null.
+   */
+  public static <T> T ofyTmOrDoNothing(Supplier<T> ofySupplier) {
+    if (tm() instanceof DatastoreTransactionManager) {
+      return ofySupplier.get();
+    } else {
+      return null;
+    }
+  }
+
+  private TransactionManagerUtil() {}
+}

core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -480,8 +480,7 @@ public class ExpandRecurringBillingEventsActionTest
           .setSyntheticCreationTime(testTime)
           .build());
     }
-    assertBillingEventsForResource(
+    assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class));
-        domain, Iterables.toArray(expectedEvents, BillingEvent.class));
     assertCursorAt(testTime);
   }
 

core/src/test/java/google/registry/model/billing/BillingEventTest.java

@@ -17,11 +17,12 @@ package google.registry.model.billing;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
 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.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistActiveDomain;
 import static google.registry.testing.DatastoreHelper.persistResource;
-import static google.registry.testing.SqlHelper.saveRegistrar;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static org.joda.money.CurrencyUnit.USD;
 import static org.joda.time.DateTimeZone.UTC;
@@ -39,13 +40,17 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import google.registry.util.DateTimeUtils;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link BillingEvent}. */
+@DualDatabaseTest
 public class BillingEventTest extends EntityTestCase {
   private final DateTime now = DateTime.now(UTC);
 
@@ -67,16 +72,26 @@ public class BillingEventTest extends EntityTestCase {
   void setUp() {
     createTld("tld");
     domain = persistActiveDomain("foo.tld");
-    historyEntry = persistResource(
+    historyEntry =
-        new HistoryEntry.Builder()
+        persistResource(
-            .setParent(domain)
+            new HistoryEntry.Builder()
-            .setModificationTime(now)
+                .setParent(domain)
-            .build());
+                .setModificationTime(now)
-    historyEntry2 = persistResource(
+                .setRequestedByRegistrar(false)
-        new HistoryEntry.Builder()
+                .setType(HistoryEntry.Type.DOMAIN_CREATE)
-        .setParent(domain)
+                .setXmlBytes(new byte[0])
-        .setModificationTime(now.plusDays(1))
+                .build()
-        .build());
+                .toChildHistoryEntity());
+    historyEntry2 =
+        persistResource(
+            new HistoryEntry.Builder()
+                .setParent(domain)
+                .setModificationTime(now.plusDays(1))
+                .setRequestedByRegistrar(false)
+                .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                .setXmlBytes(new byte[0])
+                .build()
+                .toChildHistoryEntity());
 
     AllocationToken allocationToken =
         persistResource(
@@ -149,74 +164,38 @@ public class BillingEventTest extends EntityTestCase {
                     .setEventTime(now.plusDays(1))
                     .setBillingTime(now.plusYears(1).plusDays(45))
                     .setRecurringEventKey(recurring.createVKey())));
-    modification = persistResource(commonInit(
+    modification =
-        new BillingEvent.Modification.Builder()
+        ofyTmOrDoNothing(
-            .setParent(historyEntry2)
+            () ->
-            .setReason(Reason.CREATE)
+                persistResource(
-            .setCost(Money.of(USD, 1))
+                    commonInit(
-            .setDescription("Something happened")
+                        new BillingEvent.Modification.Builder()
-            .setEventTime(now.plusDays(1))
+                            .setParent(historyEntry2)
-            .setEventKey(Key.create(oneTime))));
+                            .setReason(Reason.CREATE)
+                            .setCost(Money.of(USD, 1))
+                            .setDescription("Something happened")
+                            .setEventTime(now.plusDays(1))
+                            .setEventKey(Key.create(oneTime)))));
   }
 
   private <E extends BillingEvent, B extends BillingEvent.Builder<E, B>> E commonInit(B builder) {
-    return builder
+    return builder.setClientId("TheRegistrar").setTargetId("foo.tld").build();
-        .setClientId("a registrar")
-        .setTargetId("foo.tld")
-        .build();
   }
 
-  private void saveNewBillingEvent(BillingEvent billingEvent) {
+  @TestOfyAndSql
-    jpaTm().transact(() -> jpaTm().insert(billingEvent));
-  }
-
-  @Test
-  void testCloudSqlPersistence_OneTime() {
-    saveRegistrar("a registrar");
-    saveNewBillingEvent(oneTime);
-
-    BillingEvent.OneTime persisted = jpaTm().transact(() -> jpaTm().load(oneTime.createVKey()));
-    // TODO(b/168325240): Remove this fix after VKeyConverter generates symmetric key for
-    // AllocationToken.
-    BillingEvent.OneTime fixed =
-        persisted.asBuilder().setAllocationToken(oneTime.getAllocationToken().get()).build();
-    assertThat(fixed).isEqualTo(oneTime);
-  }
-
-  @Test
-  void testCloudSqlPersistence_Cancellation() {
-    saveRegistrar("a registrar");
-    saveNewBillingEvent(oneTime);
-    saveNewBillingEvent(cancellationOneTime);
-
-    BillingEvent.Cancellation persisted =
-        jpaTm().transact(() -> jpaTm().load(cancellationOneTime.createVKey()));
-    // TODO(b/168537779): Remove this fix after VKey<OneTime> can be reconstructed correctly.
-    BillingEvent.Cancellation fixed =
-        persisted.asBuilder().setOneTimeEventKey(oneTime.createVKey()).build();
-    assertThat(fixed).isEqualTo(cancellationOneTime);
-  }
-
-  @Test
-  void testCloudSqlPersistence_Recurring() {
-    saveRegistrar("a registrar");
-    saveNewBillingEvent(recurring);
-
-    BillingEvent.Recurring persisted = jpaTm().transact(() -> jpaTm().load(recurring.createVKey()));
-    assertThat(persisted).isEqualTo(recurring);
-  }
-
-  @Test
   void testPersistence() {
-    assertThat(ofy().load().entity(oneTime).now()).isEqualTo(oneTime);
+    assertThat(transactIfJpaTm(() -> tm().load(oneTime))).isEqualTo(oneTime);
-    assertThat(ofy().load().entity(oneTimeSynthetic).now()).isEqualTo(oneTimeSynthetic);
+    assertThat(transactIfJpaTm(() -> tm().load(oneTimeSynthetic))).isEqualTo(oneTimeSynthetic);
-    assertThat(ofy().load().entity(recurring).now()).isEqualTo(recurring);
+    assertThat(transactIfJpaTm(() -> tm().load(recurring))).isEqualTo(recurring);
-    assertThat(ofy().load().entity(cancellationOneTime).now()).isEqualTo(cancellationOneTime);
+    assertThat(transactIfJpaTm(() -> tm().load(cancellationOneTime)))
-    assertThat(ofy().load().entity(cancellationRecurring).now()).isEqualTo(cancellationRecurring);
+        .isEqualTo(cancellationOneTime);
-    assertThat(ofy().load().entity(modification).now()).isEqualTo(modification);
+    assertThat(transactIfJpaTm(() -> tm().load(cancellationRecurring)))
+        .isEqualTo(cancellationRecurring);
+
+    ofyTmOrDoNothing(() -> assertThat(tm().load(modification)).isEqualTo(modification));
   }
 
-  @Test
+  @TestOfyOnly
   void testParenting() {
     // Note that these are all tested separately because BillingEvent is an abstract base class that
     // lacks the @Entity annotation, and thus we cannot call .type(BillingEvent.class)
@@ -238,19 +217,14 @@ public class BillingEventTest extends EntityTestCase {
         .containsExactly(modification);
   }
 
-  @Test
+  @TestOfyAndSql
   void testCancellationMatching() {
-    Key<?> recurringKey =
+    VKey<?> recurringKey =
-        ofy()
+        transactIfJpaTm(() -> tm().load(oneTimeSynthetic).getCancellationMatchingBillingEvent());
-            .load()
+    assertThat(transactIfJpaTm(() -> tm().load(recurringKey))).isEqualTo(recurring);
-            .entity(oneTimeSynthetic)
-            .now()
-            .getCancellationMatchingBillingEvent()
-            .getOfyKey();
-    assertThat(ofy().load().key(recurringKey).now()).isEqualTo(recurring);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
     verifyIndexing(
         oneTime,
@@ -272,7 +246,7 @@ public class BillingEventTest extends EntityTestCase {
     verifyIndexing(modification, "clientId", "eventTime");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_syntheticFlagWithoutCreationTime() {
     IllegalStateException thrown =
         assertThrows(
@@ -288,7 +262,7 @@ public class BillingEventTest extends EntityTestCase {
         .contains("Synthetic creation time must be set if and only if the SYNTHETIC flag is set.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_syntheticCreationTimeWithoutFlag() {
     IllegalStateException thrown =
         assertThrows(
@@ -299,7 +273,7 @@ public class BillingEventTest extends EntityTestCase {
         .contains("Synthetic creation time must be set if and only if the SYNTHETIC flag is set");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_syntheticFlagWithoutCancellationMatchingKey() {
     IllegalStateException thrown =
         assertThrows(
@@ -317,7 +291,7 @@ public class BillingEventTest extends EntityTestCase {
                 + "if and only if the SYNTHETIC flag is set");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cancellationMatchingKeyWithoutFlag() {
     IllegalStateException thrown =
         assertThrows(
@@ -334,7 +308,7 @@ public class BillingEventTest extends EntityTestCase {
                 + "if and only if the SYNTHETIC flag is set");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_cancellation_forGracePeriod_withOneTime() {
     BillingEvent.Cancellation newCancellation =
         BillingEvent.Cancellation.forGracePeriod(
@@ -343,10 +317,15 @@ public class BillingEventTest extends EntityTestCase {
             "foo.tld");
     // Set ID to be the same to ignore for the purposes of comparison.
     newCancellation = newCancellation.asBuilder().setId(cancellationOneTime.getId()).build();
-    assertThat(newCancellation).isEqualTo(cancellationOneTime);
+
+    // TODO(b/168537779): Remove setRecurringEventKey after symmetric VKey can be reconstructed
+    // correctly.
+    assertThat(newCancellation)
+        .isEqualTo(
+            cancellationOneTime.asBuilder().setOneTimeEventKey(oneTime.createVKey()).build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_cancellation_forGracePeriod_withRecurring() {
     BillingEvent.Cancellation newCancellation =
         BillingEvent.Cancellation.forGracePeriod(
@@ -354,16 +333,21 @@ public class BillingEventTest extends EntityTestCase {
                 GracePeriodStatus.AUTO_RENEW,
                 domain.getRepoId(),
                 now.plusYears(1).plusDays(45),
-                "a registrar",
+                "TheRegistrar",
                 recurring.createVKey()),
             historyEntry2,
             "foo.tld");
     // Set ID to be the same to ignore for the purposes of comparison.
     newCancellation = newCancellation.asBuilder().setId(cancellationRecurring.getId()).build();
-    assertThat(newCancellation).isEqualTo(cancellationRecurring);
+
+    // TODO(b/168537779): Remove setRecurringEventKey after symmetric VKey can be reconstructed
+    // correctly.
+    assertThat(newCancellation)
+        .isEqualTo(
+            cancellationRecurring.asBuilder().setRecurringEventKey(recurring.createVKey()).build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cancellation_forGracePeriodWithoutBillingEvent() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -380,7 +364,7 @@ public class BillingEventTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("grace period without billing event");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cancellationWithNoBillingEvent() {
     IllegalStateException thrown =
         assertThrows(
@@ -394,7 +378,7 @@ public class BillingEventTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("exactly one billing event");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cancellationWithBothBillingEvents() {
     IllegalStateException thrown =
         assertThrows(
@@ -408,7 +392,7 @@ public class BillingEventTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("exactly one billing event");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDeadCodeThatDeletedScrapCommandsReference() {
     assertThat(recurring.getParentKey()).isEqualTo(Key.create(historyEntry));
     new BillingEvent.OneTime.Builder().setParent(Key.create(historyEntry));

core/src/test/java/google/registry/model/rde/RdeRevisionTest.java

@@ -23,9 +23,9 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import google.registry.model.EntityTestCase;
 import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.TestTemplate;
 
 /** Unit tests for {@link RdeRevision}. */
 @DualDatabaseTest
@@ -40,20 +40,20 @@ public class RdeRevisionTest extends EntityTestCase {
     fakeClock.setTo(DateTime.parse("1984-12-18TZ"));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testGetNextRevision_objectDoesntExist_returnsZero() {
     tm().transact(
             () -> assertThat(getNextRevision("torment", fakeClock.nowUtc(), FULL)).isEqualTo(0));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testGetNextRevision_objectExistsAtZero_returnsOne() {
     save("sorrow", fakeClock.nowUtc(), FULL, 0);
     tm().transact(
             () -> assertThat(getNextRevision("sorrow", fakeClock.nowUtc(), FULL)).isEqualTo(1));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_objectDoesntExist_newRevisionIsZero_nextRevIsOne() {
     tm().transact(() -> saveRevision("despondency", fakeClock.nowUtc(), FULL, 0));
     tm().transact(
@@ -61,7 +61,7 @@ public class RdeRevisionTest extends EntityTestCase {
                 assertThat(getNextRevision("despondency", fakeClock.nowUtc(), FULL)).isEqualTo(1));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_objectDoesntExist_newRevisionIsOne_throwsVe() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -74,7 +74,7 @@ public class RdeRevisionTest extends EntityTestCase {
                 + "when trying to save new revision 1");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_objectExistsAtZero_newRevisionIsZero_throwsVe() {
     save("melancholy", fakeClock.nowUtc(), FULL, 0);
     IllegalArgumentException thrown =
@@ -84,7 +84,7 @@ public class RdeRevisionTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("object already created");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_objectExistsAtZero_newRevisionIsOne_nextRevIsTwo() {
     DateTime startOfDay = fakeClock.nowUtc().withTimeAtStartOfDay();
     save("melancholy", startOfDay, FULL, 0);
@@ -93,7 +93,7 @@ public class RdeRevisionTest extends EntityTestCase {
     tm().transact(() -> assertThat(getNextRevision("melancholy", startOfDay, FULL)).isEqualTo(2));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_objectExistsAtZero_newRevisionIsTwo_throwsVe() {
     save("melancholy", fakeClock.nowUtc(), FULL, 0);
     IllegalArgumentException thrown =
@@ -105,7 +105,7 @@ public class RdeRevisionTest extends EntityTestCase {
         .contains("RDE revision object should be at revision 1 but was");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_negativeRevision_throwsIae() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -114,7 +114,7 @@ public class RdeRevisionTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("Negative revision");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void testSaveRevision_callerNotInTransaction_throwsIse() {
     IllegalStateException thrown =
         assertThrows(

core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java

@@ -33,6 +33,8 @@ 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 java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Set;
@@ -40,7 +42,6 @@ import java.util.stream.Stream;
 import javax.persistence.Embeddable;
 import javax.persistence.MappedSuperclass;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /**
@@ -78,28 +79,28 @@ public class TransactionManagerTest {
     fakeClock.advanceOneMilli();
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void inTransaction_returnsCorrespondingResult() {
     assertThat(tm().inTransaction()).isFalse();
     tm().transact(() -> assertThat(tm().inTransaction()).isTrue());
     assertThat(tm().inTransaction()).isFalse();
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void assertInTransaction_throwsExceptionWhenNotInTransaction() {
     assertThrows(IllegalStateException.class, () -> tm().assertInTransaction());
     tm().transact(() -> tm().assertInTransaction());
     assertThrows(IllegalStateException.class, () -> tm().assertInTransaction());
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void getTransactionTime_throwsExceptionWhenNotInTransaction() {
     assertThrows(IllegalStateException.class, () -> tm().getTransactionTime());
     tm().transact(() -> assertThat(tm().getTransactionTime()).isEqualTo(fakeClock.nowUtc()));
     assertThrows(IllegalStateException.class, () -> tm().getTransactionTime());
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void transact_hasNoEffectWithPartialSuccess() {
     assertEntityNotExist(theEntity);
     assertThrows(
@@ -113,21 +114,21 @@ public class TransactionManagerTest {
     assertEntityNotExist(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void transact_reusesExistingTransaction() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().transact(() -> tm().insert(theEntity)));
     assertEntityExists(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void transactNew_succeeds() {
     assertEntityNotExist(theEntity);
     tm().transactNew(() -> tm().insert(theEntity));
     assertEntityExists(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void transactNewReadOnly_succeeds() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().insert(theEntity));
@@ -136,7 +137,7 @@ public class TransactionManagerTest {
     assertThat(persisted).isEqualTo(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void transactNewReadOnly_throwsWhenWritingEntity() {
     assertEntityNotExist(theEntity);
     assertThrows(
@@ -144,7 +145,7 @@ public class TransactionManagerTest {
     assertEntityNotExist(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void saveNew_succeeds() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().insert(theEntity));
@@ -152,14 +153,14 @@ public class TransactionManagerTest {
     assertThat(tm().transact(() -> tm().load(theEntity.key()))).isEqualTo(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void saveAllNew_succeeds() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
     assertAllEntitiesExist(moreEntities);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void saveNewOrUpdate_persistsNewEntity() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().put(theEntity));
@@ -167,7 +168,7 @@ public class TransactionManagerTest {
     assertThat(tm().transact(() -> tm().load(theEntity.key()))).isEqualTo(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void saveNewOrUpdate_updatesExistingEntity() {
     tm().transact(() -> tm().insert(theEntity));
     TestEntity persisted = tm().transact(() -> tm().load(theEntity.key()));
@@ -179,14 +180,14 @@ public class TransactionManagerTest {
     assertThat(persisted.data).isEqualTo("bar");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void saveNewOrUpdateAll_succeeds() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().putAll(moreEntities));
     assertAllEntitiesExist(moreEntities);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void update_succeeds() {
     tm().transact(() -> tm().insert(theEntity));
     TestEntity persisted =
@@ -202,7 +203,7 @@ public class TransactionManagerTest {
     assertThat(persisted.data).isEqualTo("bar");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void load_succeeds() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().insert(theEntity));
@@ -211,14 +212,14 @@ public class TransactionManagerTest {
     assertThat(persisted.data).isEqualTo("foo");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void load_throwsOnMissingElement() {
     assertEntityNotExist(theEntity);
     assertThrows(
         NoSuchElementException.class, () -> tm().transact(() -> tm().load(theEntity.key())));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void maybeLoad_succeeds() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().insert(theEntity));
@@ -227,13 +228,13 @@ public class TransactionManagerTest {
     assertThat(persisted.data).isEqualTo("foo");
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void maybeLoad_nonExistentObject() {
     assertEntityNotExist(theEntity);
     assertThat(tm().transact(() -> tm().maybeLoad(theEntity.key())).isPresent()).isFalse();
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void delete_succeeds() {
     tm().transact(() -> tm().insert(theEntity));
     assertEntityExists(theEntity);
@@ -242,14 +243,14 @@ public class TransactionManagerTest {
     assertEntityNotExist(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void delete_doNothingWhenEntityNotExist() {
     assertEntityNotExist(theEntity);
     tm().transact(() -> tm().delete(theEntity.key()));
     assertEntityNotExist(theEntity);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void delete_succeedsForEntitySet() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
@@ -261,7 +262,7 @@ public class TransactionManagerTest {
     assertAllEntitiesNotExist(moreEntities);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void delete_ignoreNonExistentEntity() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
@@ -276,7 +277,16 @@ public class TransactionManagerTest {
     assertAllEntitiesNotExist(moreEntities);
   }
 
-  @TestTemplate
+  @TestOfyAndSql
+  void delete_deletesTheGivenEntity() {
+    tm().transact(() -> tm().insert(theEntity));
+    assertEntityExists(theEntity);
+    fakeClock.advanceOneMilli();
+    tm().transact(() -> tm().delete(theEntity));
+    assertEntityNotExist(theEntity);
+  }
+
+  @TestOfyAndSql
   void load_multi() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
@@ -286,7 +296,7 @@ public class TransactionManagerTest {
         .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void load_multiWithDuplicateKeys() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
@@ -298,7 +308,7 @@ public class TransactionManagerTest {
         .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
   }
 
-  @TestTemplate
+  @TestOfyAndSql
   void load_multiMissingKeys() {
     assertAllEntitiesNotExist(moreEntities);
     tm().transact(() -> tm().insertAll(moreEntities));
@@ -310,6 +320,14 @@ public class TransactionManagerTest {
         .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
   }
 
+  @TestOfyOnly
+  void loadAllForOfyTm_throwsExceptionInTransaction() {
+    assertAllEntitiesNotExist(moreEntities);
+    tm().transact(() -> tm().insertAll(moreEntities));
+    assertThrows(
+        IllegalArgumentException.class, () -> tm().transact(() -> tm().loadAll(TestEntity.class)));
+  }
+
   private static void assertEntityExists(TestEntity entity) {
     assertThat(tm().transact(() -> tm().exists(entity))).isTrue();
   }

core/src/test/java/google/registry/server/Fixture.java

@@ -61,7 +61,7 @@ public enum Fixture {
       createTlds("xn--q9jyb4c", "example");
 
       // Used for OT&E TLDs
-      persistPremiumList("default_sandbox_list");
+      persistPremiumList("default_sandbox_list", "sandbox,USD 1000");
 
       try {
         OteStatsTestHelper.setupCompleteOte("otefinished");

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

@@ -49,7 +49,7 @@ abstract class AbstractEppResourceSubject<
     this.actual = subject;
   }
 
-  private List<HistoryEntry> getHistoryEntries() {
+  private List<? extends HistoryEntry> getHistoryEntries() {
     return DatastoreHelper.getHistoryEntries(actual);
   }
 
@@ -120,7 +120,7 @@ abstract class AbstractEppResourceSubject<
   // TODO(weiminyu): Remove after next Truth update
   @SuppressWarnings("UnnecessaryParentheses")
   public Which<HistoryEntrySubject> hasHistoryEntryAtIndex(int index) {
-    List<HistoryEntry> historyEntries = getHistoryEntries();
+    List<? extends HistoryEntry> historyEntries = getHistoryEntries();
     check("getHistoryEntries().size()").that(historyEntries.size()).isAtLeast(index + 1);
     return new Which<>(
         check("getHistoryEntries(%s)", index)

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

@@ -17,6 +17,7 @@ package google.registry.testing;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.io.Files.asCharSink;
 import static com.google.common.truth.Truth.assertWithMessage;
+import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
 import static google.registry.testing.DatastoreHelper.persistSimpleResources;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.injectTmForDualDatabaseTest;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.restoreTmAfterDualDatabaseTest;
@@ -384,6 +385,17 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (isWithDatastoreAndCloudSql()) {
       injectTmForDualDatabaseTest(context);
     }
+    ofyOrJpaTm(
+        () -> {
+          if (withDatastore && !withoutCannedData) {
+            loadInitialData();
+          }
+        },
+        () -> {
+          if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
+            loadInitialData();
+          }
+        });
   }
 
   /**
@@ -458,9 +470,6 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
       ObjectifyService.resetNextTestId();
-      if (!withoutCannedData) {
-        loadInitialData();
-      }
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }

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

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Suppliers.memoize;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.Iterables.toArray;
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static com.google.common.truth.Truth.assertThat;
@@ -33,6 +34,9 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.label.PremiumListUtils.parentPremiumListEntriesOnRevision;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.difference;
 import static google.registry.util.CollectionUtils.union;
@@ -44,6 +48,7 @@ import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 import static java.util.Arrays.asList;
 import static org.joda.money.CurrencyUnit.USD;
+import static org.junit.jupiter.api.Assertions.fail;
 
 import com.google.common.base.Ascii;
 import com.google.common.base.Splitter;
@@ -60,22 +65,27 @@ import google.registry.dns.writer.VoidDnsWriter;
 import google.registry.model.Buildable;
 import google.registry.model.EppResource;
 import google.registry.model.EppResource.ForeignKeyedEppResource;
-import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.contact.ContactAuthInfo;
+import google.registry.model.contact.ContactBase;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DesignatedContact;
 import google.registry.model.domain.DesignatedContact.Type;
 import google.registry.model.domain.DomainAuthInfo;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostBase;
+import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.EppResourceIndexBucket;
@@ -101,10 +111,14 @@ import google.registry.persistence.VKey;
 import google.registry.tmch.LordnTaskUtils;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
+import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
+import org.joda.time.DateTimeComparator;
+import org.joda.time.DateTimeZone;
 
 /** Static utils for setting up test resources. */
 public class DatastoreHelper {
@@ -325,18 +339,40 @@ public class DatastoreHelper {
    * the requirement to have monotonically increasing timestamps.
    */
   public static PremiumList persistPremiumList(String listName, String... lines) {
-    PremiumList premiumList = new PremiumList.Builder().setName(listName).build();
+    checkState(lines.length != 0, "Must provide at least one premium entry");
-    ImmutableMap<String, PremiumListEntry> entries = premiumList.parse(asList(lines));
+    PremiumList partialPremiumList = new PremiumList.Builder().setName(listName).build();
+    ImmutableMap<String, PremiumListEntry> entries = partialPremiumList.parse(asList(lines));
+    CurrencyUnit currencyUnit =
+        entries.entrySet().iterator().next().getValue().getValue().getCurrencyUnit();
+    PremiumList premiumList =
+        partialPremiumList
+            .asBuilder()
+            .setCreationTime(DateTime.now(DateTimeZone.UTC))
+            .setCurrency(currencyUnit)
+            .setLabelsToPrices(
+                entries.entrySet().stream()
+                    .collect(
+                        toImmutableMap(
+                            Map.Entry::getKey, entry -> entry.getValue().getValue().getAmount())))
+            .build();
     PremiumListRevision revision = PremiumListRevision.create(premiumList, entries.keySet());
-    ofy()
+
-        .saveWithoutBackup()
+    ofyOrJpaTm(
-        .entities(premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision)
+        () -> {
-        .now();
+          tm().putAllWithoutBackup(
-    ofy()
+                  ImmutableList.of(
-        .saveWithoutBackup()
+                      premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision));
-        .entities(parentPremiumListEntriesOnRevision(entries.values(), Key.create(revision)))
+          tm().putAllWithoutBackup(
-        .now();
+                  parentPremiumListEntriesOnRevision(entries.values(), Key.create(revision)));
-    return ofy().load().entity(premiumList).now();
+        },
+        () -> tm().transact(() -> tm().insert(premiumList)));
+    // The above premiumList is in the session cache and it is different from the corresponding
+    // entity stored in Datastore because it has some @Ignore fields set dedicated for SQL. This
+    // breaks the assumption we have in our application code, see
+    // PremiumListUtils.savePremiumListAndEntries(). Clearing the session cache can help make sure
+    // we always get the same list.
+    tm().clearSessionCache();
+    return transactIfJpaTm(() -> tm().load(premiumList));
   }
 
   /** Creates and persists a tld. */
@@ -672,17 +708,28 @@ public class DatastoreHelper {
   }
 
   private static Iterable<BillingEvent> getBillingEvents() {
-    return Iterables.concat(
+    return transactIfJpaTm(
-        ofy().load().type(BillingEvent.OneTime.class),
+        () ->
-        ofy().load().type(BillingEvent.Recurring.class),
+            Iterables.concat(
-        ofy().load().type(BillingEvent.Cancellation.class));
+                tm().loadAll(BillingEvent.OneTime.class),
+                tm().loadAll(BillingEvent.Recurring.class),
+                tm().loadAll(BillingEvent.Cancellation.class)));
   }
 
   private static Iterable<BillingEvent> getBillingEvents(EppResource resource) {
-    return Iterables.concat(
+    return transactIfJpaTm(
-        ofy().load().type(BillingEvent.OneTime.class).ancestor(resource),
+        () ->
-        ofy().load().type(BillingEvent.Recurring.class).ancestor(resource),
+            Iterables.concat(
-        ofy().load().type(BillingEvent.Cancellation.class).ancestor(resource));
+                tm().loadAll(BillingEvent.OneTime.class).stream()
+                    .filter(oneTime -> oneTime.getDomainRepoId().equals(resource.getRepoId()))
+                    .collect(toImmutableList()),
+                tm().loadAll(BillingEvent.Recurring.class).stream()
+                    .filter(recurring -> recurring.getDomainRepoId().equals(resource.getRepoId()))
+                    .collect(toImmutableList()),
+                tm().loadAll(BillingEvent.Cancellation.class).stream()
+                    .filter(
+                        cancellation -> cancellation.getDomainRepoId().equals(resource.getRepoId()))
+                    .collect(toImmutableList())));
   }
 
   /** Assert that the actual billing event matches the expected one, ignoring IDs. */
@@ -751,41 +798,63 @@ public class DatastoreHelper {
     assertPollMessagesEqual(getPollMessages(), Arrays.asList(expected));
   }
 
-  public static void assertPollMessagesForResource(EppResource resource, PollMessage... expected) {
+  public static void assertPollMessagesForResource(DomainContent domain, PollMessage... expected) {
-    assertPollMessagesEqual(getPollMessages(resource), Arrays.asList(expected));
+    assertPollMessagesEqual(getPollMessages(domain), Arrays.asList(expected));
   }
 
   public static ImmutableList<PollMessage> getPollMessages() {
-    return ImmutableList.copyOf(ofy().load().type(PollMessage.class));
+    return ImmutableList.copyOf(transactIfJpaTm(() -> tm().loadAll(PollMessage.class)));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId) {
-    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).filter("clientId", clientId));
+    return transactIfJpaTm(
+        () ->
+            tm().loadAll(PollMessage.class).stream()
+                .filter(pollMessage -> pollMessage.getClientId().equals(clientId))
+                .collect(toImmutableList()));
   }
 
-  public static ImmutableList<PollMessage> getPollMessages(EppResource resource) {
+  public static ImmutableList<PollMessage> getPollMessages(DomainContent domain) {
-    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).ancestor(resource));
+    return transactIfJpaTm(
+        () ->
+            tm().loadAll(PollMessage.class).stream()
+                .filter(
+                    pollMessage ->
+                        pollMessage.getParentKey().getParent().getName().equals(domain.getRepoId()))
+                .collect(toImmutableList()));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId, DateTime now) {
-    return ImmutableList.copyOf(
+    return transactIfJpaTm(
-        ofy()
+        () ->
-            .load()
+            tm().loadAll(PollMessage.class).stream()
-            .type(PollMessage.class)
+                .filter(pollMessage -> pollMessage.getClientId().equals(clientId))
-            .filter("clientId", clientId)
+                .filter(
-            .filter("eventTime <=", now.toDate()));
+                    pollMessage ->
+                        pollMessage.getEventTime().isEqual(now)
+                            || pollMessage.getEventTime().isBefore(now))
+                .collect(toImmutableList()));
   }
 
   /** Gets all PollMessages associated with the given EppResource. */
   public static ImmutableList<PollMessage> getPollMessages(
       EppResource resource, String clientId, DateTime now) {
-    return ImmutableList.copyOf(
+    return transactIfJpaTm(
-        ofy()
+        () ->
-            .load()
+            tm().loadAll(PollMessage.class).stream()
-            .type(PollMessage.class)
+                .filter(
-            .ancestor(resource)
+                    pollMessage ->
-            .filter("clientId", clientId)
+                        pollMessage
-            .filter("eventTime <=", now.toDate()));
+                            .getParentKey()
+                            .getParent()
+                            .getName()
+                            .equals(resource.getRepoId()))
+                .filter(pollMessage -> pollMessage.getClientId().equals(clientId))
+                .filter(
+                    pollMessage ->
+                        pollMessage.getEventTime().isEqual(now)
+                            || pollMessage.getEventTime().isBefore(now))
+                .collect(toImmutableList()));
   }
 
   public static PollMessage getOnlyPollMessage(String clientId) {
@@ -808,19 +877,15 @@ public class DatastoreHelper {
   }
 
   public static PollMessage getOnlyPollMessage(
-      EppResource resource,
+      DomainContent domain, String clientId, DateTime now, Class<? extends PollMessage> subType) {
-      String clientId,
+    return getPollMessages(domain, clientId, now).stream()
-      DateTime now,
-      Class<? extends PollMessage> subType) {
-    return getPollMessages(resource, clientId, now)
-        .stream()
         .filter(subType::isInstance)
         .map(subType::cast)
         .collect(onlyElement());
   }
 
   public static void assertAllocationTokens(AllocationToken... expectedTokens) {
-    assertThat(ofy().load().type(AllocationToken.class).list())
+    assertThat(transactIfJpaTm(() -> tm().loadAll(AllocationToken.class)))
         .comparingElementsUsing(immutableObjectCorrespondence("updateTimestamp", "creationTime"))
         .containsExactlyElementsIn(expectedTokens);
   }
@@ -861,13 +926,17 @@ public class DatastoreHelper {
   }
 
   private static <R> void saveResource(R resource, boolean wantBackup) {
-    Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup();
+    ofyOrJpaTm(
-    saver.entity(resource);
+        () -> {
-    if (resource instanceof EppResource) {
+          Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup();
-      EppResource eppResource = (EppResource) resource;
+          saver.entity(resource);
-      persistEppResourceExtras(
+          if (resource instanceof EppResource) {
-          eppResource, EppResourceIndex.create(Key.create(eppResource)), saver);
+            EppResource eppResource = (EppResource) resource;
-    }
+            persistEppResourceExtras(
+                eppResource, EppResourceIndex.create(Key.create(eppResource)), saver);
+          }
+        },
+        () -> tm().put(resource));
   }
 
   private static <R extends EppResource> void persistEppResourceExtras(
@@ -890,23 +959,25 @@ public class DatastoreHelper {
     // Datastore and not from the session cache. This is needed to trigger Objectify's load process
     // (unmarshalling entity protos to POJOs, nulling out empty collections, calling @OnLoad
     // methods, etc.) which is bypassed for entities loaded from the session cache.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
-    return ofy().load().entity(resource).now();
+    return transactIfJpaTm(() -> tm().load(resource));
   }
 
   /** Persists an EPP resource with the {@link EppResourceIndex} always going into bucket one. */
   public static <R extends EppResource> R persistEppResourceInFirstBucket(final R resource) {
     final EppResourceIndex eppResourceIndex =
         EppResourceIndex.create(Key.create(EppResourceIndexBucket.class, 1), Key.create(resource));
-    tm()
+    tm().transact(
-        .transact(
+            () ->
-            () -> {
+                ofyOrJpaTm(
-              Saver saver = ofy().save();
+                    () -> {
-              saver.entity(resource);
+                      Saver saver = ofy().save();
-              persistEppResourceExtras(resource, eppResourceIndex, saver);
+                      saver.entity(resource);
-            });
+                      persistEppResourceExtras(resource, eppResourceIndex, saver);
-    ofy().clearSessionCache();
+                    },
-    return ofy().load().entity(resource).now();
+                    () -> tm().put(resource)));
+    tm().clearSessionCache();
+    return transactIfJpaTm(() -> tm().load(resource));
   }
 
   public static <R> void persistResources(final Iterable<R> resources) {
@@ -925,9 +996,9 @@ public class DatastoreHelper {
     }
     // 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();
     for (R resource : resources) {
-      ofy().load().entity(resource).now();
+      transactIfJpaTm(() -> tm().load(resource));
     }
   }
 
@@ -943,31 +1014,57 @@ public class DatastoreHelper {
    */
   public static <R extends EppResource> R persistEppResource(final R resource) {
     checkState(!tm().inTransaction());
-    tm()
+    tm().transact(
-        .transact(
             () -> {
-              ofy()
+              tm().put(resource);
-                  .save()
+              tm().put(
-                  .<ImmutableObject>entities(
-                      resource,
                       new HistoryEntry.Builder()
                           .setParent(resource)
                           .setType(getHistoryEntryType(resource))
                           .setModificationTime(tm().getTransactionTime())
-                          .build());
+                          .build()
-              ofy().save().entity(ForeignKeyIndex.create(resource, resource.getDeletionTime()));
+                          .toChildHistoryEntity());
+              ofyTmOrDoNothing(
+                  () -> tm().put(ForeignKeyIndex.create(resource, resource.getDeletionTime())));
             });
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
-    return ofy().load().entity(resource).safe();
+    return transactIfJpaTm(() -> tm().load(resource));
   }
 
   /** Returns all of the history entries that are parented off the given EppResource. */
-  public static List<HistoryEntry> getHistoryEntries(EppResource resource) {
+  public static List<? extends HistoryEntry> getHistoryEntries(EppResource resource) {
-    return ofy().load()
+    return ofyOrJpaTm(
-        .type(HistoryEntry.class)
+        () ->
-        .ancestor(resource)
+            ofy()
-        .order("modificationTime")
+                .load()
-        .list();
+                .type(HistoryEntry.class)
+                .ancestor(resource)
+                .order("modificationTime")
+                .list(),
+        () ->
+            tm().transact(
+                    () -> {
+                      ImmutableList<? extends HistoryEntry> unsorted = null;
+                      if (resource instanceof ContactBase) {
+                        unsorted = tm().loadAll(ContactHistory.class);
+                      } else if (resource instanceof HostBase) {
+                        unsorted = tm().loadAll(HostHistory.class);
+                      } else if (resource instanceof DomainContent) {
+                        unsorted = tm().loadAll(DomainHistory.class);
+                      } else {
+                        fail("Expected an EppResource instance, but got " + resource.getClass());
+                      }
+                      ImmutableList<? extends HistoryEntry> filtered =
+                          unsorted.stream()
+                              .filter(
+                                  historyEntry ->
+                                      historyEntry
+                                          .getParent()
+                                          .getName()
+                                          .equals(resource.getRepoId()))
+                              .collect(toImmutableList());
+                      return ImmutableList.sortedCopyOf(DateTimeComparator.getInstance(), filtered);
+                    }));
   }
 
   /**
@@ -1009,9 +1106,13 @@ public class DatastoreHelper {
   }
 
   public static PollMessage getOnlyPollMessageForHistoryEntry(HistoryEntry historyEntry) {
-    return Iterables.getOnlyElement(ofy().load()
+    return Iterables.getOnlyElement(
-        .type(PollMessage.class)
+        transactIfJpaTm(
-        .ancestor(historyEntry));
+            () ->
+                tm().loadAll(PollMessage.class).stream()
+                    .filter(
+                        pollMessage -> pollMessage.getParentKey().equals(Key.create(historyEntry)))
+                    .collect(toImmutableList())));
   }
 
   public static <T extends EppResource> HistoryEntry createHistoryEntryForEppResource(
@@ -1029,23 +1130,30 @@ public class DatastoreHelper {
    * ForeignKeyedEppResources.
    */
   public static <R> ImmutableList<R> persistSimpleResources(final Iterable<R> resources) {
-    tm().transact(() -> ofy().saveWithoutBackup().entities(resources));
+    tm().transact(() -> tm().putAllWithoutBackup(ImmutableList.copyOf(resources)));
     // 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 ImmutableList.copyOf(ofy().load().entities(resources).values());
+    return transactIfJpaTm(() -> tm().loadAll(resources));
   }
 
   public static void deleteResource(final Object resource) {
-    ofy().deleteWithoutBackup().entity(resource).now();
+    transactIfJpaTm(() -> tm().deleteWithoutBackup(resource));
     // 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();
   }
 
   /** Force the create and update timestamps to get written into the resource. **/
   public static <R> R cloneAndSetAutoTimestamps(final R resource) {
-    return tm().transact(() -> ofy().load().fromEntity(ofy().save().toEntity(resource)));
+    return tm().transact(
+            () ->
+                ofyOrJpaTm(
+                    () -> ofy().load().fromEntity(ofy().save().toEntity(resource)),
+                    () -> {
+                      tm().put(resource);
+                      return tm().load(resource);
+                    }));
   }
 
   /** Returns the entire map of {@link PremiumListEntry}s for the given {@link PremiumList}. */

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

@@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableList;
 import google.registry.persistence.transaction.TransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import java.lang.reflect.Field;
+import java.lang.reflect.Method;
 import java.util.List;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -52,9 +53,21 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
   @Override
   public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(
       ExtensionContext context) {
-    return Stream.of(
+    TestTemplateInvocationContext ofyContext =
-        createInvocationContext("Test Datastore", TransactionManagerFactory::ofyTm),
+        createInvocationContext("Test Datastore", TransactionManagerFactory::ofyTm);
-        createInvocationContext("Test PostgreSQL", TransactionManagerFactory::jpaTm));
+    TestTemplateInvocationContext sqlContext =
+        createInvocationContext("Test PostgreSQL", TransactionManagerFactory::jpaTm);
+    Method testMethod = context.getTestMethod().orElseThrow(IllegalStateException::new);
+    if (testMethod.isAnnotationPresent(TestOfyAndSql.class)) {
+      return Stream.of(ofyContext, sqlContext);
+    } else if (testMethod.isAnnotationPresent(TestOfyOnly.class)) {
+      return Stream.of(ofyContext);
+    } else if (testMethod.isAnnotationPresent(TestSqlOnly.class)) {
+      return Stream.of(sqlContext);
+    } else {
+      throw new IllegalStateException(
+          "Test method must be annotated with @TestOfyAndSql, @TestOfyOnly or @TestSqlOnly");
+    }
   }
 
   private TestTemplateInvocationContext createInvocationContext(
@@ -104,6 +117,18 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
 
   static void injectTmForDualDatabaseTest(ExtensionContext context) {
     if (isDualDatabaseTest(context)) {
+      context
+          .getTestMethod()
+          .ifPresent(
+              testMethod -> {
+                if (!testMethod.isAnnotationPresent(TestOfyAndSql.class)
+                    && !testMethod.isAnnotationPresent(TestOfyOnly.class)
+                    && !testMethod.isAnnotationPresent(TestSqlOnly.class)) {
+                  throw new IllegalStateException(
+                      "Test method must be annotated with @TestOfyAndSql, @TestOfyOnly or"
+                          + " @TestSqlOnly");
+                }
+              });
       context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm());
       Supplier<? extends TransactionManager> tmSupplier =
           (Supplier<? extends TransactionManager>)

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

@@ -19,37 +19,67 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.persistence.transaction.TransactionManager;
 import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /**
- * Test to verify that {@link DualDatabaseTestInvocationContextProvider} extension executes {@link
+ * Test to verify that {@link DualDatabaseTestInvocationContextProvider} extension executes tests
- * TestTemplate} test twice with different databases.
+ * with corresponding {@link TransactionManager}.
  */
 @DualDatabaseTest
 public class DualDatabaseTestInvocationContextProviderTest {
 
-  private static int datastoreTestCounter = 0;
+  private static int testBothDbsOfyCounter = 0;
-  private static int postgresqlTestCounter = 0;
+  private static int testBothDbsSqlCounter = 0;
+  private static int testOfyOnlyOfyCounter = 0;
+  private static int testOfyOnlySqlCounter = 0;
+  private static int testSqlOnlyOfyCounter = 0;
+  private static int testSqlOnlySqlCounter = 0;
 
   @RegisterExtension
   public final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
-  @TestTemplate
+  @TestOfyAndSql
-  void testToUseTransactionManager() {
+  void testToVerifyBothOfyAndSqlTmAreUsed() {
     if (tm() instanceof DatastoreTransactionManager) {
-      datastoreTestCounter++;
+      testBothDbsOfyCounter++;
     }
     if (tm() instanceof JpaTransactionManager) {
-      postgresqlTestCounter++;
+      testBothDbsSqlCounter++;
+    }
+  }
+
+  @TestOfyOnly
+  void testToVerifyOnlyOfyTmIsUsed() {
+    if (tm() instanceof DatastoreTransactionManager) {
+      testOfyOnlyOfyCounter++;
+    }
+    if (tm() instanceof JpaTransactionManager) {
+      testOfyOnlySqlCounter++;
+    }
+  }
+
+  @TestSqlOnly
+  void testToVerifyOnlySqlTmIsUsed() {
+    if (tm() instanceof DatastoreTransactionManager) {
+      testSqlOnlyOfyCounter++;
+    }
+    if (tm() instanceof JpaTransactionManager) {
+      testSqlOnlySqlCounter++;
     }
   }
 
   @AfterAll
   static void assertEachTransactionManagerIsUsed() {
-    assertThat(datastoreTestCounter).isEqualTo(1);
+    assertThat(testBothDbsOfyCounter).isEqualTo(1);
-    assertThat(postgresqlTestCounter).isEqualTo(1);
+    assertThat(testBothDbsSqlCounter).isEqualTo(1);
+
+    assertThat(testOfyOnlyOfyCounter).isEqualTo(1);
+    assertThat(testOfyOnlySqlCounter).isEqualTo(0);
+
+    assertThat(testSqlOnlyOfyCounter).isEqualTo(0);
+    assertThat(testSqlOnlySqlCounter).isEqualTo(1);
   }
 }

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

@@ -0,0 +1,31 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.testing;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import org.junit.jupiter.api.TestTemplate;
+
+/**
+ * Annotation to indicate a test method will be executed twice with Datastore and Postgresql
+ * respectively.
+ */
+@Target({METHOD})
+@Retention(RUNTIME)
+@TestTemplate
+public @interface TestOfyAndSql {}

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

@@ -0,0 +1,28 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.testing;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import org.junit.jupiter.api.TestTemplate;
+
+/** Annotation to indicate a test method will be executed only with Datastore. */
+@Target({METHOD})
+@Retention(RUNTIME)
+@TestTemplate
+public @interface TestOfyOnly {}

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

@@ -0,0 +1,28 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.testing;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import org.junit.jupiter.api.TestTemplate;
+
+/** Annotation to indicate a test method will be executed only with Postgresql. */
+@Target({METHOD})
+@Retention(RUNTIME)
+@TestTemplate
+public @interface TestSqlOnly {}

db/src/main/resources/sql/er_diagram/brief_er_diagram.html

@@ -261,11 +261,11 @@ td.section {
     </tr> 
     <tr> 
      <td class="property_name">generated on</td> 
-     <td class="property_value">2020-10-26 15:49:09.291097</td> 
+     <td class="property_value">2020-10-26 18:19:17.415062</td> 
     </tr> 
     <tr>
      <td class="property_name">last flyway file</td>
-     <td id="lastFlywayFile" class="property_value">V67__grace_period_history_ids.sql</td>
+     <td id="lastFlywayFile" class="property_value">V68__make_reserved_list_nullable_in_registry.sql</td>
     </tr>
    </tbody>
   </table> 
@@ -284,7 +284,7 @@ td.section {
      generated on
     </text> 
     <text text-anchor="start" x="3310.05" y="-10.8" font-family="Helvetica,sans-Serif" font-size="14.00">
-     2020-10-26 15:49:09.291097
+     2020-10-26 18:19:17.415062
     </text> 
     <polygon fill="none" stroke="#888888" points="3222.55,-4 3222.55,-44 3487.55,-44 3487.55,-4 3222.55,-4" /> <!-- allocationtoken_a08ccbef --> 
     <g id="node1" class="node"> 

db/src/main/resources/sql/er_diagram/full_er_diagram.html

@@ -261,11 +261,11 @@ td.section {
     </tr> 
     <tr> 
      <td class="property_name">generated on</td> 
-     <td class="property_value">2020-10-26 15:49:07.386754</td> 
+     <td class="property_value">2020-10-26 18:19:15.617359</td> 
     </tr> 
     <tr>
      <td class="property_name">last flyway file</td>
-     <td id="lastFlywayFile" class="property_value">V67__grace_period_history_ids.sql</td>
+     <td id="lastFlywayFile" class="property_value">V68__make_reserved_list_nullable_in_registry.sql</td>
     </tr>
    </tbody>
   </table> 
@@ -284,7 +284,7 @@ td.section {
      generated on
     </text> 
     <text text-anchor="start" x="3735.5" y="-10.8" font-family="Helvetica,sans-Serif" font-size="14.00">
-     2020-10-26 15:49:07.386754
+     2020-10-26 18:19:15.617359
     </text> 
     <polygon fill="none" stroke="#888888" points="3648,-4 3648,-44 3913,-44 3913,-4 3648,-4" /> <!-- allocationtoken_a08ccbef --> 
     <g id="node1" class="node"> 
@@ -5600,7 +5600,7 @@ td.section {
      <text text-anchor="start" x="3780.5" y="-6672.4" font-family="Helvetica,sans-Serif" font-size="14.00"> 
      </text> 
      <text text-anchor="start" x="3788.5" y="-6672.4" font-family="Helvetica,sans-Serif" font-size="14.00">
-      _text not null
+      _text
      </text> 
      <text text-anchor="start" x="3509.5" y="-6653.4" font-family="Helvetica,sans-Serif" font-size="14.00">
       restore_billing_cost_amount
@@ -11740,7 +11740,7 @@ td.section {
     <tr> 
      <td class="spacer"></td> 
      <td class="minwidth">reserved_list_names</td> 
-     <td class="minwidth">_text not null</td> 
+     <td class="minwidth">_text</td> 
     </tr> 
     <tr> 
      <td class="spacer"></td> 

db/src/main/resources/sql/flyway.txt

@@ -65,3 +65,4 @@ V64__transfer_history_columns.sql
 V65__local_date_date_type.sql
 V66__create_rde_revision.sql
 V67__grace_period_history_ids.sql
+V68__make_reserved_list_nullable_in_registry.sql

db/src/main/resources/sql/flyway/V68__make_reserved_list_nullable_in_registry.sql

@@ -0,0 +1,15 @@
+-- Copyright 2020 The Nomulus Authors. All Rights Reserved.
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+ALTER TABLE "Tld" ALTER COLUMN reserved_list_names DROP NOT NULL;

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -646,7 +646,7 @@
         registry_lock_or_unlock_cost_currency text,
         renew_billing_cost_transitions hstore not null,
         renew_grace_period_length interval not null,
-        reserved_list_names text[] not null,
+        reserved_list_names text[],
         restore_billing_cost_amount numeric(19, 2),
         restore_billing_cost_currency text,
         roid_suffix text,

db/src/main/resources/sql/schema/nomulus.golden.sql

@@ -904,7 +904,7 @@ CREATE TABLE public."Tld" (
     registry_lock_or_unlock_cost_currency text,
     renew_billing_cost_transitions public.hstore NOT NULL,
     renew_grace_period_length interval NOT NULL,
-    reserved_list_names text[] NOT NULL,
+    reserved_list_names text[],
     restore_billing_cost_amount numeric(19,2),
     restore_billing_cost_currency text,
     roid_suffix text,