Convert DomainTAF and DomainFlowUtils to SQL (#1045)

* Convert DomainTAF and DomainFlowUtils to SQL

The only tricky part to this is that the order of entities that we're
saving during the DomainTransferApproveFlow matters -- some entities
have dependencies on others so we need to save the latter first. We
change `entitiesToSave` to be a list to reinforce this.

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

@@ -27,13 +27,13 @@ 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.approvePendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
@@ -132,6 +132,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setBillingTime(now.plus(Registry.get(tld).getTransferGracePeriodLength()))
                     .setParent(historyEntry)
                     .build());
+    ImmutableList.Builder<ImmutableObject> entitiesToSave = new ImmutableList.Builder<>();
+    entitiesToSave.add(historyEntry);
     // If we are within an autorenew grace period, cancel the autorenew billing event and don't
     // increase the registration time, since the transfer subsumes the autorenew's extra year.
     GracePeriod autorenewGrace =
@@ -143,7 +145,7 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
       // then the gaining registrar is not charged for the one year renewal and the losing registrar
       // still needs to be charged for the auto-renew.
       if (billingEvent.isPresent()) {
-        ofy().save().entity(
+        entitiesToSave.add(
             BillingEvent.Cancellation.forGracePeriod(autorenewGrace, historyEntry, targetId));
       }
     }
@@ -190,31 +192,26 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
-                billingEvent.isPresent()
+                billingEvent
-                    ? ImmutableSet.of(
+                    .map(
-                        GracePeriod.forBillingEvent(
+                        oneTime ->
-                            GracePeriodStatus.TRANSFER,
+                            ImmutableSet.of(
-                            existingDomain.getRepoId(),
+                                GracePeriod.forBillingEvent(
-                            billingEvent.get()))
+                                    GracePeriodStatus.TRANSFER,
-                    : ImmutableSet.of())
+                                    existingDomain.getRepoId(),
+                                    oneTime)))
+                    .orElseGet(ImmutableSet::of))
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .build();
     // Create a poll message for the gaining client.
-    PollMessage gainingClientPollMessage = createGainingTransferPollMessage(
+    PollMessage gainingClientPollMessage =
-        targetId,
+        createGainingTransferPollMessage(
-        newDomain.getTransferData(),
+            targetId, newDomain.getTransferData(), newExpirationTime, historyEntry);
-        newExpirationTime,
-        historyEntry);
-    ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(
-        newDomain,
-        historyEntry,
-        autorenewEvent,
-        gainingClientPollMessage,
-        gainingClientAutorenewPollMessage);
     billingEvent.ifPresent(entitiesToSave::add);
-    ofy().save().entities(entitiesToSave.build());
+    entitiesToSave.add(
+        autorenewEvent, gainingClientPollMessage, gainingClientAutorenewPollMessage, newDomain);
+    tm().putAll(entitiesToSave.build());
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingDomain.getTransferData().getServerApproveEntities());

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

@@ -33,9 +33,11 @@ import google.registry.flows.domain.DomainFlowUtils.TldDoesNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.TrailingDashException;
 import google.registry.model.domain.DomainBase;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
+@DualDatabaseTest
 class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase> {
 
   @BeforeEach
@@ -45,12 +47,12 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     persistResource(AppEngineExtension.makeRegistrar1().asBuilder().build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainNameAcceptsValidName() throws EppException {
     assertThat(DomainFlowUtils.validateDomainName("example.tld")).isNotNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_IllegalCharacters() {
     BadDomainNameCharacterException thrown =
         assertThrows(
@@ -62,7 +64,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_DomainNameWithEmptyParts() {
     EmptyDomainNamePartException thrown =
         assertThrows(
@@ -72,7 +74,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_DomainNameWithLessThanTwoParts() {
     BadDomainNamePartsCountException thrown =
         assertThrows(
@@ -84,7 +86,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_invalidTLD() {
     TldDoesNotExistException thrown =
         assertThrows(
@@ -96,7 +98,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_DomainNameIsTooLong() {
     DomainLabelTooLongException thrown =
         assertThrows(
@@ -110,7 +112,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_leadingDash() {
     LeadingDashException thrown =
         assertThrows(
@@ -119,7 +121,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_trailingDash() {
     TrailingDashException thrown =
         assertThrows(
@@ -128,7 +130,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_invalidIDN() {
     InvalidPunycodeException thrown =
         assertThrows(
@@ -140,7 +142,7 @@ class DomainFlowUtilsTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBas
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testValidateDomainName_containsInvalidDashes() {
     DashesInThirdAndFourthException thrown =
         assertThrows(

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

@@ -21,12 +21,12 @@ import static google.registry.model.reporting.DomainTransactionRecord.Transactio
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.assertBillingEventsForResource;
 import static google.registry.testing.DatabaseHelper.createTld;
 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.loadByKey;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DomainBaseSubject.assertAboutDomains;
@@ -70,7 +70,9 @@ import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferResponse.DomainTransferResponse;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.ReplayExtension;
+import google.registry.testing.TestOfyAndSql;
 import java.util.Arrays;
 import java.util.stream.Stream;
 import org.joda.money.Money;
@@ -78,10 +80,10 @@ import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DomainTransferApproveFlow}. */
+@DualDatabaseTest
 class DomainTransferApproveFlowTest
     extends DomainTransferFlowTestCase<DomainTransferApproveFlow, DomainBase> {
 
@@ -197,7 +199,7 @@ class DomainTransferApproveFlowTest
     assertAboutHistoryEntries().that(historyEntryTransferApproved).hasOtherClientId("NewRegistrar");
     assertTransferApproved(domain, originalTransferData);
     assertAboutDomains().that(domain).hasRegistrationExpirationTime(expectedExpirationTime);
-    assertThat(tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
+    assertThat(loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // The poll message (in the future) to the losing registrar for implicit ack should be gone.
     assertThat(getPollMessages(domain, "TheRegistrar", clock.nowUtc().plusMonths(1))).isEmpty();
@@ -346,18 +348,18 @@ class DomainTransferApproveFlowTest
     runFlow();
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     setEppLoader("domain_transfer_approve.xml");
     dryRunFlowAssertResponse(loadFile("domain_transfer_approve_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     doSuccessfulTest("tld", "domain_transfer_approve.xml", "domain_transfer_approve_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_nonDefaultTransferGracePeriod() throws Exception {
     // We have to set up a new domain in a different TLD so that the billing event will be persisted
     // with the new transfer grace period in mind.
@@ -372,7 +374,7 @@ class DomainTransferApproveFlowTest
         "net", "domain_transfer_approve_net.xml", "domain_transfer_approve_response_net.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_domainAuthInfo() throws Exception {
     doSuccessfulTest(
         "tld",
@@ -380,7 +382,7 @@ class DomainTransferApproveFlowTest
         "domain_transfer_approve_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_contactAuthInfo() throws Exception {
     doSuccessfulTest(
         "tld",
@@ -388,7 +390,7 @@ class DomainTransferApproveFlowTest
         "domain_transfer_approve_response.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_autorenewBeforeTransfer() throws Exception {
     domain = reloadResourceByForeignKey();
     DateTime oldExpirationTime = clock.nowUtc().minusDays(1);
@@ -412,7 +414,7 @@ class DomainTransferApproveFlowTest
             .setRecurringEventKey(domain.getAutorenewBillingEvent()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badContactPassword() {
     // Change the contact's password so it does not match the password in the file.
     contact =
@@ -428,7 +430,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badDomainPassword() {
     // Change the domain's password so it does not match the password in the file.
     persistResource(
@@ -443,7 +445,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neverBeenTransferred() {
     changeTransferStatus(null);
     EppException thrown =
@@ -452,7 +454,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientApproved() {
     changeTransferStatus(TransferStatus.CLIENT_APPROVED);
     EppException thrown =
@@ -461,7 +463,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientRejected() {
     changeTransferStatus(TransferStatus.CLIENT_REJECTED);
     EppException thrown =
@@ -470,7 +472,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientCancelled() {
     changeTransferStatus(TransferStatus.CLIENT_CANCELLED);
     EppException thrown =
@@ -479,7 +481,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_serverApproved() {
     changeTransferStatus(TransferStatus.SERVER_APPROVED);
     EppException thrown =
@@ -488,7 +490,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_serverCancelled() {
     changeTransferStatus(TransferStatus.SERVER_CANCELLED);
     EppException thrown =
@@ -497,7 +499,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_gainingClient() {
     setClientIdForFlow("NewRegistrar");
     EppException thrown =
@@ -506,7 +508,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_unrelatedClient() {
     setClientIdForFlow("ClientZ");
     EppException thrown =
@@ -515,7 +517,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_deletedDomain() throws Exception {
     persistResource(domain.asBuilder().setDeletionTime(clock.nowUtc().minusDays(1)).build());
     ResourceDoesNotExistException thrown =
@@ -525,7 +527,7 @@ class DomainTransferApproveFlowTest
     assertThat(thrown).hasMessageThat().contains(String.format("(%s)", getUniqueIdFromCommand()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonexistentDomain() throws Exception {
     deleteTestDomain(domain);
     ResourceDoesNotExistException thrown =
@@ -535,7 +537,7 @@ class DomainTransferApproveFlowTest
     assertThat(thrown).hasMessageThat().contains(String.format("(%s)", getUniqueIdFromCommand()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_notAuthorizedForTld() {
     persistResource(
         loadRegistrar("TheRegistrar").asBuilder().setAllowedTlds(ImmutableSet.of()).build());
@@ -548,7 +550,7 @@ class DomainTransferApproveFlowTest
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserNotAuthorizedForTld() throws Exception {
     persistResource(
         loadRegistrar("TheRegistrar").asBuilder().setAllowedTlds(ImmutableSet.of()).build());
@@ -561,7 +563,7 @@ class DomainTransferApproveFlowTest
   // 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-approve");
@@ -577,7 +579,7 @@ class DomainTransferApproveFlowTest
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannTransactionRecord_noRecordsToCancel() throws Exception {
     setUpGracePeriodDurations();
     clock.advanceOneMilli();
@@ -590,7 +592,7 @@ class DomainTransferApproveFlowTest
                 "tld", clock.nowUtc().plusDays(3), TRANSFER_SUCCESSFUL, 1));
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannTransactionRecord_cancelsPreviousRecords() throws Exception {
     clock.advanceOneMilli();
     setUpGracePeriodDurations();
@@ -618,7 +620,7 @@ class DomainTransferApproveFlowTest
                 "tld", clock.nowUtc().plusDays(3), TRANSFER_SUCCESSFUL, 1));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserExtension_transferPeriodZero() throws Exception {
     domain = reloadResourceByForeignKey();
     DomainTransferData.Builder transferDataBuilder = domain.getTransferData().asBuilder();
@@ -637,7 +639,7 @@ class DomainTransferApproveFlowTest
     assertHistoryEntriesDoNotContainTransferBillingEventsOrGracePeriods();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserExtension_transferPeriodZero_autorenewGraceActive() throws Exception {
     DomainBase domain = reloadResourceByForeignKey();
     VKey<Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();

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

@@ -23,12 +23,14 @@ import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.deleteResource;
 import static google.registry.testing.DatabaseHelper.getBillingEvents;
 import static google.registry.testing.DatabaseHelper.getOnlyHistoryEntryOfType;
+import static google.registry.testing.DatabaseHelper.loadAllOf;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistDomainWithDependentResources;
 import static google.registry.testing.DatabaseHelper.persistDomainWithPendingTransfer;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DomainBaseSubject.assertAboutDomains;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableSet;
@@ -45,6 +47,7 @@ import google.registry.model.host.HostResource;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.testing.AppEngineExtension;
@@ -168,8 +171,9 @@ abstract class DomainTransferFlowTestCase<F extends Flow, R extends EppResource>
    */
   void deleteTestDomain(DomainBase domain) {
     Iterable<BillingEvent> billingEvents = getBillingEvents();
-    Iterable<HistoryEntry> historyEntries = tm().loadAllOf(HistoryEntry.class);
+    Iterable<? extends HistoryEntry> historyEntries =
-    Iterable<PollMessage> pollMessages = tm().loadAllOf(PollMessage.class);
+        HistoryEntryDao.loadAllHistoryObjects(START_OF_TIME, END_OF_TIME);
+    Iterable<PollMessage> pollMessages = loadAllOf(PollMessage.class);
     tm().transact(
             () -> {
               deleteResource(domain);