Convert DomainTransferRejectFlow to use tm() methods (#977)

* Convert DomainTransferRejectFlow to use tm() methods

This change includes a few other necessary dependencies to converting
DomainTransferRejectFlowTest to be a dual-database test. Namely:

- The basic "use tm() instead of ofy()" and "branching database
selection on what were previously raw ofy queries"
- Modification of the PollMessage convertVKey methods to do what they
say they do
- Filling the generic pending / response fields in PollMessage based on what type of
poll message it is (this has to be done because SQL is not very good at
storing ambiguous superclasses)
- Setting the generic pending / repsonse fields in PollMessage upon
build
- Filling out the serverApproveEntities field in DomainTransferData with
all necessary poll messages / billing events that should be cancelled on
rejection
- Scattered changes in DatabaseHelper to make sure that we're saving and
loading entities correctly where we weren't before

core/src/main/java/google/registry/flows/ResourceFlowUtils.java

@@ -18,12 +18,11 @@ import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.EppResourceUtils.getLinkedDomainKeys;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.InvalidAuthorizationInformationErrorException;
 import google.registry.flows.EppException.ObjectDoesNotExistException;
@@ -185,16 +184,15 @@ public final class ResourceFlowUtils {
       return;
     }
     // The roid should match one of the contacts.
-    Optional<Key<ContactResource>> foundContact =
+    Optional<VKey<ContactResource>> foundContact =
         domain.getReferencedContacts().stream()
-            .map(VKey::getOfyKey)
+            .filter(key -> key.getOfyKey().getName().equals(authRepoId))
-            .filter(key -> key.getName().equals(authRepoId))
             .findFirst();
     if (!foundContact.isPresent()) {
       throw new BadAuthInfoForResourceException();
     }
     // Check the authInfo against the contact.
-    verifyAuthInfo(authInfo, ofy().load().key(foundContact.get()).now());
+    verifyAuthInfo(authInfo, transactIfJpaTm(() -> tm().loadByKey(foundContact.get())));
   }
 
   /** Check that the given {@link AuthInfo} is valid for the given contact. */

core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java

@@ -25,6 +25,9 @@ import static com.google.common.collect.Iterables.any;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
 import static com.google.common.collect.Sets.union;
+import static google.registry.model.DatabaseMigrationUtils.getPrimaryDatabase;
+import static google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase.DATASTORE;
+import static google.registry.model.common.DatabaseTransitionSchedule.TransitionId.REPLAYED_ENTITIES;
 import static google.registry.model.domain.DomainBase.MAX_REGISTRATION_YEARS;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registries.findTldForName;
@@ -38,6 +41,7 @@ import static google.registry.model.registry.label.ReservationType.FULLY_BLOCKED
 import static google.registry.model.registry.label.ReservationType.NAME_COLLISION;
 import static google.registry.model.registry.label.ReservationType.RESERVED_FOR_ANCHOR_TENANT;
 import static google.registry.model.registry.label.ReservationType.RESERVED_FOR_SPECIFIC_USE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.isDomainPremium;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -88,6 +92,7 @@ import google.registry.model.domain.DomainCommand.Create;
 import google.registry.model.domain.DomainCommand.CreateOrUpdate;
 import google.registry.model.domain.DomainCommand.InvalidReferencesException;
 import google.registry.model.domain.DomainCommand.Update;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.ForeignKeyedDesignatedContact;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.BaseFee;
@@ -355,22 +360,23 @@ public class DomainFlowUtils {
 
   static void validateNoDuplicateContacts(Set<DesignatedContact> contacts)
       throws ParameterValuePolicyErrorException {
-    ImmutableMultimap<Type, Key<ContactResource>> contactsByType =
+    ImmutableMultimap<Type, VKey<ContactResource>> contactsByType =
         contacts.stream()
             .collect(
                 toImmutableSetMultimap(
-                    DesignatedContact::getType, contact -> contact.getContactKey().getOfyKey()));
+                    DesignatedContact::getType, contact -> contact.getContactKey()));
 
     // If any contact type has multiple contacts:
     if (contactsByType.asMap().values().stream().anyMatch(v -> v.size() > 1)) {
       // Find the duplicates.
-      Map<Type, Collection<Key<ContactResource>>> dupeKeysMap =
+      Map<Type, Collection<VKey<ContactResource>>> dupeKeysMap =
           Maps.filterEntries(contactsByType.asMap(), e -> e.getValue().size() > 1);
-      ImmutableList<Key<ContactResource>> dupeKeys =
+      ImmutableList<VKey<ContactResource>> dupeKeys =
           dupeKeysMap.values().stream().flatMap(Collection::stream).collect(toImmutableList());
       // Load the duplicates in one batch.
-      Map<Key<ContactResource>, ContactResource> dupeContacts = ofy().load().keys(dupeKeys);
+      Map<VKey<? extends ContactResource>, ContactResource> dupeContacts =
-      ImmutableMultimap.Builder<Type, Key<ContactResource>> typesMap =
+          tm().loadByKeys(dupeKeys);
+      ImmutableMultimap.Builder<Type, VKey<ContactResource>> typesMap =
           new ImmutableMultimap.Builder<>();
       dupeKeysMap.forEach(typesMap::putAll);
       // Create an error message showing the type and contact IDs of the duplicates.
@@ -537,13 +543,13 @@ public class DomainFlowUtils {
     // If the resultant autorenew poll message would have no poll messages to deliver, then just
     // delete it. Otherwise save it with the new end time.
     if (isAtOrAfter(updatedAutorenewPollMessage.getEventTime(), newEndTime)) {
-      autorenewPollMessage.ifPresent(autorenew -> ofy().delete().entity(autorenew));
+      autorenewPollMessage.ifPresent(autorenew -> tm().delete(autorenew));
     } else {
-      ofy().save().entity(updatedAutorenewPollMessage);
+      tm().put(updatedAutorenewPollMessage);
     }
 
     Recurring recurring = tm().loadByKey(domain.getAutorenewBillingEvent());
-    ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
+    tm().put(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 
   /**
@@ -1045,18 +1051,11 @@ public class DomainFlowUtils {
       Duration maxSearchPeriod,
       final ImmutableSet<TransactionReportField> cancelableFields) {
 
-    List<HistoryEntry> recentHistoryEntries =
+    List<? extends HistoryEntry> recentHistoryEntries =
-        ofy()
+        findRecentHistoryEntries(domainBase, now, maxSearchPeriod);
-            .load()
+    Optional<? extends HistoryEntry> entryToCancel =
-            .type(HistoryEntry.class)
-            .ancestor(domainBase)
-            .filter("modificationTime >=", now.minus(maxSearchPeriod))
-            .order("modificationTime")
-            .list();
-    Optional<HistoryEntry> entryToCancel =
         Streams.findLast(
-            recentHistoryEntries
+            recentHistoryEntries.stream()
-                .stream()
                 .filter(
                     historyEntry -> {
                       // Look for add and renew transaction records that have yet to be reported
@@ -1082,6 +1081,28 @@ public class DomainFlowUtils {
     return recordsBuilder.build();
   }
 
+  private static List<? extends HistoryEntry> findRecentHistoryEntries(
+      DomainBase domainBase, DateTime now, Duration maxSearchPeriod) {
+    if (getPrimaryDatabase(REPLAYED_ENTITIES).equals(DATASTORE)) {
+      return ofy()
+          .load()
+          .type(HistoryEntry.class)
+          .ancestor(domainBase)
+          .filter("modificationTime >=", now.minus(maxSearchPeriod))
+          .order("modificationTime")
+          .list();
+    } else {
+      return jpaTm()
+          .getEntityManager()
+          .createQuery(
+              "FROM DomainHistory WHERE modificationTime >= :beginning "
+                  + "ORDER BY modificationTime ASC",
+              DomainHistory.class)
+          .setParameter("beginning", now.minus(maxSearchPeriod))
+          .getResultList();
+    }
+  }
+
   /** Resource linked to this domain does not exist. */
   static class LinkedResourcesDoNotExistException extends ObjectDoesNotExistException {
     public LinkedResourcesDoNotExistException(Class<?> type, ImmutableSet<String> resourceIds) {

core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java

@@ -25,7 +25,6 @@ import static google.registry.flows.domain.DomainFlowUtils.updateAutorenewRecurr
 import static google.registry.flows.domain.DomainTransferUtils.createGainingTransferPollMessage;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_NACKED;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -41,7 +40,6 @@ import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
-import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -102,11 +100,11 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
     }
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_REJECTED, now, clientId);
-    ofy().save().<ImmutableObject>entities(
+    tm().putAll(
-        newDomain,
+            newDomain,
-        historyEntry,
+            historyEntry,
-        createGainingTransferPollMessage(
+            createGainingTransferPollMessage(
-            targetId, newDomain.getTransferData(), null, historyEntry));
+                targetId, newDomain.getTransferData(), null, historyEntry));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may end up recreating the poll message if it was deleted upon the transfer request.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);

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

@@ -118,9 +118,7 @@ public final class ResourceTransferUtils {
     if (resource.getStatusValues().contains(StatusValue.PENDING_TRANSFER)) {
       TransferData oldTransferData = resource.getTransferData();
       tm().delete(oldTransferData.getServerApproveEntities());
-      ofy()
+      tm().put(
-          .save()
-          .entity(
               new PollMessage.OneTime.Builder()
                   .setClientId(oldTransferData.getGainingClientId())
                   .setEventTime(now)

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

@@ -131,6 +131,11 @@ public class DatastoreTransactionManager implements TransactionManager {
     saveEntity(entity);
   }
 
+  @Override
+  public void putAll(Object... entities) {
+    syncIfTransactionless(getOfy().save().entities(entities));
+  }
+
   @Override
   public void putAll(ImmutableCollection<?> entities) {
     syncIfTransactionless(getOfy().save().entities(entities));

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -359,8 +359,13 @@ public abstract class PollMessage extends ImmutableObject
     }
 
     /** Converts an unspecialized VKey<PollMessage> to a VKey of the derived class. */
-    public static @Nullable VKey<OneTime> convertVKey(@Nullable VKey<OneTime> key) {
+    public static @Nullable VKey<OneTime> convertVKey(@Nullable VKey<? extends PollMessage> key) {
-      return key == null ? null : VKey.create(OneTime.class, key.getSqlKey(), key.getOfyKey());
+      if (key == null) {
+        return null;
+      }
+      Key<OneTime> ofyKey =
+          Key.create(key.getOfyKey().getParent(), OneTime.class, key.getOfyKey().getId());
+      return VKey.create(OneTime.class, key.getSqlKey(), ofyKey);
     }
 
     @Override
@@ -383,6 +388,7 @@ public abstract class PollMessage extends ImmutableObject
     @OnLoad
     void onLoad() {
       super.onLoad();
+      // Take the Objectify-specific fields and map them to the SQL-specific fields, if applicable
       if (!isNullOrEmpty(contactPendingActionNotificationResponses)) {
         pendingActionNotificationResponse = contactPendingActionNotificationResponses.get(0);
       }
@@ -390,36 +396,70 @@ public abstract class PollMessage extends ImmutableObject
         contactId = contactTransferResponses.get(0).getContactId();
         transferResponse = contactTransferResponses.get(0);
       }
+      if (!isNullOrEmpty(domainPendingActionNotificationResponses)) {
+        pendingActionNotificationResponse = domainPendingActionNotificationResponses.get(0);
+      }
+      if (!isNullOrEmpty(domainTransferResponses)) {
+        fullyQualifiedDomainName = domainTransferResponses.get(0).getFullyQualifiedDomainName();
+        transferResponse = domainTransferResponses.get(0);
+        extendedRegistrationExpirationTime =
+            domainTransferResponses.get(0).getExtendedRegistrationExpirationTime();
+      }
     }
 
     @Override
     @PostLoad
     void postLoad() {
       super.postLoad();
+      // Take the SQL-specific fields and map them to the Objectify-specific fields, if applicable
       if (pendingActionNotificationResponse != null) {
-        contactPendingActionNotificationResponses =
+        if (contactId != null) {
-            ImmutableList.of(
+          contactPendingActionNotificationResponses =
-                ContactPendingActionNotificationResponse.create(
+              ImmutableList.of(
-                    pendingActionNotificationResponse.nameOrId.value,
+                  ContactPendingActionNotificationResponse.create(
-                    pendingActionNotificationResponse.getActionResult(),
+                      pendingActionNotificationResponse.nameOrId.value,
-                    pendingActionNotificationResponse.getTrid(),
+                      pendingActionNotificationResponse.getActionResult(),
-                    pendingActionNotificationResponse.processedDate));
+                      pendingActionNotificationResponse.getTrid(),
+                      pendingActionNotificationResponse.processedDate));
+        } else if (fullyQualifiedDomainName != null) {
+          domainPendingActionNotificationResponses =
+              ImmutableList.of(
+                  DomainPendingActionNotificationResponse.create(
+                      pendingActionNotificationResponse.nameOrId.value,
+                      pendingActionNotificationResponse.getActionResult(),
+                      pendingActionNotificationResponse.getTrid(),
+                      pendingActionNotificationResponse.processedDate));
+        }
       }
-      if (contactId != null && transferResponse != null) {
+      if (transferResponse != null) {
-        // The transferResponse is currently an unspecialized TransferResponse instance, create a
+        // The transferResponse is currently an unspecialized TransferResponse instance, create the
-        // ContactTransferResponse so that the value is consistently specialized and store it in the
+        // appropriate subclass so that the value is consistently specialized
-        // list representation for datastore.
+        if (contactId != null) {
-        transferResponse =
+          transferResponse =
-            new ContactTransferResponse.Builder()
+              new ContactTransferResponse.Builder()
-                .setContactId(contactId)
+                  .setContactId(contactId)
-                .setGainingClientId(transferResponse.getGainingClientId())
+                  .setGainingClientId(transferResponse.getGainingClientId())
-                .setLosingClientId(transferResponse.getLosingClientId())
+                  .setLosingClientId(transferResponse.getLosingClientId())
-                .setTransferStatus(transferResponse.getTransferStatus())
+                  .setTransferStatus(transferResponse.getTransferStatus())
-                .setTransferRequestTime(transferResponse.getTransferRequestTime())
+                  .setTransferRequestTime(transferResponse.getTransferRequestTime())
-                .setPendingTransferExpirationTime(
+                  .setPendingTransferExpirationTime(
-                    transferResponse.getPendingTransferExpirationTime())
+                      transferResponse.getPendingTransferExpirationTime())
-                .build();
+                  .build();
-        contactTransferResponses = ImmutableList.of((ContactTransferResponse) transferResponse);
+          contactTransferResponses = ImmutableList.of((ContactTransferResponse) transferResponse);
+        } else if (fullyQualifiedDomainName != null) {
+          transferResponse =
+              new DomainTransferResponse.Builder()
+                  .setFullyQualifiedDomainName(fullyQualifiedDomainName)
+                  .setGainingClientId(transferResponse.getGainingClientId())
+                  .setLosingClientId(transferResponse.getLosingClientId())
+                  .setTransferStatus(transferResponse.getTransferStatus())
+                  .setTransferRequestTime(transferResponse.getTransferRequestTime())
+                  .setPendingTransferExpirationTime(
+                      transferResponse.getPendingTransferExpirationTime())
+                  .setExtendedRegistrationExpirationTime(extendedRegistrationExpirationTime)
+                  .build();
+          domainTransferResponses = ImmutableList.of((DomainTransferResponse) transferResponse);
+        }
       }
     }
 
@@ -441,10 +481,7 @@ public abstract class PollMessage extends ImmutableObject
                     .filter(ContactPendingActionNotificationResponse.class::isInstance)
                     .map(ContactPendingActionNotificationResponse.class::cast)
                     .collect(toImmutableList()));
-        if (getInstance().contactPendingActionNotificationResponses != null) {
+
-          getInstance().pendingActionNotificationResponse =
-              getInstance().contactPendingActionNotificationResponses.get(0);
-        }
         getInstance().contactTransferResponses =
             forceEmptyToNull(
                 responseData
@@ -452,10 +489,6 @@ public abstract class PollMessage extends ImmutableObject
                     .filter(ContactTransferResponse.class::isInstance)
                     .map(ContactTransferResponse.class::cast)
                     .collect(toImmutableList()));
-        if (getInstance().contactTransferResponses != null) {
-          getInstance().contactId = getInstance().contactTransferResponses.get(0).getContactId();
-          getInstance().transferResponse = getInstance().contactTransferResponses.get(0);
-        }
 
         getInstance().domainPendingActionNotificationResponses =
             forceEmptyToNull(
@@ -471,13 +504,36 @@ public abstract class PollMessage extends ImmutableObject
                     .filter(DomainTransferResponse.class::isInstance)
                     .map(DomainTransferResponse.class::cast)
                     .collect(toImmutableList()));
+
         getInstance().hostPendingActionNotificationResponses =
             forceEmptyToNull(
-                responseData
+                responseData.stream()
-                    .stream()
                     .filter(HostPendingActionNotificationResponse.class::isInstance)
                     .map(HostPendingActionNotificationResponse.class::cast)
                     .collect(toImmutableList()));
+
+        // Set the generic pending-action field as appropriate
+        if (getInstance().contactPendingActionNotificationResponses != null) {
+          getInstance().pendingActionNotificationResponse =
+              getInstance().contactPendingActionNotificationResponses.get(0);
+        } else if (getInstance().domainPendingActionNotificationResponses != null) {
+          getInstance().pendingActionNotificationResponse =
+              getInstance().domainPendingActionNotificationResponses.get(0);
+        } else if (getInstance().hostPendingActionNotificationResponses != null) {
+          getInstance().pendingActionNotificationResponse =
+              getInstance().hostPendingActionNotificationResponses.get(0);
+        }
+        // Set the generic transfer response field as appropriate
+        if (getInstance().contactTransferResponses != null) {
+          getInstance().contactId = getInstance().contactTransferResponses.get(0).getContactId();
+          getInstance().transferResponse = getInstance().contactTransferResponses.get(0);
+        } else if (getInstance().domainTransferResponses != null) {
+          getInstance().fullyQualifiedDomainName =
+              getInstance().domainTransferResponses.get(0).getFullyQualifiedDomainName();
+          getInstance().transferResponse = getInstance().domainTransferResponses.get(0);
+          getInstance().extendedRegistrationExpirationTime =
+              getInstance().domainTransferResponses.get(0).getExtendedRegistrationExpirationTime();
+        }
         return this;
       }
     }
@@ -518,8 +574,13 @@ public abstract class PollMessage extends ImmutableObject
     }
 
     /** Converts an unspecialized VKey&lt;PollMessage&gt; to a VKey of the derived class. */
-    public static @Nullable VKey<Autorenew> convertVKey(VKey<Autorenew> key) {
+    public static @Nullable VKey<Autorenew> convertVKey(VKey<? extends PollMessage> key) {
-      return key == null ? null : VKey.create(Autorenew.class, key.getSqlKey(), key.getOfyKey());
+      if (key == null) {
+        return null;
+      }
+      Key<Autorenew> ofyKey =
+          Key.create(key.getOfyKey().getParent(), Autorenew.class, key.getOfyKey().getId());
+      return VKey.create(Autorenew.class, key.getSqlKey(), ofyKey);
     }
 
     @Override

core/src/main/java/google/registry/model/transfer/DomainTransferData.java

@@ -14,7 +14,10 @@
 
 package google.registry.model.transfer;
 
+import static google.registry.util.CollectionUtils.forceEmptyToNull;
+
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.AlsoLoad;
 import com.googlecode.objectify.annotation.Embed;
@@ -34,6 +37,7 @@ import javax.persistence.AttributeOverrides;
 import javax.persistence.Column;
 import javax.persistence.Embeddable;
 import javax.persistence.Embedded;
+import javax.persistence.PostLoad;
 import org.joda.time.DateTime;
 
 /** Transfer data for domain. */
@@ -214,6 +218,28 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
     return serverApproveAutorenewPollMessageHistoryId;
   }
 
+  @PostLoad
+  @Override
+  void postLoad() {
+    // The superclass's serverApproveEntities should include the billing events if present
+    super.postLoad();
+    ImmutableSet.Builder<VKey<? extends TransferServerApproveEntity>> serverApproveEntitiesBuilder =
+        new ImmutableSet.Builder<>();
+    if (serverApproveEntities != null) {
+      serverApproveEntitiesBuilder.addAll(serverApproveEntities);
+    }
+    if (serverApproveBillingEvent != null) {
+      serverApproveEntitiesBuilder.add(serverApproveBillingEvent);
+    }
+    if (serverApproveAutorenewEvent != null) {
+      serverApproveEntitiesBuilder.add(serverApproveAutorenewEvent);
+    }
+    if (serverApproveAutorenewPollMessage != null) {
+      serverApproveEntitiesBuilder.add(serverApproveAutorenewPollMessage);
+    }
+    serverApproveEntities = forceEmptyToNull(serverApproveEntitiesBuilder.build());
+  }
+
   @Override
   public boolean isEmpty() {
     return EMPTY.equals(this);

core/src/main/java/google/registry/model/transfer/TransferData.java

@@ -167,7 +167,7 @@ public abstract class TransferData<
       return;
     }
     Key<? extends EppResource> eppKey;
-    if (getClass().equals(DomainBase.class)) {
+    if (getClass().equals(DomainTransferData.class)) {
       eppKey = Key.create(DomainBase.class, repoId);
     } else {
       eppKey = Key.create(ContactResource.class, repoId);

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

@@ -288,6 +288,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     transactionInfo.get().addUpdate(toPersist);
   }
 
+  @Override
+  public void putAll(Object... entities) {
+    checkArgumentNotNull(entities, "entities must be specified");
+    assertInTransaction();
+    for (Object entity : entities) {
+      put(entity);
+    }
+  }
+
   @Override
   public void putAll(ImmutableCollection<?> entities) {
     checkArgumentNotNull(entities, "entities must be specified");

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

@@ -123,7 +123,10 @@ public interface TransactionManager {
   /** 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. */
+  /** Persists all new entities or updates the existing entities in the database. */
+  void putAll(Object... entities);
+
+  /** Persists all new entities or updates the existing entities in the database. */
   void putAll(ImmutableCollection<?> entities);
 
   /**

core/src/test/java/google/registry/flows/domain/DomainTransferRejectFlowTest.java

@@ -23,11 +23,11 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REJECT;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.testing.DatabaseHelper.assertBillingEvents;
-import static google.registry.testing.DatabaseHelper.deleteResource;
 import static google.registry.testing.DatabaseHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatabaseHelper.getOnlyPollMessage;
 import static google.registry.testing.DatabaseHelper.getPollMessages;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
+import static google.registry.testing.DatabaseHelper.persistDomainAsDeleted;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DomainBaseSubject.assertAboutDomains;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
@@ -56,12 +56,14 @@ import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferResponse;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link DomainTransferRejectFlow}. */
+@DualDatabaseTest
 class DomainTransferRejectFlowTest
     extends DomainTransferFlowTestCase<DomainTransferRejectFlow, DomainBase> {
 
@@ -151,31 +153,31 @@ class DomainTransferRejectFlowTest
     runFlow();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     doSuccessfulTest("domain_transfer_reject.xml", "domain_transfer_reject_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     setEppInput("domain_transfer_reject.xml");
     eppLoader.replaceAll("JD1234-REP", contact.getRepoId());
     dryRunFlowAssertResponse(loadFile("domain_transfer_reject_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_domainAuthInfo() throws Exception {
     doSuccessfulTest(
         "domain_transfer_reject_domain_authinfo.xml", "domain_transfer_reject_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_contactAuthInfo() throws Exception {
     doSuccessfulTest(
         "domain_transfer_reject_contact_authinfo.xml", "domain_transfer_reject_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_notAuthorizedForTld() {
     persistResource(
         loadRegistrar("TheRegistrar").asBuilder().setAllowedTlds(ImmutableSet.of()).build());
@@ -188,7 +190,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserNotAuthorizedForTld() throws Exception {
     persistResource(
         loadRegistrar("TheRegistrar").asBuilder().setAllowedTlds(ImmutableSet.of()).build());
@@ -196,7 +198,7 @@ class DomainTransferRejectFlowTest
         CommitMode.LIVE, UserPrivileges.SUPERUSER, loadFile("domain_transfer_reject_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badContactPassword() {
     // Change the contact's password so it does not match the password in the file.
     contact =
@@ -212,7 +214,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badDomainPassword() {
     // Change the domain's password so it does not match the password in the file.
     domain =
@@ -228,7 +230,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neverBeenTransferred() {
     changeTransferStatus(null);
     EppException thrown =
@@ -237,7 +239,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientApproved() {
     changeTransferStatus(TransferStatus.CLIENT_APPROVED);
     EppException thrown =
@@ -246,7 +248,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientRejected() {
     changeTransferStatus(TransferStatus.CLIENT_REJECTED);
     EppException thrown =
@@ -255,7 +257,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientCancelled() {
     changeTransferStatus(TransferStatus.CLIENT_CANCELLED);
     EppException thrown =
@@ -264,7 +266,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_serverApproved() {
     changeTransferStatus(TransferStatus.SERVER_APPROVED);
     EppException thrown =
@@ -273,7 +275,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_serverCancelled() {
     changeTransferStatus(TransferStatus.SERVER_CANCELLED);
     EppException thrown =
@@ -282,7 +284,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_gainingClient() {
     setClientIdForFlow("NewRegistrar");
     EppException thrown =
@@ -291,7 +293,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_unrelatedClient() {
     setClientIdForFlow("ClientZ");
     EppException thrown =
@@ -300,7 +302,7 @@ class DomainTransferRejectFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_deletedDomain() throws Exception {
     domain =
         persistResource(domain.asBuilder().setDeletionTime(clock.nowUtc().minusDays(1)).build());
@@ -310,9 +312,9 @@ class DomainTransferRejectFlowTest
     assertThat(thrown).hasMessageThat().contains(String.format("(%s)", getUniqueIdFromCommand()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonexistentDomain() throws Exception {
-    deleteResource(domain);
+    persistDomainAsDeleted(domain, clock.nowUtc());
     ResourceDoesNotExistException thrown =
         assertThrows(
             ResourceDoesNotExistException.class, () -> doFailingTest("domain_transfer_reject.xml"));
@@ -322,7 +324,7 @@ class DomainTransferRejectFlowTest
   // NB: No need to test pending delete status since pending transfers will get cancelled upon
   // entering pending delete phase. So it's already handled in that test case.
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-dom-transfer-reject");
@@ -338,7 +340,7 @@ class DomainTransferRejectFlowTest
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannTransactionRecord_noRecordsToCancel() throws Exception {
     setUpGracePeriodDurations();
     runFlow();
@@ -348,7 +350,7 @@ class DomainTransferRejectFlowTest
         .containsExactly(DomainTransactionRecord.create("tld", clock.nowUtc(), TRANSFER_NACKED, 1));
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannTransactionRecord_cancelsPreviousRecords() throws Exception {
     setUpGracePeriodDurations();
     DomainTransactionRecord previousSuccessRecord =

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

@@ -21,6 +21,7 @@ import static com.google.common.truth.Truth.assertThat;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import com.google.common.truth.Correspondence;
 import com.google.common.truth.Correspondence.BinaryPredicate;
@@ -50,6 +51,11 @@ public final class ImmutableObjectSubject extends Subject {
     this.actual = actual;
   }
 
+  public void isEqualExceptFields(
+      @Nullable ImmutableObject expected, Iterable<String> ignoredFields) {
+    isEqualExceptFields(expected, Iterables.toArray(ignoredFields, String.class));
+  }
+
   public void isEqualExceptFields(@Nullable ImmutableObject expected, String... ignoredFields) {
     if (actual == null) {
       assertThat(expected).isNull();

core/src/test/java/google/registry/model/domain/DomainBaseSqlTest.java

@@ -15,6 +15,7 @@
 package google.registry.model.domain;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.flows.domain.DomainTransferUtils.createPendingTransferData;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -34,11 +35,13 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DesignatedContact.Type;
+import google.registry.model.domain.Period.Unit;
 import google.registry.model.domain.launch.LaunchNotice;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.eppcommon.Trid;
 import google.registry.model.host.HostResource;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
@@ -54,6 +57,7 @@ import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableList;
 
 /** Verify that we can store/retrieve DomainBase objects from a SQL database. */
 @DualDatabaseTest
@@ -617,15 +621,15 @@ public class DomainBaseSqlTest {
                       .setParent(historyEntry)
                       .build();
               DomainTransferData transferData =
-                  new DomainTransferData.Builder()
+                  createPendingTransferData(
-                      .setServerApproveBillingEvent(
+                      new DomainTransferData.Builder()
-                          createLegacyVKey(BillingEvent.OneTime.class, oneTimeBillingEvent.getId()))
+                          .setTransferRequestTrid(Trid.create("foo", "bar"))
-                      .setServerApproveAutorenewEvent(
+                          .setTransferRequestTime(fakeClock.nowUtc())
-                          createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId()))
+                          .setGainingClientId("registrar2")
-                      .setServerApproveAutorenewPollMessage(
+                          .setLosingClientId("registrar1")
-                          createLegacyVKey(
+                          .setPendingTransferExpirationTime(fakeClock.nowUtc().plusDays(1)),
-                              PollMessage.Autorenew.class, autorenewPollMessage.getId()))
+                      ImmutableSet.of(oneTimeBillingEvent, billEvent, autorenewPollMessage),
-                      .build();
+                      Period.create(0, Unit.YEARS));
               gracePeriods =
                   ImmutableSet.of(
                       GracePeriod.create(
@@ -708,10 +712,17 @@ public class DomainBaseSqlTest {
     // Fix the original creation timestamp (this gets initialized on first write)
     DomainBase org = domain.asBuilder().setCreationTime(thatDomain.getCreationTime()).build();
 
-    String[] moreExcepts = Arrays.copyOf(excepts, excepts.length + 1);
+    ImmutableList<String> moreExcepts =
-    moreExcepts[moreExcepts.length - 1] = "updateTimestamp";
+        new ImmutableList.Builder<String>()
-
+            .addAll(Arrays.asList(excepts))
+            .add("updateTimestamp")
+            .add("transferData")
+            .build();
     // Note that the equality comparison forces a lazy load of all fields.
     assertAboutImmutableObjects().that(thatDomain).isEqualExceptFields(org, moreExcepts);
+    // Transfer data cannot be directly compared due to serverApproveEtities inequalities
+    assertAboutImmutableObjects()
+        .that(domain.getTransferData())
+        .isEqualExceptFields(org.getTransferData(), "serverApproveEntities");
   }
 }

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

@@ -326,8 +326,7 @@ public class DatabaseHelper {
    * Returns a persisted domain that is the passed-in domain modified to be deleted at the specified
    * time.
    */
-  public static DomainBase persistDomainAsDeleted(
+  public static DomainBase persistDomainAsDeleted(DomainBase domain, DateTime deletionTime) {
-      DomainBase domain, DateTime deletionTime) {
     return persistResource(domain.asBuilder().setDeletionTime(deletionTime).build());
   }
 
@@ -347,7 +346,7 @@ public class DatabaseHelper {
   }
 
   public static ReservedList persistReservedList(
-    String listName, boolean shouldPublish, String... lines) {
+      String listName, boolean shouldPublish, String... lines) {
     ReservedList reservedList =
         new ReservedList.Builder()
             .setName(listName)
@@ -512,10 +511,7 @@ public class DatabaseHelper {
   }
 
   public static BillingEvent.OneTime createBillingEventForTransfer(
-      DomainBase domain,
+      DomainBase domain, HistoryEntry historyEntry, DateTime costLookupTime, DateTime eventTime) {
-      HistoryEntry historyEntry,
-      DateTime costLookupTime,
-      DateTime eventTime) {
     return new BillingEvent.OneTime.Builder()
         .setReason(Reason.TRANSFER)
         .setTargetId(domain.getDomainName())
@@ -530,10 +526,7 @@ public class DatabaseHelper {
   }
 
   public static ContactResource persistContactWithPendingTransfer(
-      ContactResource contact,
+      ContactResource contact, DateTime requestTime, DateTime expirationTime, DateTime now) {
-      DateTime requestTime,
-      DateTime expirationTime,
-      DateTime now) {
     HistoryEntry historyEntryContactTransfer =
         persistResource(
             new HistoryEntry.Builder()
@@ -587,26 +580,29 @@ public class DatabaseHelper {
     String domainName = String.format("%s.%s", label, tld);
     String repoId = generateNewDomainRoid(tld);
     DomainBase domain =
-        new DomainBase.Builder()
-            .setRepoId(repoId)
-            .setDomainName(domainName)
-            .setPersistedCurrentSponsorClientId("TheRegistrar")
-            .setCreationClientId("TheRegistrar")
-            .setCreationTimeForTest(creationTime)
-            .setRegistrationExpirationTime(expirationTime)
-            .setRegistrant(contact.createVKey())
-            .setContacts(
-                ImmutableSet.of(
-                    DesignatedContact.create(Type.ADMIN, contact.createVKey()),
-                    DesignatedContact.create(Type.TECH, contact.createVKey())))
-            .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("fooBAR")))
-            .addGracePeriod(
-                GracePeriod.create(GracePeriodStatus.ADD, repoId, now.plusDays(10), "foo", null))
-            .build();
-    HistoryEntry historyEntryDomainCreate =
         persistResource(
-            new HistoryEntry.Builder()
+            new DomainBase.Builder()
+                .setRepoId(repoId)
+                .setDomainName(domainName)
+                .setPersistedCurrentSponsorClientId("TheRegistrar")
+                .setCreationClientId("TheRegistrar")
+                .setCreationTimeForTest(creationTime)
+                .setRegistrationExpirationTime(expirationTime)
+                .setRegistrant(contact.createVKey())
+                .setContacts(
+                    ImmutableSet.of(
+                        DesignatedContact.create(Type.ADMIN, contact.createVKey()),
+                        DesignatedContact.create(Type.TECH, contact.createVKey())))
+                .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("fooBAR")))
+                .addGracePeriod(
+                    GracePeriod.create(
+                        GracePeriodStatus.ADD, repoId, now.plusDays(10), "TheRegistrar", null))
+                .build());
+    DomainHistory historyEntryDomainCreate =
+        persistResource(
+            new DomainHistory.Builder()
                 .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                .setModificationTime(now)
                 .setParent(domain)
                 .build());
     BillingEvent.Recurring autorenewEvent =
@@ -643,18 +639,17 @@ public class DatabaseHelper {
       DateTime requestTime,
       DateTime expirationTime,
       DateTime extendedRegistrationExpirationTime) {
-    HistoryEntry historyEntryDomainTransfer =
+    DomainHistory historyEntryDomainTransfer =
         persistResource(
-            new HistoryEntry.Builder()
+            new DomainHistory.Builder()
                 .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
                 .setModificationTime(tm().transact(() -> tm().getTransactionTime()))
                 .setParent(domain)
                 .build());
-    BillingEvent.OneTime transferBillingEvent = persistResource(createBillingEventForTransfer(
+    BillingEvent.OneTime transferBillingEvent =
-            domain,
+        persistResource(
-            historyEntryDomainTransfer,
+            createBillingEventForTransfer(
-            requestTime,
+                domain, historyEntryDomainTransfer, requestTime, expirationTime));
-            expirationTime));
     BillingEvent.Recurring gainingClientAutorenewEvent =
         persistResource(
             new BillingEvent.Recurring.Builder()
@@ -678,18 +673,18 @@ public class DatabaseHelper {
                 .build());
     // Modify the existing autorenew event to reflect the pending transfer.
     persistResource(
-        tm().loadByKey(domain.getAutorenewBillingEvent())
+        transactIfJpaTm(
-            .asBuilder()
+            () ->
-            .setRecurrenceEndTime(expirationTime)
+                tm().loadByKey(domain.getAutorenewBillingEvent())
-            .build());
+                    .asBuilder()
+                    .setRecurrenceEndTime(expirationTime)
+                    .build()));
     // Update the end time of the existing autorenew poll message. We must delete it if it has no
     // events left in it.
-    PollMessage.Autorenew autorenewPollMessage = tm().loadByKey(domain.getAutorenewPollMessage());
+    PollMessage.Autorenew autorenewPollMessage =
+        transactIfJpaTm(() -> tm().loadByKey(domain.getAutorenewPollMessage()));
     if (autorenewPollMessage.getEventTime().isBefore(expirationTime)) {
-      persistResource(
+      persistResource(autorenewPollMessage.asBuilder().setAutorenewEndTime(expirationTime).build());
-          autorenewPollMessage.asBuilder()
-              .setAutorenewEndTime(expirationTime)
-              .build());
     } else {
       deleteResource(autorenewPollMessage);
     }
@@ -928,11 +923,8 @@ public class DatabaseHelper {
   }
 
   public static PollMessage getOnlyPollMessage(
-      String clientId,
+      String clientId, DateTime now, Class<? extends PollMessage> subType) {
-      DateTime now,
+    return getPollMessages(clientId, now).stream()
-      Class<? extends PollMessage> subType) {
-    return getPollMessages(clientId, now)
-        .stream()
         .filter(subType::isInstance)
         .map(subType::cast)
         .collect(onlyElement());
@@ -957,10 +949,7 @@ public class DatabaseHelper {
     return createDomainRepoId(ObjectifyService.allocateId(), tld);
   }
 
-  /**
+  /** Returns a newly allocated, globally unique contact/host repoId of the format HEX_TLD-ROID. */
-   * Returns a newly allocated, globally unique contact/host repoId of the format
-   * HEX_TLD-ROID.
-   */
   public static String generateNewContactHostRoid() {
     return createRepoId(ObjectifyService.allocateId(), getContactAndHostRoidSuffix());
   }
@@ -1131,8 +1120,7 @@ public class DatabaseHelper {
    */
   public static ImmutableList<HistoryEntry> getHistoryEntriesOfType(
       EppResource resource, final HistoryEntry.Type type) {
-    return getHistoryEntries(resource)
+    return getHistoryEntries(resource).stream()
-        .stream()
         .filter(entry -> entry.getType() == type)
         .collect(toImmutableList());
   }
@@ -1143,7 +1131,7 @@ public class DatabaseHelper {
    */
   public static HistoryEntry getOnlyHistoryEntryOfType(
       EppResource resource, final HistoryEntry.Type type) {
-    List<HistoryEntry> historyEntries =  getHistoryEntriesOfType(resource, type);
+    List<HistoryEntry> historyEntries = getHistoryEntriesOfType(resource, type);
     assertThat(historyEntries).hasSize(1);
     return historyEntries.get(0);
   }
@@ -1151,13 +1139,16 @@ public class DatabaseHelper {
   private static HistoryEntry.Type getHistoryEntryType(EppResource resource) {
     if (resource instanceof ContactResource) {
       return resource.getRepoId() != null
-          ? HistoryEntry.Type.CONTACT_CREATE : HistoryEntry.Type.CONTACT_UPDATE;
+          ? HistoryEntry.Type.CONTACT_CREATE
+          : HistoryEntry.Type.CONTACT_UPDATE;
     } else if (resource instanceof HostResource) {
       return resource.getRepoId() != null
-          ? HistoryEntry.Type.HOST_CREATE : HistoryEntry.Type.HOST_UPDATE;
+          ? HistoryEntry.Type.HOST_CREATE
+          : HistoryEntry.Type.HOST_UPDATE;
     } else if (resource instanceof DomainBase) {
       return resource.getRepoId() != null
-          ? HistoryEntry.Type.DOMAIN_CREATE : HistoryEntry.Type.DOMAIN_UPDATE;
+          ? HistoryEntry.Type.DOMAIN_CREATE
+          : HistoryEntry.Type.DOMAIN_UPDATE;
     } else {
       throw new AssertionError();
     }
@@ -1215,7 +1206,7 @@ public class DatabaseHelper {
     tm().clearSessionCache();
   }
 
-  /** Force the create and update timestamps to get written into the resource. **/
+  /** Force the create and update timestamps to get written into the resource. */
   public static <R> R cloneAndSetAutoTimestamps(final R resource) {
     R result;
     if (tm().isOfy()) {