Add a nomulus command to resave entity with unique id (#656)

* Add a nomulus command to resave entity with unique id

* Remove PollMessage

* Remove logic for PollMessage and resolve comments

* Resolve comments

core/build.gradle

@@ -63,6 +63,8 @@ def dockerIncompatibleTestPatterns = [
     // methods, so we exclude the whole test class.
     "google/registry/tools/params/PathParameterTest.*",
     "google/registry/persistence/PersistenceModuleTest.*",
+    // This test is failing in docker when using Java 11. The cause is unclear.
+    "google/registry/tools/DomainLockUtilsTest.*",
 ]
 
 // Tests that conflict with members of both the main test suite and the

core/src/main/java/google/registry/beam/invoicing/InvoicingPipeline.java

@@ -14,6 +14,8 @@
 
 package google.registry.beam.invoicing;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
 import com.google.auth.oauth2.GoogleCredentials;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey.InvoiceGroupingKeyCoder;
@@ -22,7 +24,9 @@ import google.registry.config.RegistryConfig.Config;
 import google.registry.reporting.billing.BillingModule;
 import google.registry.reporting.billing.GenerateInvoicesAction;
 import google.registry.util.GoogleCredentialsBundle;
+import java.io.File;
 import java.io.Serializable;
+import java.util.Arrays;
 import javax.inject.Inject;
 import org.apache.beam.runners.dataflow.DataflowRunner;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
@@ -111,6 +115,14 @@ public class InvoicingPipeline implements Serializable {
     // So, make sure the credential has write permission to GCS in that project.
     options.setGcpCredential(googleCredentials);
 
+    // The BEAM files-to-stage classpath loader is broken past Java 8, so we do this manually.
+    // See https://issues.apache.org/jira/browse/BEAM-2530
+    options.setFilesToStage(
+        Arrays.stream(System.getProperty("java.class.path").split(File.separator))
+            .map(File::new)
+            .map(File::getPath)
+            .collect(toImmutableList()));
+
     Pipeline p = Pipeline.create(options);
 
     PCollection<BillingEvent> billingEvents =

core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java

@@ -15,17 +15,26 @@
 package google.registry.beam.spec11;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.beam.BeamUtils.getQueryFromFile;
 
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableSet;
+import google.registry.backup.AppEngineEnvironment;
+import google.registry.beam.initsql.Transforms.SerializableSupplier;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.CredentialModule.LocalCredential;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.model.reporting.Spec11ThreatMatch;
+import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
 import google.registry.util.SqlTemplate;
+import java.io.File;
 import java.io.Serializable;
+import java.util.Arrays;
 import javax.inject.Inject;
 import org.apache.beam.runners.dataflow.DataflowRunner;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
@@ -37,6 +46,7 @@ import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
+import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -46,6 +56,7 @@ import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.joda.time.LocalDate;
 import org.joda.time.YearMonth;
+import org.joda.time.format.ISODateTimeFormat;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -86,6 +97,7 @@ public class Spec11Pipeline implements Serializable {
   private final String reportingBucketUrl;
   private final GoogleCredentials googleCredentials;
   private final Retrier retrier;
+  private final SerializableSupplier<JpaTransactionManager> jpaSupplierFactory;
 
   @Inject
   public Spec11Pipeline(
@@ -93,12 +105,14 @@ public class Spec11Pipeline implements Serializable {
       @Config("beamStagingUrl") String beamStagingUrl,
       @Config("spec11TemplateUrl") String spec11TemplateUrl,
       @Config("reportingBucketUrl") String reportingBucketUrl,
+      SerializableSupplier<JpaTransactionManager> jpaSupplierFactory,
       @LocalCredential GoogleCredentialsBundle googleCredentialsBundle,
       Retrier retrier) {
     this.projectId = projectId;
     this.beamStagingUrl = beamStagingUrl;
     this.spec11TemplateUrl = spec11TemplateUrl;
     this.reportingBucketUrl = reportingBucketUrl;
+    this.jpaSupplierFactory = jpaSupplierFactory;
     this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
     this.retrier = retrier;
   }
@@ -143,6 +157,14 @@ public class Spec11Pipeline implements Serializable {
     // So, make sure the credential has write permission to GCS in that project.
     options.setGcpCredential(googleCredentials);
 
+    // The BEAM files-to-stage classpath loader is broken past Java 8, so we do this manually.
+    // See https://issues.apache.org/jira/browse/BEAM-2530
+    options.setFilesToStage(
+        Arrays.stream(System.getProperty("java.class.path").split(File.separator))
+            .map(File::new)
+            .map(File::getPath)
+            .collect(toImmutableList()));
+
     Pipeline p = Pipeline.create(options);
     PCollection<Subdomain> domains =
         p.apply(
@@ -177,12 +199,40 @@ public class Spec11Pipeline implements Serializable {
       EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
       ValueProvider<String> dateProvider) {
 
+    PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
+        domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
+    /* Store ThreatMatch objects in SQL. */
+    subdomainsSql.apply(
+        ParDo.of(
+            new DoFn<KV<Subdomain, ThreatMatch>, Void>() {
+              @ProcessElement
+              public void processElement(ProcessContext context) {
+                // create the Spec11ThreatMatch from Subdomain and ThreatMatch
+                try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+                  Subdomain subdomain = context.element().getKey();
+                  Spec11ThreatMatch threatMatch =
+                      new Spec11ThreatMatch.Builder()
+                          .setThreatTypes(
+                              ImmutableSet.of(
+                                  ThreatType.valueOf(context.element().getValue().threatType())))
+                          .setCheckDate(
+                              LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
+                          .setDomainName(subdomain.domainName())
+                          .setDomainRepoId(subdomain.domainRepoId())
+                          .setRegistrarId(subdomain.registrarId())
+                          .build();
+                  JpaTransactionManager jpaTransactionManager = jpaSupplierFactory.get();
+                  jpaTransactionManager.transact(() -> jpaTransactionManager.saveNew(threatMatch));
+                }
+              }
+            }));
+
     /* Store ThreatMatch objects in JSON. */
     PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =
         domains.apply("Run through SafeBrowsingAPI", ParDo.of(evaluateSafeBrowsingFn));
     subdomainsJson
         .apply(
-            "Map registrar client ID to email/ThreatMatch pair",
+            "Map registrar ID to email/ThreatMatch pair",
             MapElements.into(
                     TypeDescriptors.kvs(
                         TypeDescriptors.strings(), TypeDescriptor.of(EmailAndThreatMatch.class)))

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

@@ -42,6 +42,7 @@ import google.registry.model.EppResource.ForeignKeyedEppResource;
 import google.registry.model.EppResource.ResourceWithTransferData;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.eppcommon.AuthInfo;
@@ -255,7 +256,7 @@ public final class ResourceFlowUtils {
    * @param domain is the domain already projected at approvalTime
    */
   public static DateTime computeExDateForApprovalTime(
-      DomainBase domain, DateTime approvalTime, Period period) {
+      DomainContent domain, DateTime approvalTime, Period period) {
     boolean inAutoRenew = domain.getGracePeriodStatuses().contains(GracePeriodStatus.AUTO_RENEW);
     // inAutoRenew is set to false if the period is zero because a zero-period transfer should not
     // subsume an autorenew.

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

@@ -360,7 +360,8 @@ public class DomainCreateFlow implements TransactionalFlow {
                     command.getNameservers().stream().collect(toImmutableSet()))
             .setStatusValues(statuses.build())
             .setContacts(command.getContacts())
-            .addGracePeriod(GracePeriod.forBillingEvent(GracePeriodStatus.ADD, createBillingEvent))
+            .addGracePeriod(
+                GracePeriod.forBillingEvent(GracePeriodStatus.ADD, repoId, createBillingEvent))
             .build();
     entitiesToSave.add(
         newDomain,

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

@@ -186,14 +186,18 @@ public final class DomainDeleteFlow implements TransactionalFlow {
       DateTime redemptionTime = now.plus(redemptionGracePeriodLength);
       asyncTaskEnqueuer.enqueueAsyncResave(
           existingDomain, now, ImmutableSortedSet.of(redemptionTime, deletionTime));
-      builder.setDeletionTime(deletionTime)
+      builder
+          .setDeletionTime(deletionTime)
           .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
           // Clear out all old grace periods and add REDEMPTION, which does not include a key to a
           // billing event because there isn't one for a domain delete.
-          .setGracePeriods(ImmutableSet.of(GracePeriod.createWithoutBillingEvent(
-              GracePeriodStatus.REDEMPTION,
-              redemptionTime,
-              clientId)));
+          .setGracePeriods(
+              ImmutableSet.of(
+                  GracePeriod.createWithoutBillingEvent(
+                      GracePeriodStatus.REDEMPTION,
+                      existingDomain.getRepoId(),
+                      redemptionTime,
+                      clientId)));
       // Note: The expiration time is unchanged, so if it's before the new deletion time, there will
       // be a "phantom autorenew" where the expiration time advances. No poll message will be
       // produced (since we are ending the autorenew recurrences at "now" below) and the billing

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

@@ -60,7 +60,7 @@ public final class DomainPricingLogic {
    * <p>If {@code allocationToken} is present and the domain is non-premium, that discount will be
    * applied to the first year.
    */
-  public FeesAndCredits getCreatePrice(
+  FeesAndCredits getCreatePrice(
       Registry registry,
       String domainName,
       DateTime dateTime,
@@ -104,8 +104,8 @@ public final class DomainPricingLogic {
 
   /** Returns a new renew price for the pricer. */
   @SuppressWarnings("unused")
-  public FeesAndCredits getRenewPrice(
-      Registry registry, String domainName, DateTime dateTime, int years) throws EppException {
+  FeesAndCredits getRenewPrice(Registry registry, String domainName, DateTime dateTime, int years)
+      throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
     BigDecimal renewCost = domainPrices.getRenewCost().multipliedBy(years).getAmount();
     return customLogic.customizeRenewPrice(
@@ -123,7 +123,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new restore price for the pricer. */
-  public FeesAndCredits getRestorePrice(
+  FeesAndCredits getRestorePrice(
       Registry registry, String domainName, DateTime dateTime, boolean isExpired)
       throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
@@ -147,7 +147,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new transfer price for the pricer. */
-  public FeesAndCredits getTransferPrice(Registry registry, String domainName, DateTime dateTime)
+  FeesAndCredits getTransferPrice(Registry registry, String domainName, DateTime dateTime)
       throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
     return customLogic.customizeTransferPrice(
@@ -168,7 +168,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new update price for the pricer. */
-  public FeesAndCredits getUpdatePrice(Registry registry, String domainName, DateTime dateTime)
+  FeesAndCredits getUpdatePrice(Registry registry, String domainName, DateTime dateTime)
       throws EppException {
     CurrencyUnit currency = registry.getCurrency();
     BaseFee feeOrCredit = Fee.create(zeroInCurrency(currency), FeeType.UPDATE, false);
@@ -191,16 +191,20 @@ public final class DomainPricingLogic {
       throws EppException {
     if (allocationToken.isPresent()
         && allocationToken.get().getDiscountFraction() != 0.0
-        && domainPrices.isPremium()) {
+        && domainPrices.isPremium()
+        && !allocationToken.get().shouldDiscountPremiums()) {
       throw new AllocationTokenInvalidForPremiumNameException();
     }
     Money oneYearCreateCost = domainPrices.getCreateCost();
     Money totalDomainCreateCost = oneYearCreateCost.multipliedBy(years);
-    // If a discount is applicable, apply it only to the first year
+
+    // Apply the allocation token discount, if applicable.
     if (allocationToken.isPresent()) {
+      int discountedYears = Math.min(years, allocationToken.get().getDiscountYears());
       Money discount =
           oneYearCreateCost.multipliedBy(
-              allocationToken.get().getDiscountFraction(), RoundingMode.HALF_UP);
+              discountedYears * allocationToken.get().getDiscountFraction(),
+              RoundingMode.HALF_EVEN);
       totalDomainCreateCost = totalDomainCreateCost.minus(discount);
     }
     return totalDomainCreateCost;
@@ -209,7 +213,7 @@ public final class DomainPricingLogic {
   /** An allocation token was provided that is invalid for premium domains. */
   public static class AllocationTokenInvalidForPremiumNameException
       extends CommandUseErrorException {
-    public AllocationTokenInvalidForPremiumNameException() {
+    AllocationTokenInvalidForPremiumNameException() {
       super("A nonzero discount code cannot be applied to premium domains");
     }
   }

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

@@ -183,7 +183,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
             .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .addGracePeriod(
-                GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, explicitRenewEvent))
+                GracePeriod.forBillingEvent(
+                    GracePeriodStatus.RENEW, existingDomain.getRepoId(), explicitRenewEvent))
             .build();
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(

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

@@ -192,7 +192,10 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             .setGracePeriods(
                 billingEvent.isPresent()
                     ? ImmutableSet.of(
-                        GracePeriod.forBillingEvent(GracePeriodStatus.TRANSFER, billingEvent.get()))
+                        GracePeriod.forBillingEvent(
+                            GracePeriodStatus.TRANSFER,
+                            existingDomain.getRepoId(),
+                            billingEvent.get()))
                     : ImmutableSet.of())
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)

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

@@ -14,7 +14,6 @@
 
 package google.registry.model;
 
-import com.google.common.annotations.VisibleForTesting;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.MappedSuperclass;
@@ -30,6 +29,7 @@ import javax.xml.bind.annotation.XmlTransient;
  */
 @MappedSuperclass
 public abstract class BackupGroupRoot extends ImmutableObject {
+
   /**
    * An automatically managed timestamp of when this object was last written to Datastore.
    *
@@ -40,7 +40,6 @@ public abstract class BackupGroupRoot extends ImmutableObject {
   // Prevents subclasses from unexpectedly accessing as property (e.g., HostResource), which would
   // require an unnecessary non-private setter method.
   @Access(AccessType.FIELD)
-  @VisibleForTesting
   UpdateAutoTimestamp updateTimestamp = UpdateAutoTimestamp.create(null);
 
   /** Get the {@link UpdateAutoTimestamp} for this entity. */

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -347,27 +347,11 @@ public abstract class BillingEvent extends ImmutableObject
 
     @Override
     public VKey<OneTime> createVKey() {
-      return VKey.create(
-          getClass(),
-          parent.getParent().getName() + "/" + parent.getId() + "/" + getId(),
-          Key.create(this));
+      return VKey.create(getClass(), getId(), Key.create(this));
     }
 
     public static VKey<OneTime> createVKey(Key<OneTime> key) {
-      // TODO(b/159207551): As it stands, the SQL key generated here doesn't mesh with the primary
-      // key type for the table, which is a single long integer.
-      if (key == null) {
-        return null;
-      }
-      Key parent = key.getParent();
-      Key grandparent = (parent != null) ? parent.getParent() : null;
-      String path =
-          (grandparent != null ? grandparent.getName() : "")
-              + "/"
-              + (parent != null ? parent.getId() : "")
-              + "/"
-              + key.getId();
-      return VKey.create(OneTime.class, path, key);
+      return VKey.create(OneTime.class, key.getId(), key);
     }
 
     @Override
@@ -505,21 +489,11 @@ public abstract class BillingEvent extends ImmutableObject
 
     @Override
     public VKey<Recurring> createVKey() {
-      return VKey.create(
-          getClass(),
-          parent.getParent().getName() + "/" + parent.getId() + "/" + getId(),
-          Key.create(this));
+      return VKey.create(getClass(), getId(), Key.create(this));
     }
 
     public static VKey<Recurring> createVKey(Key<Recurring> key) {
-      // TODO(b/159207551): As it stands, the SQL key generated here doesn't mesh with the primary
-      // key type for the table, which is a single long integer.
-      if (key == null) {
-        return null;
-      }
-      String path =
-          key.getParent().getParent().getName() + "/" + key.getParent().getId() + "/" + key.getId();
-      return VKey.create(Recurring.class, path, key);
+      return VKey.create(Recurring.class, key.getId(), key);
     }
 
     @Override
@@ -639,21 +613,11 @@ public abstract class BillingEvent extends ImmutableObject
 
     @Override
     public VKey<Cancellation> createVKey() {
-      return VKey.create(
-          getClass(),
-          parent.getParent().getName() + "/" + parent.getId() + "/" + getId(),
-          Key.create(this));
+      return VKey.create(getClass(), getId(), Key.create(this));
     }
 
     public static VKey<Cancellation> createVKey(Key<Cancellation> key) {
-      // TODO(b/159207551): As it stands, the SQL key generated here doesn't mesh with the primary
-      // key type for the table, which is a single long integer.
-      if (key == null) {
-        return null;
-      }
-      String path =
-          key.getParent().getParent().getName() + "/" + key.getParent().getId() + "/" + key.getId();
-      return VKey.create(Cancellation.class, path, key);
+      return VKey.create(Cancellation.class, key.getId(), key);
     }
 
     @Override
@@ -733,21 +697,11 @@ public abstract class BillingEvent extends ImmutableObject
 
     @Override
     public VKey<Modification> createVKey() {
-      return VKey.create(
-          getClass(),
-          parent.getParent().getName() + "/" + parent.getId() + "/" + getId(),
-          Key.create(this));
+      return VKey.create(getClass(), getId(), Key.create(this));
     }
 
     public static VKey<Modification> createVKey(Key<Modification> key) {
-      // TODO(b/159207551): As it stands, the SQL key generated here doesn't mesh with the primary
-      // key type for the table, which is a single long integer.
-      if (key == null) {
-        return null;
-      }
-      String path =
-          key.getParent().getParent().getName() + "/" + key.getParent().getId() + "/" + key.getId();
-      return VKey.create(Modification.class, path, key);
+      return VKey.create(Modification.class, key.getId(), key);
     }
 
     /**

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

@@ -25,8 +25,8 @@ import javax.persistence.Entity;
  * A persisted history entry representing an EPP modification to a contact.
  *
  * <p>In addition to the general history fields (e.g. action time, registrar ID) we also persist a
- * copy of the host entity at this point in time. We persist a raw {@link ContactBase} so that the
- * foreign-keyed fields in that class can refer to this object.
+ * copy of the contact entity at this point in time. We persist a raw {@link ContactBase} so that
+ * the foreign-keyed fields in that class can refer to this object.
  */
 @Entity
 @javax.persistence.Table(

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

@@ -40,7 +40,7 @@ import org.joda.time.DateTime;
       @javax.persistence.Index(columnList = "creationTime"),
       @javax.persistence.Index(columnList = "currentSponsorRegistrarId"),
       @javax.persistence.Index(columnList = "deletionTime"),
-      @javax.persistence.Index(columnList = "contactId", unique = true),
+      @javax.persistence.Index(columnList = "contactId"),
       @javax.persistence.Index(columnList = "searchName")
     })
 @ExternalMessagingName("contact")

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

@@ -14,79 +14,28 @@
 
 package google.registry.model.domain;
 
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Strings.emptyToNull;
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static com.google.common.collect.ImmutableSortedSet.toImmutableSortedSet;
-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.EppResourceUtils.projectResourceOntoBuilderAtTime;
-import static google.registry.model.EppResourceUtils.setAutomaticTransferSuccessProperties;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.util.CollectionUtils.forceEmptyToNull;
-import static google.registry.util.CollectionUtils.nullToEmpty;
-import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
-import static google.registry.util.CollectionUtils.union;
-import static google.registry.util.DateTimeUtils.END_OF_TIME;
-import static google.registry.util.DateTimeUtils.earliestOf;
-import static google.registry.util.DateTimeUtils.isBeforeOrAt;
-import static google.registry.util.DateTimeUtils.leapSafeAddYears;
-import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
-import static google.registry.util.DomainNameUtils.getTldFromDomainName;
-import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
-import com.google.common.collect.Ordering;
-import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
-import com.googlecode.objectify.annotation.Ignore;
-import com.googlecode.objectify.annotation.IgnoreSave;
-import com.googlecode.objectify.annotation.Index;
-import com.googlecode.objectify.annotation.OnLoad;
-import com.googlecode.objectify.condition.IfNull;
-import google.registry.flows.ResourceFlowUtils;
 import google.registry.model.EppResource;
 import google.registry.model.EppResource.ForeignKeyedEppResource;
-import google.registry.model.EppResource.ResourceWithTransferData;
 import google.registry.model.annotations.ExternalMessagingName;
 import google.registry.model.annotations.ReportedOn;
-import google.registry.model.billing.BillingEvent;
-import google.registry.model.contact.ContactResource;
-import google.registry.model.domain.DesignatedContact.Type;
-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.StatusValue;
 import google.registry.model.host.HostResource;
-import google.registry.model.poll.PollMessage;
-import google.registry.model.registry.Registry;
-import google.registry.model.transfer.DomainTransferData;
-import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithStringVKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
-import google.registry.util.CollectionUtils;
-import java.util.HashSet;
-import java.util.Objects;
-import java.util.Optional;
 import java.util.Set;
-import java.util.function.Predicate;
-import javax.annotation.Nullable;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
-import javax.persistence.AttributeOverride;
-import javax.persistence.AttributeOverrides;
+import javax.persistence.CascadeType;
 import javax.persistence.Column;
 import javax.persistence.ElementCollection;
-import javax.persistence.Embedded;
+import javax.persistence.FetchType;
+import javax.persistence.JoinColumn;
 import javax.persistence.JoinTable;
-import javax.persistence.PostLoad;
-import javax.persistence.Transient;
+import javax.persistence.OneToMany;
 import org.joda.time.DateTime;
-import org.joda.time.Interval;
 
 /**
  * A persistable domain resource including mutable and non-mutable fields.
@@ -112,197 +61,8 @@ import org.joda.time.Interval;
 @WithStringVKey
 @ExternalMessagingName("domain")
 @Access(AccessType.FIELD)
-public class DomainBase extends EppResource
-    implements DatastoreAndSqlEntity,
-        ForeignKeyedEppResource,
-        ResourceWithTransferData<DomainTransferData> {
-
-  /** The max number of years that a domain can be registered for, as set by ICANN policy. */
-  public static final int MAX_REGISTRATION_YEARS = 10;
-
-  /** Status values which prohibit DNS information from being published. */
-  private static final ImmutableSet<StatusValue> DNS_PUBLISHING_PROHIBITED_STATUSES =
-      ImmutableSet.of(
-          StatusValue.CLIENT_HOLD,
-          StatusValue.INACTIVE,
-          StatusValue.PENDING_DELETE,
-          StatusValue.SERVER_HOLD);
-
-  /**
-   * Fully qualified domain name (puny-coded), which serves as the foreign key for this domain.
-   *
-   * <p>This is only unique in the sense that for any given lifetime specified as the time range
-   * from (creationTime, deletionTime) there can only be one domain in Datastore with this name.
-   * However, there can be many domains with the same name and non-overlapping lifetimes.
-   *
-   * @invariant fullyQualifiedDomainName == fullyQualifiedDomainName.toLowerCase(Locale.ENGLISH)
-   */
-  // TODO(b/158858642): Rename this to domainName when we are off Datastore
-  @Column(name = "domainName")
-  @Index
-  String fullyQualifiedDomainName;
-
-  /** The top level domain this is under, dernormalized from {@link #fullyQualifiedDomainName}. */
-  @Index
-  String tld;
-
-  /** References to hosts that are the nameservers for the domain. */
-  @Index
-  @ElementCollection
-  @JoinTable(name = "DomainHost")
-  Set<VKey<HostResource>> nsHosts;
-
-  /**
-   * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.
-   *
-   * <p>These are stored in one field so that we can query across all contacts at once.
-   */
-  @Transient Set<DesignatedContact> allContacts;
-
-  /**
-   * Contacts as they are stored in cloud SQL.
-   *
-   * <p>This information is duplicated in allContacts, and must be kept in sync with it.
-   */
-  @Ignore VKey<ContactResource> adminContact;
-
-  @Ignore VKey<ContactResource> billingContact;
-  @Ignore VKey<ContactResource> techContact;
-  @Ignore VKey<ContactResource> registrantContact;
-
-  /** Authorization info (aka transfer secret) of the domain. */
-  @Embedded
-  @AttributeOverrides({
-    @AttributeOverride(name = "pw.value", column = @Column(name = "auth_info_value")),
-    @AttributeOverride(name = "pw.repoId", column = @Column(name = "auth_info_repo_id")),
-  })
-  DomainAuthInfo authInfo;
-
-  /**
-   * Data used to construct DS records for this domain.
-   *
-   * <p>This is {@literal @}XmlTransient because it needs to be returned under the "extension" tag
-   * of an info response rather than inside the "infData" tag.
-   */
-  @Transient Set<DelegationSignerData> dsData;
-
-  /**
-   * The claims notice supplied when this application or domain was created, if there was one. It's
-   * {@literal @}XmlTransient because it's not returned in an info response.
-   */
-  @IgnoreSave(IfNull.class)
-  @Embedded
-  @AttributeOverrides({
-    @AttributeOverride(name = "noticeId.tcnId", column = @Column(name = "launch_notice_tcn_id")),
-    @AttributeOverride(
-        name = "noticeId.validatorId",
-        column = @Column(name = "launch_notice_validator_id")),
-    @AttributeOverride(
-        name = "expirationTime",
-        column = @Column(name = "launch_notice_expiration_time")),
-    @AttributeOverride(
-        name = "acceptedTime",
-        column = @Column(name = "launch_notice_accepted_time")),
-  })
-  LaunchNotice launchNotice;
-
-  /**
-   * Name of first IDN table associated with TLD that matched the characters in this domain label.
-   *
-   * @see google.registry.tldconfig.idn.IdnLabelValidator#findValidIdnTableForTld
-   */
-  @IgnoreSave(IfNull.class)
-  String idnTableName;
-
-  /** Fully qualified host names of this domain's active subordinate hosts. */
-  Set<String> subordinateHosts;
-
-  /** When this domain's registration will expire. */
-  DateTime registrationExpirationTime;
-
-  /**
-   * The poll message associated with this domain being deleted.
-   *
-   * <p>This field should be null if the domain is not in pending delete. If it is, the field should
-   * refer to a {@link PollMessage} timed to when the domain is fully deleted. If the domain is
-   * restored, the message should be deleted.
-   */
-  @Column(name = "deletion_poll_message_id")
-  VKey<PollMessage.OneTime> deletePollMessage;
-
-  /**
-   * The recurring billing event associated with this domain's autorenewals.
-   *
-   * <p>The recurrence should be open ended unless the domain is in pending delete or fully deleted,
-   * in which case it should be closed at the time the delete was requested. Whenever the domain's
-   * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
-   * should be created, and this field should be updated to point to the new one.
-   */
-  @Column(name = "billing_recurrence_id")
-  VKey<BillingEvent.Recurring> autorenewBillingEvent;
-
-  /**
-   * The recurring poll message associated with this domain's autorenewals.
-   *
-   * <p>The recurrence should be open ended unless the domain is in pending delete or fully deleted,
-   * in which case it should be closed at the time the delete was requested. Whenever the domain's
-   * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
-   * should be created, and this field should be updated to point to the new one.
-   */
-  @Column(name = "autorenew_poll_message_id")
-  VKey<PollMessage.Autorenew> autorenewPollMessage;
-
-  /** The unexpired grace periods for this domain (some of which may not be active yet). */
-  @Transient @ElementCollection Set<GracePeriod> gracePeriods;
-
-  /**
-   * The id of the signed mark that was used to create this domain in sunrise.
-   *
-   * <p>Will only be populated for domains created in sunrise.
-   */
-  @IgnoreSave(IfNull.class)
-  String smdId;
-
-  /** Data about any pending or past transfers on this domain. */
-  DomainTransferData transferData;
-
-  /**
-   * The time that this resource was last transferred.
-   *
-   * <p>Can be null if the resource has never been transferred.
-   */
-  DateTime lastTransferTime;
-
-  @OnLoad
-  void load() {
-    // Reconstitute all of the contacts so that they have VKeys.
-    allContacts =
-        allContacts.stream().map(contact -> contact.reconstitute()).collect(toImmutableSet());
-    setContactFields(allContacts, true);
-  }
-
-  @PostLoad
-  void postLoad() {
-    // Reconstitute the contact list.
-    ImmutableSet.Builder<DesignatedContact> contactsBuilder =
-        new ImmutableSet.Builder<DesignatedContact>();
-
-    if (registrantContact != null) {
-      contactsBuilder.add(
-          DesignatedContact.create(DesignatedContact.Type.REGISTRANT, registrantContact));
-    }
-    if (billingContact != null) {
-      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.BILLING, billingContact));
-    }
-    if (techContact != null) {
-      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.TECH, techContact));
-    }
-    if (adminContact != null) {
-      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.ADMIN, adminContact));
-    }
-
-    allContacts = contactsBuilder.build();
-  }
+public class DomainBase extends DomainContent
+    implements DatastoreAndSqlEntity, ForeignKeyedEppResource {
 
   @Override
   @javax.persistence.Id
@@ -311,326 +71,23 @@ public class DomainBase extends EppResource
     return super.getRepoId();
   }
 
-  public ImmutableSet<String> getSubordinateHosts() {
-    return nullToEmptyImmutableCopy(subordinateHosts);
+  @ElementCollection
+  @JoinTable(name = "DomainHost")
+  @Access(AccessType.PROPERTY)
+  @Column(name = "host_repo_id")
+  public Set<VKey<HostResource>> getNsHosts() {
+    return super.nsHosts;
   }
 
-  public DateTime getRegistrationExpirationTime() {
-    return registrationExpirationTime;
-  }
-
-  public VKey<PollMessage.OneTime> getDeletePollMessage() {
-    return deletePollMessage;
-  }
-
-  public VKey<BillingEvent.Recurring> getAutorenewBillingEvent() {
-    return autorenewBillingEvent;
-  }
-
-  public VKey<PollMessage.Autorenew> getAutorenewPollMessage() {
-    return autorenewPollMessage;
-  }
-
-  public ImmutableSet<GracePeriod> getGracePeriods() {
-    return nullToEmptyImmutableCopy(gracePeriods);
-  }
-
-  public String getSmdId() {
-    return smdId;
-  }
-
-  @Override
-  public DomainTransferData getTransferData() {
-    return Optional.ofNullable(transferData).orElse(DomainTransferData.EMPTY);
-  }
-
-  @Override
-  public DateTime getLastTransferTime() {
-    return lastTransferTime;
-  }
-
-  @Override
-  public String getForeignKey() {
-    return fullyQualifiedDomainName;
-  }
-
-  public String getDomainName() {
-    return fullyQualifiedDomainName;
-  }
-
-  public ImmutableSet<DelegationSignerData> getDsData() {
-    return nullToEmptyImmutableCopy(dsData);
-  }
-
-  public LaunchNotice getLaunchNotice() {
-    return launchNotice;
-  }
-
-  public String getIdnTableName() {
-    return idnTableName;
-  }
-
-  public ImmutableSet<VKey<HostResource>> getNameservers() {
-    return nullToEmptyImmutableCopy(nsHosts);
-  }
-
-  public final String getCurrentSponsorClientId() {
-    return getPersistedCurrentSponsorClientId();
-  }
-
-  /** Returns true if DNS information should be published for the given domain. */
-  public boolean shouldPublishToDns() {
-    return intersection(getStatusValues(), DNS_PUBLISHING_PROHIBITED_STATUSES).isEmpty();
-  }
-
-  /**
-   * Returns the Registry Grace Period Statuses for this domain.
-   *
-   * <p>This collects all statuses from the domain's {@link GracePeriod} entries and also adds the
-   * PENDING_DELETE status if needed.
-   */
-  public ImmutableSet<GracePeriodStatus> getGracePeriodStatuses() {
-    Set<GracePeriodStatus> gracePeriodStatuses = new HashSet<>();
-    for (GracePeriod gracePeriod : getGracePeriods()) {
-      gracePeriodStatuses.add(gracePeriod.getType());
-    }
-    if (getStatusValues().contains(StatusValue.PENDING_DELETE)
-        && !gracePeriodStatuses.contains(GracePeriodStatus.REDEMPTION)) {
-      gracePeriodStatuses.add(GracePeriodStatus.PENDING_DELETE);
-    }
-    return ImmutableSet.copyOf(gracePeriodStatuses);
-  }
-
-  /** Returns the subset of grace periods having the specified type. */
-  public ImmutableSet<GracePeriod> getGracePeriodsOfType(GracePeriodStatus gracePeriodType) {
-    ImmutableSet.Builder<GracePeriod> builder = new ImmutableSet.Builder<>();
-    for (GracePeriod gracePeriod : getGracePeriods()) {
-      if (gracePeriod.getType() == gracePeriodType) {
-        builder.add(gracePeriod);
-      }
-    }
-    return builder.build();
-  }
-
-  /**
-   * The logic in this method, which handles implicit server approval of transfers, very closely
-   * parallels the logic in {@code DomainTransferApproveFlow} which handles explicit client
-   * approvals.
-   */
-  @Override
-  public DomainBase cloneProjectedAtTime(final DateTime now) {
-
-    DomainTransferData transferData = getTransferData();
-    DateTime transferExpirationTime = transferData.getPendingTransferExpirationTime();
-
-    // If there's a pending transfer that has expired, handle it.
-    if (TransferStatus.PENDING.equals(transferData.getTransferStatus())
-        && isBeforeOrAt(transferExpirationTime, now)) {
-      // Project until just before the transfer time. This will handle the case of an autorenew
-      // before the transfer was even requested or during the request period.
-      // If the transfer time is precisely the moment that the domain expires, there will not be an
-      // autorenew billing event (since we end the recurrence at transfer time and recurrences are
-      // exclusive of their ending), and we can just proceed with the transfer.
-      DomainBase domainAtTransferTime =
-          cloneProjectedAtTime(transferExpirationTime.minusMillis(1));
-
-      DateTime expirationDate = transferData.getTransferredRegistrationExpirationTime();
-      if (expirationDate == null) {
-        // Extend the registration by the correct number of years from the expiration time
-        // that was current on the domain right before the transfer, capped at 10 years from
-        // the moment of the transfer.
-        expirationDate =
-            ResourceFlowUtils.computeExDateForApprovalTime(
-                domainAtTransferTime, transferExpirationTime, transferData.getTransferPeriod());
-      }
-      // If we are within an autorenew grace period, the transfer will subsume the autorenew. There
-      // will already be a cancellation written in advance by the transfer request flow, so we don't
-      // need to worry about billing, but we do need to cancel out the expiration time increase.
-      // The transfer period saved in the transfer data will be one year, unless the superuser
-      // extension set the transfer period to zero.
-      // Set the expiration, autorenew events, and grace period for the transfer. (Transfer ends
-      // all other graces).
-      Builder builder =
-          domainAtTransferTime
-              .asBuilder()
-              .setRegistrationExpirationTime(expirationDate)
-              // Set the speculatively-written new autorenew events as the domain's autorenew
-              // events.
-              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
-              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
-      if (transferData.getTransferPeriod().getValue() == 1) {
-        // Set the grace period using a key to the prescheduled transfer billing event.  Not using
-        // GracePeriod.forBillingEvent() here in order to avoid the actual Datastore fetch.
-        builder.setGracePeriods(
-            ImmutableSet.of(
-                GracePeriod.create(
-                    GracePeriodStatus.TRANSFER,
-                    transferExpirationTime.plus(
-                        Registry.get(getTld()).getTransferGracePeriodLength()),
-                    transferData.getGainingClientId(),
-                    transferData.getServerApproveBillingEvent())));
-      } else {
-        // There won't be a billing event, so we don't need a grace period
-        builder.setGracePeriods(ImmutableSet.of());
-      }
-      // Set all remaining transfer properties.
-      setAutomaticTransferSuccessProperties(builder, transferData);
-      builder
-          .setLastEppUpdateTime(transferExpirationTime)
-          .setLastEppUpdateClientId(transferData.getGainingClientId());
-      // Finish projecting to now.
-      return builder.build().cloneProjectedAtTime(now);
-    }
-
-    Optional<DateTime> newLastEppUpdateTime = Optional.empty();
-
-    // There is no transfer. Do any necessary autorenews for active domains.
-
-    Builder builder = asBuilder();
-    if (isBeforeOrAt(registrationExpirationTime, now) && END_OF_TIME.equals(getDeletionTime())) {
-      // Autorenew by the number of years between the old expiration time and now.
-      DateTime lastAutorenewTime = leapSafeAddYears(
-          registrationExpirationTime,
-          new Interval(registrationExpirationTime, now).toPeriod().getYears());
-      DateTime newExpirationTime  = lastAutorenewTime.plusYears(1);
-      builder
-          .setRegistrationExpirationTime(newExpirationTime)
-          .addGracePeriod(
-              GracePeriod.createForRecurring(
-                  GracePeriodStatus.AUTO_RENEW,
-                  lastAutorenewTime.plus(Registry.get(getTld()).getAutoRenewGracePeriodLength()),
-                  getCurrentSponsorClientId(),
-                  autorenewBillingEvent));
-      newLastEppUpdateTime = Optional.of(lastAutorenewTime);
-    }
-
-    // Remove any grace periods that have expired.
-    DomainBase almostBuilt = builder.build();
-    builder = almostBuilt.asBuilder();
-    for (GracePeriod gracePeriod : almostBuilt.getGracePeriods()) {
-      if (isBeforeOrAt(gracePeriod.getExpirationTime(), now)) {
-        builder.removeGracePeriod(gracePeriod);
-        if (!newLastEppUpdateTime.isPresent()
-            || isBeforeOrAt(newLastEppUpdateTime.get(), gracePeriod.getExpirationTime())) {
-          newLastEppUpdateTime = Optional.of(gracePeriod.getExpirationTime());
-        }
-      }
-    }
-
-    // It is possible that the lastEppUpdateClientId is different from current sponsor client
-    // id, so we have to do the comparison instead of having one variable just storing the most
-    // recent time.
-    if (newLastEppUpdateTime.isPresent()) {
-      if (getLastEppUpdateTime() == null
-          || newLastEppUpdateTime.get().isAfter(getLastEppUpdateTime())) {
-        builder
-            .setLastEppUpdateTime(newLastEppUpdateTime.get())
-            .setLastEppUpdateClientId(getCurrentSponsorClientId());
-      }
-    }
-
-    // Handle common properties like setting or unsetting linked status. This also handles the
-    // general case of pending transfers for other resource types, but since we've always handled
-    // a pending transfer by this point that's a no-op for domains.
-    projectResourceOntoBuilderAtTime(almostBuilt, builder, now);
-    return builder.build();
-  }
-
-  /** Return what the expiration time would be if the given number of years were added to it. */
-  public static DateTime extendRegistrationWithCap(
-      DateTime now,
-      DateTime currentExpirationTime,
-      @Nullable Integer extendedRegistrationYears) {
-    // We must cap registration at the max years (aka 10), even if that truncates the last year.
-    return earliestOf(
-        leapSafeAddYears(
-            currentExpirationTime,
-            Optional.ofNullable(extendedRegistrationYears).orElse(0)),
-        leapSafeAddYears(now, MAX_REGISTRATION_YEARS));
-  }
-
-  /** Loads and returns the fully qualified host names of all linked nameservers. */
-  public ImmutableSortedSet<String> loadNameserverHostNames() {
-    return ofy()
-        .load()
-        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet()))
-        .values()
-        .stream()
-        .map(HostResource::getHostName)
-        .collect(toImmutableSortedSet(Ordering.natural()));
-  }
-
-  /** A key to the registrant who registered this domain. */
-  public VKey<ContactResource> getRegistrant() {
-    return registrantContact;
-  }
-
-  public VKey<ContactResource> getAdminContact() {
-    return adminContact;
-  }
-
-  public VKey<ContactResource> getBillingContact() {
-    return billingContact;
-  }
-
-  public VKey<ContactResource> getTechContact() {
-    return techContact;
-  }
-
-  /** Associated contacts for the domain (other than registrant). */
-  public ImmutableSet<DesignatedContact> getContacts() {
-    return nullToEmpty(allContacts)
-        .stream()
-        .filter(IS_REGISTRANT.negate())
-        .collect(toImmutableSet());
-  }
-
-  public DomainAuthInfo getAuthInfo() {
-    return authInfo;
-  }
-
-  /** Returns all referenced contacts from this domain or application. */
-  public ImmutableSet<VKey<ContactResource>> getReferencedContacts() {
-    return nullToEmptyImmutableCopy(allContacts)
-        .stream()
-        .map(DesignatedContact::getContactKey)
-        .filter(Objects::nonNull)
-        .collect(toImmutableSet());
-  }
-
-  public String getTld() {
-    return tld;
-  }
-
-  /**
-   * Sets the individual contact fields from {@code contacts}.
-   *
-   * <p>The registrant field is only set if {@code includeRegistrant} is true, as this field needs
-   * to be set in some circumstances but not in others.
-   */
-  private void setContactFields(Set<DesignatedContact> contacts, boolean includeRegistrant) {
-
-    // Set the individual contact fields.
-    for (DesignatedContact contact : contacts) {
-      switch (contact.getType()) {
-        case BILLING:
-          billingContact = contact.getContactKey();
-          break;
-        case TECH:
-          techContact = contact.getContactKey();
-          break;
-        case ADMIN:
-          adminContact = contact.getContactKey();
-          break;
-        case REGISTRANT:
-          if (includeRegistrant) {
-            registrantContact = contact.getContactKey();
-          }
-          break;
-        default:
-          throw new IllegalArgumentException("Unknown contact resource type: " + contact.getType());
-      }
-    }
+  @Access(AccessType.PROPERTY)
+  @OneToMany(
+      cascade = {CascadeType.ALL},
+      fetch = FetchType.EAGER,
+      orphanRemoval = true)
+  @JoinColumn(name = "domainRepoId", referencedColumnName = "repoId")
+  @SuppressWarnings("UnusedMethod")
+  private Set<GracePeriod> getInternalGracePeriods() {
+    return gracePeriods;
   }
 
   @Override
@@ -638,9 +95,14 @@ public class DomainBase extends EppResource
     return VKey.create(DomainBase.class, getRepoId(), Key.create(this));
   }
 
-  /** Predicate to determine if a given {@link DesignatedContact} is the registrant. */
-  private static final Predicate<DesignatedContact> IS_REGISTRANT =
-      (DesignatedContact contact) -> DesignatedContact.Type.REGISTRANT.equals(contact.type);
+  @Override
+  public DomainBase cloneProjectedAtTime(final DateTime now) {
+    return cloneDomainProjectedAtTime(this, now);
+  }
+
+  public static VKey<DomainBase> createVKey(Key key) {
+    return VKey.create(DomainBase.class, key.getName(), key);
+  }
 
   /** An override of {@link EppResource#asBuilder} with tighter typing. */
   @Override
@@ -649,199 +111,12 @@ public class DomainBase extends EppResource
   }
 
   /** A builder for constructing {@link DomainBase}, since it is immutable. */
-  public static class Builder extends EppResource.Builder<DomainBase, Builder>
-      implements BuilderWithTransferData<DomainTransferData, Builder> {
+  public static class Builder extends DomainContent.Builder<DomainBase, Builder> {
 
     public Builder() {}
 
     Builder(DomainBase instance) {
       super(instance);
     }
-
-    @Override
-    public DomainBase build() {
-      DomainBase instance = getInstance();
-      // If TransferData is totally empty, set it to null.
-      if (DomainTransferData.EMPTY.equals(getInstance().transferData)) {
-        setTransferData(null);
-      }
-      // A DomainBase has status INACTIVE if there are no nameservers.
-      if (getInstance().getNameservers().isEmpty()) {
-        addStatusValue(StatusValue.INACTIVE);
-      } else {  // There are nameservers, so make sure INACTIVE isn't there.
-        removeStatusValue(StatusValue.INACTIVE);
-      }
-
-      checkArgumentNotNull(emptyToNull(instance.fullyQualifiedDomainName), "Missing domainName");
-      if (instance.getRegistrant() == null
-          && instance.allContacts.stream().anyMatch(IS_REGISTRANT)) {
-        throw new IllegalArgumentException("registrant is null but is in allContacts");
-      }
-      checkArgumentNotNull(instance.getRegistrant(), "Missing registrant");
-      instance.tld = getTldFromDomainName(instance.fullyQualifiedDomainName);
-      return super.build();
-    }
-
-    public Builder setDomainName(String domainName) {
-      checkArgument(
-          domainName.equals(canonicalizeDomainName(domainName)),
-          "Domain name must be in puny-coded, lower-case form");
-      getInstance().fullyQualifiedDomainName = domainName;
-      return thisCastToDerived();
-    }
-
-    public Builder setDsData(ImmutableSet<DelegationSignerData> dsData) {
-      getInstance().dsData = dsData;
-      return thisCastToDerived();
-    }
-
-    public Builder setRegistrant(VKey<ContactResource> registrant) {
-      // Replace the registrant contact inside allContacts.
-      getInstance().allContacts = union(
-          getInstance().getContacts(),
-          DesignatedContact.create(Type.REGISTRANT, checkArgumentNotNull(registrant)));
-
-      // Set the registrant field specifically.
-      getInstance().registrantContact = registrant;
-      return thisCastToDerived();
-    }
-
-    public Builder setAuthInfo(DomainAuthInfo authInfo) {
-      getInstance().authInfo = authInfo;
-      return thisCastToDerived();
-    }
-
-    public Builder setNameservers(VKey<HostResource> nameserver) {
-      getInstance().nsHosts = ImmutableSet.of(nameserver);
-      return thisCastToDerived();
-    }
-
-    public Builder setNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
-      getInstance().nsHosts = forceEmptyToNull(nameservers);
-      return thisCastToDerived();
-    }
-
-    public Builder addNameserver(VKey<HostResource> nameserver) {
-      return addNameservers(ImmutableSet.of(nameserver));
-    }
-
-    public Builder addNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
-      return setNameservers(
-          ImmutableSet.copyOf(union(getInstance().getNameservers(), nameservers)));
-    }
-
-    public Builder removeNameserver(VKey<HostResource> nameserver) {
-      return removeNameservers(ImmutableSet.of(nameserver));
-    }
-
-    public Builder removeNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
-      return setNameservers(
-          ImmutableSet.copyOf(difference(getInstance().getNameservers(), nameservers)));
-    }
-
-    public Builder setContacts(DesignatedContact contact) {
-      return setContacts(ImmutableSet.of(contact));
-    }
-
-    public Builder setContacts(ImmutableSet<DesignatedContact> contacts) {
-      checkArgument(contacts.stream().noneMatch(IS_REGISTRANT), "Registrant cannot be a contact");
-
-      // Replace the non-registrant contacts inside allContacts.
-      getInstance().allContacts =
-          Streams.concat(
-                  nullToEmpty(getInstance().allContacts).stream().filter(IS_REGISTRANT),
-                  contacts.stream())
-              .collect(toImmutableSet());
-
-      // Set the individual fields.
-      getInstance().setContactFields(contacts, false);
-      return thisCastToDerived();
-    }
-
-    public Builder addContacts(ImmutableSet<DesignatedContact> contacts) {
-      return setContacts(ImmutableSet.copyOf(union(getInstance().getContacts(), contacts)));
-    }
-
-    public Builder removeContacts(ImmutableSet<DesignatedContact> contacts) {
-      return setContacts(ImmutableSet.copyOf(difference(getInstance().getContacts(), contacts)));
-    }
-
-    public Builder setLaunchNotice(LaunchNotice launchNotice) {
-      getInstance().launchNotice = launchNotice;
-      return thisCastToDerived();
-    }
-
-    public Builder setIdnTableName(String idnTableName) {
-      getInstance().idnTableName = idnTableName;
-      return thisCastToDerived();
-    }
-
-    public Builder setSubordinateHosts(ImmutableSet<String> subordinateHosts) {
-      getInstance().subordinateHosts = subordinateHosts;
-      return thisCastToDerived();
-    }
-
-    public Builder addSubordinateHost(String hostToAdd) {
-      return setSubordinateHosts(ImmutableSet.copyOf(
-          union(getInstance().getSubordinateHosts(), hostToAdd)));
-    }
-
-    public Builder removeSubordinateHost(String hostToRemove) {
-      return setSubordinateHosts(ImmutableSet.copyOf(
-          CollectionUtils.difference(getInstance().getSubordinateHosts(), hostToRemove)));
-    }
-
-    public Builder setRegistrationExpirationTime(DateTime registrationExpirationTime) {
-      getInstance().registrationExpirationTime = registrationExpirationTime;
-      return this;
-    }
-
-    public Builder setDeletePollMessage(VKey<PollMessage.OneTime> deletePollMessage) {
-      getInstance().deletePollMessage = deletePollMessage;
-      return this;
-    }
-
-    public Builder setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
-      getInstance().autorenewBillingEvent = autorenewBillingEvent;
-      return this;
-    }
-
-    public Builder setAutorenewPollMessage(VKey<PollMessage.Autorenew> autorenewPollMessage) {
-      getInstance().autorenewPollMessage = autorenewPollMessage;
-      return this;
-    }
-
-    public Builder setSmdId(String smdId) {
-      getInstance().smdId = smdId;
-      return this;
-    }
-
-    public Builder setGracePeriods(ImmutableSet<GracePeriod> gracePeriods) {
-      getInstance().gracePeriods = gracePeriods;
-      return this;
-    }
-
-    public Builder addGracePeriod(GracePeriod gracePeriod) {
-      getInstance().gracePeriods = union(getInstance().getGracePeriods(), gracePeriod);
-      return this;
-    }
-
-    public Builder removeGracePeriod(GracePeriod gracePeriod) {
-      getInstance().gracePeriods = CollectionUtils
-          .difference(getInstance().getGracePeriods(), gracePeriod);
-      return this;
-    }
-
-    @Override
-    public Builder setTransferData(DomainTransferData transferData) {
-      getInstance().transferData = transferData;
-      return thisCastToDerived();
-    }
-
-    @Override
-    public Builder setLastTransferTime(DateTime lastTransferTime) {
-      getInstance().lastTransferTime = lastTransferTime;
-      return thisCastToDerived();
-    }
   }
 }

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

@@ -0,0 +1,853 @@
+// Copyright 2017 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.model.domain;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Strings.emptyToNull;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.ImmutableSortedSet.toImmutableSortedSet;
+import static com.google.common.collect.Sets.difference;
+import static com.google.common.collect.Sets.intersection;
+import static google.registry.model.EppResourceUtils.projectResourceOntoBuilderAtTime;
+import static google.registry.model.EppResourceUtils.setAutomaticTransferSuccessProperties;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.util.CollectionUtils.forceEmptyToNull;
+import static google.registry.util.CollectionUtils.nullToEmpty;
+import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
+import static google.registry.util.CollectionUtils.union;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
+import static google.registry.util.DateTimeUtils.earliestOf;
+import static google.registry.util.DateTimeUtils.isBeforeOrAt;
+import static google.registry.util.DateTimeUtils.leapSafeAddYears;
+import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
+import static google.registry.util.DomainNameUtils.getTldFromDomainName;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
+
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
+import com.google.common.collect.Ordering;
+import com.google.common.collect.Sets;
+import com.google.common.collect.Streams;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.Ignore;
+import com.googlecode.objectify.annotation.IgnoreSave;
+import com.googlecode.objectify.annotation.Index;
+import com.googlecode.objectify.annotation.OnLoad;
+import com.googlecode.objectify.condition.IfNull;
+import google.registry.flows.ResourceFlowUtils;
+import google.registry.model.EppResource;
+import google.registry.model.EppResource.ResourceWithTransferData;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.contact.ContactResource;
+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.StatusValue;
+import google.registry.model.host.HostResource;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registry.Registry;
+import google.registry.model.transfer.DomainTransferData;
+import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
+import google.registry.util.CollectionUtils;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Predicate;
+import javax.annotation.Nullable;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.AttributeOverride;
+import javax.persistence.AttributeOverrides;
+import javax.persistence.Column;
+import javax.persistence.Embeddable;
+import javax.persistence.Embedded;
+import javax.persistence.MappedSuperclass;
+import javax.persistence.PostLoad;
+import javax.persistence.Transient;
+import org.joda.time.DateTime;
+import org.joda.time.Interval;
+
+/**
+ * A persistable domain resource including mutable and non-mutable fields.
+ *
+ * <p>This class deliberately does not include an {@link javax.persistence.Id} so that any
+ * foreign-keyed fields can refer to the proper parent entity's ID, whether we're storing this in
+ * the DB itself or as part of another entity.
+ *
+ * <p>For historical reasons, the name of this class is "DomainContent". Ideally it would be
+ * "DomainBase" for parallelism with the other {@link EppResource} entity classes, but because that
+ * name is already taken by {@link DomainBase} (also for historical reasons), we can't use it. Once
+ * we are no longer on Datastore, we can rename the classes.
+ *
+ * @see <a href="https://tools.ietf.org/html/rfc5731">RFC 5731</a>
+ */
+@MappedSuperclass
+@Embeddable
+@Access(AccessType.FIELD)
+public class DomainContent extends EppResource
+    implements ResourceWithTransferData<DomainTransferData> {
+
+  /** The max number of years that a domain can be registered for, as set by ICANN policy. */
+  public static final int MAX_REGISTRATION_YEARS = 10;
+
+  /** Status values which prohibit DNS information from being published. */
+  private static final ImmutableSet<StatusValue> DNS_PUBLISHING_PROHIBITED_STATUSES =
+      ImmutableSet.of(
+          StatusValue.CLIENT_HOLD,
+          StatusValue.INACTIVE,
+          StatusValue.PENDING_DELETE,
+          StatusValue.SERVER_HOLD);
+
+  /**
+   * Fully qualified domain name (puny-coded), which serves as the foreign key for this domain.
+   *
+   * <p>This is only unique in the sense that for any given lifetime specified as the time range
+   * from (creationTime, deletionTime) there can only be one domain in Datastore with this name.
+   * However, there can be many domains with the same name and non-overlapping lifetimes.
+   *
+   * @invariant fullyQualifiedDomainName == fullyQualifiedDomainName.toLowerCase(Locale.ENGLISH)
+   */
+  // TODO(b/158858642): Rename this to domainName when we are off Datastore
+  @Column(name = "domainName")
+  @Index
+  String fullyQualifiedDomainName;
+
+  /** The top level domain this is under, dernormalized from {@link #fullyQualifiedDomainName}. */
+  @Index String tld;
+
+  /** References to hosts that are the nameservers for the domain. */
+  @Index @Transient Set<VKey<HostResource>> nsHosts;
+
+  /**
+   * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.
+   *
+   * <p>These are stored in one field so that we can query across all contacts at once.
+   */
+  @Transient Set<DesignatedContact> allContacts;
+
+  /**
+   * Contacts as they are stored in cloud SQL.
+   *
+   * <p>This information is duplicated in allContacts, and must be kept in sync with it.
+   */
+  @Ignore VKey<ContactResource> adminContact;
+
+  @Ignore VKey<ContactResource> billingContact;
+  @Ignore VKey<ContactResource> techContact;
+  @Ignore VKey<ContactResource> registrantContact;
+
+  /** Authorization info (aka transfer secret) of the domain. */
+  @Embedded
+  @AttributeOverrides({
+    @AttributeOverride(name = "pw.value", column = @Column(name = "auth_info_value")),
+    @AttributeOverride(name = "pw.repoId", column = @Column(name = "auth_info_repo_id")),
+  })
+  DomainAuthInfo authInfo;
+
+  /**
+   * Data used to construct DS records for this domain.
+   *
+   * <p>This is {@literal @}XmlTransient because it needs to be returned under the "extension" tag
+   * of an info response rather than inside the "infData" tag.
+   */
+  @Transient Set<DelegationSignerData> dsData;
+
+  /**
+   * The claims notice supplied when this application or domain was created, if there was one. It's
+   * {@literal @}XmlTransient because it's not returned in an info response.
+   */
+  @IgnoreSave(IfNull.class)
+  @Embedded
+  @AttributeOverrides({
+    @AttributeOverride(name = "noticeId.tcnId", column = @Column(name = "launch_notice_tcn_id")),
+    @AttributeOverride(
+        name = "noticeId.validatorId",
+        column = @Column(name = "launch_notice_validator_id")),
+    @AttributeOverride(
+        name = "expirationTime",
+        column = @Column(name = "launch_notice_expiration_time")),
+    @AttributeOverride(
+        name = "acceptedTime",
+        column = @Column(name = "launch_notice_accepted_time")),
+  })
+  LaunchNotice launchNotice;
+
+  /**
+   * Name of first IDN table associated with TLD that matched the characters in this domain label.
+   *
+   * @see google.registry.tldconfig.idn.IdnLabelValidator#findValidIdnTableForTld
+   */
+  @IgnoreSave(IfNull.class)
+  String idnTableName;
+
+  /** Fully qualified host names of this domain's active subordinate hosts. */
+  Set<String> subordinateHosts;
+
+  /** When this domain's registration will expire. */
+  DateTime registrationExpirationTime;
+
+  /**
+   * The poll message associated with this domain being deleted.
+   *
+   * <p>This field should be null if the domain is not in pending delete. If it is, the field should
+   * refer to a {@link PollMessage} timed to when the domain is fully deleted. If the domain is
+   * restored, the message should be deleted.
+   */
+  @Column(name = "deletion_poll_message_id")
+  VKey<PollMessage.OneTime> deletePollMessage;
+
+  /**
+   * The recurring billing event associated with this domain's autorenewals.
+   *
+   * <p>The recurrence should be open ended unless the domain is in pending delete or fully deleted,
+   * in which case it should be closed at the time the delete was requested. Whenever the domain's
+   * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
+   * should be created, and this field should be updated to point to the new one.
+   */
+  @Column(name = "billing_recurrence_id")
+  VKey<BillingEvent.Recurring> autorenewBillingEvent;
+
+  /**
+   * The recurring poll message associated with this domain's autorenewals.
+   *
+   * <p>The recurrence should be open ended unless the domain is in pending delete or fully deleted,
+   * in which case it should be closed at the time the delete was requested. Whenever the domain's
+   * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
+   * should be created, and this field should be updated to point to the new one.
+   */
+  @Column(name = "autorenew_poll_message_id")
+  VKey<PollMessage.Autorenew> autorenewPollMessage;
+
+  /** The unexpired grace periods for this domain (some of which may not be active yet). */
+  @Transient Set<GracePeriod> gracePeriods;
+
+  /**
+   * The id of the signed mark that was used to create this domain in sunrise.
+   *
+   * <p>Will only be populated for domains created in sunrise.
+   */
+  @IgnoreSave(IfNull.class)
+  String smdId;
+
+  /** Data about any pending or past transfers on this domain. */
+  DomainTransferData transferData;
+
+  /**
+   * The time that this resource was last transferred.
+   *
+   * <p>Can be null if the resource has never been transferred.
+   */
+  DateTime lastTransferTime;
+
+  @OnLoad
+  void load() {
+    // Reconstitute all of the contacts so that they have VKeys.
+    allContacts =
+        allContacts.stream().map(contact -> contact.reconstitute()).collect(toImmutableSet());
+    setContactFields(allContacts, true);
+
+    // We have to return the cloned object here because the original object's
+    // hashcode is not correct due to the change to its domainRepoId. The cloned
+    // object will have a null hashcode so that it can get a recalculated hashcode
+    // when its hashCode() is invoked.
+    // TODO(b/162739503): Remove this after fully migrating to Cloud SQL.
+    if (gracePeriods != null) {
+      gracePeriods =
+          gracePeriods.stream()
+              .map(gracePeriod -> gracePeriod.cloneWithDomainRepoId(getRepoId()))
+              .collect(toImmutableSet());
+    }
+  }
+
+  @PostLoad
+  void postLoad() {
+    // Reconstitute the contact list.
+    ImmutableSet.Builder<DesignatedContact> contactsBuilder =
+        new ImmutableSet.Builder<DesignatedContact>();
+
+    if (registrantContact != null) {
+      contactsBuilder.add(
+          DesignatedContact.create(DesignatedContact.Type.REGISTRANT, registrantContact));
+    }
+    if (billingContact != null) {
+      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.BILLING, billingContact));
+    }
+    if (techContact != null) {
+      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.TECH, techContact));
+    }
+    if (adminContact != null) {
+      contactsBuilder.add(DesignatedContact.create(DesignatedContact.Type.ADMIN, adminContact));
+    }
+
+    allContacts = contactsBuilder.build();
+  }
+
+  public ImmutableSet<String> getSubordinateHosts() {
+    return nullToEmptyImmutableCopy(subordinateHosts);
+  }
+
+  public DateTime getRegistrationExpirationTime() {
+    return registrationExpirationTime;
+  }
+
+  public VKey<PollMessage.OneTime> getDeletePollMessage() {
+    return deletePollMessage;
+  }
+
+  public VKey<BillingEvent.Recurring> getAutorenewBillingEvent() {
+    return autorenewBillingEvent;
+  }
+
+  public VKey<PollMessage.Autorenew> getAutorenewPollMessage() {
+    return autorenewPollMessage;
+  }
+
+  public ImmutableSet<GracePeriod> getGracePeriods() {
+    return nullToEmptyImmutableCopy(gracePeriods);
+  }
+
+  public String getSmdId() {
+    return smdId;
+  }
+
+  @Override
+  public DomainTransferData getTransferData() {
+    return Optional.ofNullable(transferData).orElse(DomainTransferData.EMPTY);
+  }
+
+  @Override
+  public DateTime getLastTransferTime() {
+    return lastTransferTime;
+  }
+
+  @Override
+  public String getForeignKey() {
+    return fullyQualifiedDomainName;
+  }
+
+  public String getDomainName() {
+    return fullyQualifiedDomainName;
+  }
+
+  public ImmutableSet<DelegationSignerData> getDsData() {
+    return nullToEmptyImmutableCopy(dsData);
+  }
+
+  public LaunchNotice getLaunchNotice() {
+    return launchNotice;
+  }
+
+  public String getIdnTableName() {
+    return idnTableName;
+  }
+
+  public ImmutableSet<VKey<HostResource>> getNameservers() {
+    return nullToEmptyImmutableCopy(nsHosts);
+  }
+
+  // Hibernate needs this in order to populate nsHosts but no one else should ever use it
+  @SuppressWarnings("UnusedMethod")
+  private void setNsHosts(Set<VKey<HostResource>> nsHosts) {
+    this.nsHosts = nsHosts;
+  }
+
+  // Hibernate needs this in order to populate gracePeriods but no one else should ever use it
+  @SuppressWarnings("UnusedMethod")
+  private void setInternalGracePeriods(Set<GracePeriod> gracePeriods) {
+    this.gracePeriods = gracePeriods;
+  }
+
+  public final String getCurrentSponsorClientId() {
+    return getPersistedCurrentSponsorClientId();
+  }
+
+  /** Returns true if DNS information should be published for the given domain. */
+  public boolean shouldPublishToDns() {
+    return intersection(getStatusValues(), DNS_PUBLISHING_PROHIBITED_STATUSES).isEmpty();
+  }
+
+  /**
+   * Returns the Registry Grace Period Statuses for this domain.
+   *
+   * <p>This collects all statuses from the domain's {@link GracePeriod} entries and also adds the
+   * PENDING_DELETE status if needed.
+   */
+  public ImmutableSet<GracePeriodStatus> getGracePeriodStatuses() {
+    Set<GracePeriodStatus> gracePeriodStatuses = new HashSet<>();
+    for (GracePeriod gracePeriod : getGracePeriods()) {
+      gracePeriodStatuses.add(gracePeriod.getType());
+    }
+    if (getStatusValues().contains(StatusValue.PENDING_DELETE)
+        && !gracePeriodStatuses.contains(GracePeriodStatus.REDEMPTION)) {
+      gracePeriodStatuses.add(GracePeriodStatus.PENDING_DELETE);
+    }
+    return ImmutableSet.copyOf(gracePeriodStatuses);
+  }
+
+  /** Returns the subset of grace periods having the specified type. */
+  public ImmutableSet<GracePeriod> getGracePeriodsOfType(GracePeriodStatus gracePeriodType) {
+    ImmutableSet.Builder<GracePeriod> builder = new ImmutableSet.Builder<>();
+    for (GracePeriod gracePeriod : getGracePeriods()) {
+      if (gracePeriod.getType() == gracePeriodType) {
+        builder.add(gracePeriod);
+      }
+    }
+    return builder.build();
+  }
+
+  @Override
+  public DomainContent cloneProjectedAtTime(final DateTime now) {
+    return cloneDomainProjectedAtTime(this, now);
+  }
+
+  /**
+   * The logic in this method, which handles implicit server approval of transfers, very closely
+   * parallels the logic in {@code DomainTransferApproveFlow} which handles explicit client
+   * approvals.
+   */
+  protected static <T extends DomainContent> T cloneDomainProjectedAtTime(T domain, DateTime now) {
+    DomainTransferData transferData = domain.getTransferData();
+    DateTime transferExpirationTime = transferData.getPendingTransferExpirationTime();
+
+    // If there's a pending transfer that has expired, handle it.
+    if (TransferStatus.PENDING.equals(transferData.getTransferStatus())
+        && isBeforeOrAt(transferExpirationTime, now)) {
+      // Project until just before the transfer time. This will handle the case of an autorenew
+      // before the transfer was even requested or during the request period.
+      // If the transfer time is precisely the moment that the domain expires, there will not be an
+      // autorenew billing event (since we end the recurrence at transfer time and recurrences are
+      // exclusive of their ending), and we can just proceed with the transfer.
+      T domainAtTransferTime =
+          cloneDomainProjectedAtTime(domain, transferExpirationTime.minusMillis(1));
+
+      DateTime expirationDate = transferData.getTransferredRegistrationExpirationTime();
+      if (expirationDate == null) {
+        // Extend the registration by the correct number of years from the expiration time
+        // that was current on the domain right before the transfer, capped at 10 years from
+        // the moment of the transfer.
+        expirationDate =
+            ResourceFlowUtils.computeExDateForApprovalTime(
+                domainAtTransferTime, transferExpirationTime, transferData.getTransferPeriod());
+      }
+      // If we are within an autorenew grace period, the transfer will subsume the autorenew. There
+      // will already be a cancellation written in advance by the transfer request flow, so we don't
+      // need to worry about billing, but we do need to cancel out the expiration time increase.
+      // The transfer period saved in the transfer data will be one year, unless the superuser
+      // extension set the transfer period to zero.
+      // Set the expiration, autorenew events, and grace period for the transfer. (Transfer ends
+      // all other graces).
+      Builder builder =
+          domainAtTransferTime
+              .asBuilder()
+              .setRegistrationExpirationTime(expirationDate)
+              // Set the speculatively-written new autorenew events as the domain's autorenew
+              // events.
+              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
+              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
+      if (transferData.getTransferPeriod().getValue() == 1) {
+        // Set the grace period using a key to the prescheduled transfer billing event.  Not using
+        // GracePeriod.forBillingEvent() here in order to avoid the actual Datastore fetch.
+        builder.setGracePeriods(
+            ImmutableSet.of(
+                GracePeriod.create(
+                    GracePeriodStatus.TRANSFER,
+                    domain.getRepoId(),
+                    transferExpirationTime.plus(
+                        Registry.get(domain.getTld()).getTransferGracePeriodLength()),
+                    transferData.getGainingClientId(),
+                    transferData.getServerApproveBillingEvent())));
+      } else {
+        // There won't be a billing event, so we don't need a grace period
+        builder.setGracePeriods(ImmutableSet.of());
+      }
+      // Set all remaining transfer properties.
+      setAutomaticTransferSuccessProperties(builder, transferData);
+      builder
+          .setLastEppUpdateTime(transferExpirationTime)
+          .setLastEppUpdateClientId(transferData.getGainingClientId());
+      // Finish projecting to now.
+      return (T) builder.build().cloneProjectedAtTime(now);
+    }
+
+    Optional<DateTime> newLastEppUpdateTime = Optional.empty();
+
+    // There is no transfer. Do any necessary autorenews for active domains.
+
+    Builder builder = domain.asBuilder();
+    if (isBeforeOrAt(domain.getRegistrationExpirationTime(), now)
+        && END_OF_TIME.equals(domain.getDeletionTime())) {
+      // Autorenew by the number of years between the old expiration time and now.
+      DateTime lastAutorenewTime =
+          leapSafeAddYears(
+              domain.getRegistrationExpirationTime(),
+              new Interval(domain.getRegistrationExpirationTime(), now).toPeriod().getYears());
+      DateTime newExpirationTime = lastAutorenewTime.plusYears(1);
+      builder
+          .setRegistrationExpirationTime(newExpirationTime)
+          .addGracePeriod(
+              GracePeriod.createForRecurring(
+                  GracePeriodStatus.AUTO_RENEW,
+                  domain.getRepoId(),
+                  lastAutorenewTime.plus(
+                      Registry.get(domain.getTld()).getAutoRenewGracePeriodLength()),
+                  domain.getCurrentSponsorClientId(),
+                  domain.getAutorenewBillingEvent()));
+      newLastEppUpdateTime = Optional.of(lastAutorenewTime);
+    }
+
+    // Remove any grace periods that have expired.
+    T almostBuilt = (T) builder.build();
+    builder = almostBuilt.asBuilder();
+    for (GracePeriod gracePeriod : almostBuilt.getGracePeriods()) {
+      if (isBeforeOrAt(gracePeriod.getExpirationTime(), now)) {
+        builder.removeGracePeriod(gracePeriod);
+        if (!newLastEppUpdateTime.isPresent()
+            || isBeforeOrAt(newLastEppUpdateTime.get(), gracePeriod.getExpirationTime())) {
+          newLastEppUpdateTime = Optional.of(gracePeriod.getExpirationTime());
+        }
+      }
+    }
+
+    // It is possible that the lastEppUpdateClientId is different from current sponsor client
+    // id, so we have to do the comparison instead of having one variable just storing the most
+    // recent time.
+    if (newLastEppUpdateTime.isPresent()) {
+      if (domain.getLastEppUpdateTime() == null
+          || newLastEppUpdateTime.get().isAfter(domain.getLastEppUpdateTime())) {
+        builder
+            .setLastEppUpdateTime(newLastEppUpdateTime.get())
+            .setLastEppUpdateClientId(domain.getCurrentSponsorClientId());
+      }
+    }
+
+    // Handle common properties like setting or unsetting linked status. This also handles the
+    // general case of pending transfers for other resource types, but since we've always handled
+    // a pending transfer by this point that's a no-op for domains.
+    projectResourceOntoBuilderAtTime(almostBuilt, builder, now);
+    return (T) builder.build();
+  }
+
+  /** Return what the expiration time would be if the given number of years were added to it. */
+  public static DateTime extendRegistrationWithCap(
+      DateTime now, DateTime currentExpirationTime, @Nullable Integer extendedRegistrationYears) {
+    // We must cap registration at the max years (aka 10), even if that truncates the last year.
+    return earliestOf(
+        leapSafeAddYears(
+            currentExpirationTime, Optional.ofNullable(extendedRegistrationYears).orElse(0)),
+        leapSafeAddYears(now, MAX_REGISTRATION_YEARS));
+  }
+
+  /** Loads and returns the fully qualified host names of all linked nameservers. */
+  public ImmutableSortedSet<String> loadNameserverHostNames() {
+    return ofy()
+        .load()
+        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet()))
+        .values()
+        .stream()
+        .map(HostResource::getHostName)
+        .collect(toImmutableSortedSet(Ordering.natural()));
+  }
+
+  /** A key to the registrant who registered this domain. */
+  public VKey<ContactResource> getRegistrant() {
+    return registrantContact;
+  }
+
+  public VKey<ContactResource> getAdminContact() {
+    return adminContact;
+  }
+
+  public VKey<ContactResource> getBillingContact() {
+    return billingContact;
+  }
+
+  public VKey<ContactResource> getTechContact() {
+    return techContact;
+  }
+
+  /** Associated contacts for the domain (other than registrant). */
+  public ImmutableSet<DesignatedContact> getContacts() {
+    return nullToEmpty(allContacts).stream()
+        .filter(IS_REGISTRANT.negate())
+        .collect(toImmutableSet());
+  }
+
+  public DomainAuthInfo getAuthInfo() {
+    return authInfo;
+  }
+
+  /** Returns all referenced contacts from this domain or application. */
+  public ImmutableSet<VKey<ContactResource>> getReferencedContacts() {
+    return nullToEmptyImmutableCopy(allContacts).stream()
+        .map(DesignatedContact::getContactKey)
+        .filter(Objects::nonNull)
+        .collect(toImmutableSet());
+  }
+
+  public String getTld() {
+    return tld;
+  }
+
+  /**
+   * Sets the individual contact fields from {@code contacts}.
+   *
+   * <p>The registrant field is only set if {@code includeRegistrant} is true, as this field needs
+   * to be set in some circumstances but not in others.
+   */
+  protected void setContactFields(Set<DesignatedContact> contacts, boolean includeRegistrant) {
+    // Set the individual contact fields.
+    for (DesignatedContact contact : contacts) {
+      switch (contact.getType()) {
+        case BILLING:
+          billingContact = contact.getContactKey();
+          break;
+        case TECH:
+          techContact = contact.getContactKey();
+          break;
+        case ADMIN:
+          adminContact = contact.getContactKey();
+          break;
+        case REGISTRANT:
+          if (includeRegistrant) {
+            registrantContact = contact.getContactKey();
+          }
+          break;
+        default:
+          throw new IllegalArgumentException("Unknown contact resource type: " + contact.getType());
+      }
+    }
+  }
+
+  @Override
+  public VKey<DomainBase> createVKey() {
+    return VKey.create(DomainBase.class, getRepoId(), Key.create(this));
+  }
+
+  public static VKey<DomainBase> createVKey(Key key) {
+    return VKey.create(DomainBase.class, key.getName(), key);
+  }
+
+  /** Predicate to determine if a given {@link DesignatedContact} is the registrant. */
+  protected static final Predicate<DesignatedContact> IS_REGISTRANT =
+      (DesignatedContact contact) -> DesignatedContact.Type.REGISTRANT.equals(contact.type);
+
+  /** An override of {@link EppResource#asBuilder} with tighter typing. */
+  @Override
+  public Builder asBuilder() {
+    return new Builder<>(clone(this));
+  }
+
+  /** A builder for constructing {@link DomainBase}, since it is immutable. */
+  public static class Builder<T extends DomainContent, B extends Builder<T, B>>
+      extends EppResource.Builder<T, B> implements BuilderWithTransferData<DomainTransferData, B> {
+
+    public Builder() {}
+
+    Builder(T instance) {
+      super(instance);
+    }
+
+    @Override
+    public T build() {
+      T instance = getInstance();
+      // If TransferData is totally empty, set it to null.
+      if (DomainTransferData.EMPTY.equals(getInstance().transferData)) {
+        setTransferData(null);
+      }
+      // A DomainBase has status INACTIVE if there are no nameservers.
+      if (getInstance().getNameservers().isEmpty()) {
+        addStatusValue(StatusValue.INACTIVE);
+      } else { // There are nameservers, so make sure INACTIVE isn't there.
+        removeStatusValue(StatusValue.INACTIVE);
+      }
+
+      checkArgumentNotNull(emptyToNull(instance.fullyQualifiedDomainName), "Missing domainName");
+      if (instance.getRegistrant() == null
+          && instance.allContacts.stream().anyMatch(IS_REGISTRANT)) {
+        throw new IllegalArgumentException("registrant is null but is in allContacts");
+      }
+      checkArgumentNotNull(instance.getRegistrant(), "Missing registrant");
+      instance.tld = getTldFromDomainName(instance.fullyQualifiedDomainName);
+      return super.build();
+    }
+
+    public B setDomainName(String domainName) {
+      checkArgument(
+          domainName.equals(canonicalizeDomainName(domainName)),
+          "Domain name must be in puny-coded, lower-case form");
+      getInstance().fullyQualifiedDomainName = domainName;
+      return thisCastToDerived();
+    }
+
+    public B setDsData(ImmutableSet<DelegationSignerData> dsData) {
+      getInstance().dsData = dsData;
+      return thisCastToDerived();
+    }
+
+    public B setRegistrant(VKey<ContactResource> registrant) {
+      // Replace the registrant contact inside allContacts.
+      getInstance().allContacts =
+          union(
+              getInstance().getContacts(),
+              DesignatedContact.create(
+                  DesignatedContact.Type.REGISTRANT, checkArgumentNotNull(registrant)));
+
+      // Set the registrant field specifically.
+      getInstance().registrantContact = registrant;
+      return thisCastToDerived();
+    }
+
+    public B setAuthInfo(DomainAuthInfo authInfo) {
+      getInstance().authInfo = authInfo;
+      return thisCastToDerived();
+    }
+
+    public B setNameservers(VKey<HostResource> nameserver) {
+      getInstance().nsHosts = ImmutableSet.of(nameserver);
+      return thisCastToDerived();
+    }
+
+    public B setNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
+      getInstance().nsHosts = forceEmptyToNull(nameservers);
+      return thisCastToDerived();
+    }
+
+    public B addNameserver(VKey<HostResource> nameserver) {
+      return addNameservers(ImmutableSet.of(nameserver));
+    }
+
+    public B addNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
+      return setNameservers(
+          ImmutableSet.copyOf(Sets.union(getInstance().getNameservers(), nameservers)));
+    }
+
+    public B removeNameserver(VKey<HostResource> nameserver) {
+      return removeNameservers(ImmutableSet.of(nameserver));
+    }
+
+    public B removeNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
+      return setNameservers(
+          ImmutableSet.copyOf(difference(getInstance().getNameservers(), nameservers)));
+    }
+
+    public B setContacts(DesignatedContact contact) {
+      return setContacts(ImmutableSet.of(contact));
+    }
+
+    public B setContacts(ImmutableSet<DesignatedContact> contacts) {
+      checkArgument(contacts.stream().noneMatch(IS_REGISTRANT), "Registrant cannot be a contact");
+
+      // Replace the non-registrant contacts inside allContacts.
+      getInstance().allContacts =
+          Streams.concat(
+                  nullToEmpty(getInstance().allContacts).stream().filter(IS_REGISTRANT),
+                  contacts.stream())
+              .collect(toImmutableSet());
+
+      // Set the individual fields.
+      getInstance().setContactFields(contacts, false);
+      return thisCastToDerived();
+    }
+
+    public B addContacts(ImmutableSet<DesignatedContact> contacts) {
+      return setContacts(ImmutableSet.copyOf(Sets.union(getInstance().getContacts(), contacts)));
+    }
+
+    public B removeContacts(ImmutableSet<DesignatedContact> contacts) {
+      return setContacts(ImmutableSet.copyOf(difference(getInstance().getContacts(), contacts)));
+    }
+
+    public B setLaunchNotice(LaunchNotice launchNotice) {
+      getInstance().launchNotice = launchNotice;
+      return thisCastToDerived();
+    }
+
+    public B setIdnTableName(String idnTableName) {
+      getInstance().idnTableName = idnTableName;
+      return thisCastToDerived();
+    }
+
+    public B setSubordinateHosts(ImmutableSet<String> subordinateHosts) {
+      getInstance().subordinateHosts = subordinateHosts;
+      return thisCastToDerived();
+    }
+
+    public B addSubordinateHost(String hostToAdd) {
+      return setSubordinateHosts(
+          ImmutableSet.copyOf(union(getInstance().getSubordinateHosts(), hostToAdd)));
+    }
+
+    public B removeSubordinateHost(String hostToRemove) {
+      return setSubordinateHosts(
+          ImmutableSet.copyOf(
+              CollectionUtils.difference(getInstance().getSubordinateHosts(), hostToRemove)));
+    }
+
+    public B setRegistrationExpirationTime(DateTime registrationExpirationTime) {
+      getInstance().registrationExpirationTime = registrationExpirationTime;
+      return thisCastToDerived();
+    }
+
+    public B setDeletePollMessage(VKey<PollMessage.OneTime> deletePollMessage) {
+      getInstance().deletePollMessage = deletePollMessage;
+      return thisCastToDerived();
+    }
+
+    public B setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
+      getInstance().autorenewBillingEvent = autorenewBillingEvent;
+      return thisCastToDerived();
+    }
+
+    public B setAutorenewPollMessage(VKey<PollMessage.Autorenew> autorenewPollMessage) {
+      getInstance().autorenewPollMessage = autorenewPollMessage;
+      return thisCastToDerived();
+    }
+
+    public B setSmdId(String smdId) {
+      getInstance().smdId = smdId;
+      return thisCastToDerived();
+    }
+
+    public B setGracePeriods(ImmutableSet<GracePeriod> gracePeriods) {
+      getInstance().gracePeriods = gracePeriods;
+      return thisCastToDerived();
+    }
+
+    public B addGracePeriod(GracePeriod gracePeriod) {
+      getInstance().gracePeriods = union(getInstance().getGracePeriods(), gracePeriod);
+      return thisCastToDerived();
+    }
+
+    public B removeGracePeriod(GracePeriod gracePeriod) {
+      getInstance().gracePeriods =
+          CollectionUtils.difference(getInstance().getGracePeriods(), gracePeriod);
+      return thisCastToDerived();
+    }
+
+    @Override
+    public B setTransferData(DomainTransferData transferData) {
+      getInstance().transferData = transferData;
+      return thisCastToDerived();
+    }
+
+    @Override
+    public B setLastTransferTime(DateTime lastTransferTime) {
+      getInstance().lastTransferTime = lastTransferTime;
+      return thisCastToDerived();
+    }
+  }
+}

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

@@ -0,0 +1,111 @@
+// 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.model.domain;
+
+import com.googlecode.objectify.Key;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.host.HostResource;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import java.util.Set;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.Column;
+import javax.persistence.ElementCollection;
+import javax.persistence.Entity;
+import javax.persistence.JoinTable;
+
+/**
+ * A persisted history entry representing an EPP modification to a domain.
+ *
+ * <p>In addition to the general history fields (e.g. action time, registrar ID) we also persist a
+ * copy of the domain entity at this point in time. We persist a raw {@link DomainContent} so that
+ * the foreign-keyed fields in that class can refer to this object.
+ */
+@Entity
+@javax.persistence.Table(
+    indexes = {
+      @javax.persistence.Index(columnList = "creationTime"),
+      @javax.persistence.Index(columnList = "historyRegistrarId"),
+      @javax.persistence.Index(columnList = "historyType"),
+      @javax.persistence.Index(columnList = "historyModificationTime")
+    })
+public class DomainHistory extends HistoryEntry {
+  // Store DomainContent instead of DomainBase so we don't pick up its @Id
+  DomainContent domainContent;
+
+  @Column(nullable = false)
+  VKey<DomainBase> domainRepoId;
+
+  @ElementCollection
+  @JoinTable(name = "DomainHistoryHost")
+  @Access(AccessType.PROPERTY)
+  @Column(name = "host_repo_id")
+  public Set<VKey<HostResource>> getNsHosts() {
+    return domainContent.nsHosts;
+  }
+
+  /** The state of the {@link DomainContent} object at this point in time. */
+  public DomainContent getDomainContent() {
+    return domainContent;
+  }
+
+  /** The key to the {@link ContactResource} this is based off of. */
+  public VKey<DomainBase> getDomainRepoId() {
+    return domainRepoId;
+  }
+
+  // Hibernate needs this in order to populate nsHosts but no one else should ever use it
+  @SuppressWarnings("UnusedMethod")
+  private void setNsHosts(Set<VKey<HostResource>> nsHosts) {
+    if (domainContent != null) {
+      domainContent.nsHosts = nsHosts;
+    }
+  }
+
+  @Override
+  public Builder asBuilder() {
+    return new Builder(clone(this));
+  }
+
+  public static class Builder extends HistoryEntry.Builder<DomainHistory, DomainHistory.Builder> {
+
+    public Builder() {}
+
+    public Builder(DomainHistory instance) {
+      super(instance);
+    }
+
+    public Builder setDomainContent(DomainContent domainContent) {
+      getInstance().domainContent = domainContent;
+      return this;
+    }
+
+    public Builder setDomainRepoId(VKey<DomainBase> domainRepoId) {
+      getInstance().domainRepoId = domainRepoId;
+      domainRepoId.maybeGetOfyKey().ifPresent(parent -> getInstance().parent = parent);
+      return this;
+    }
+
+    // We can remove this once all HistoryEntries are converted to History objects
+    @Override
+    public Builder setParent(Key<? extends EppResource> parent) {
+      super.setParent(parent);
+      getInstance().domainRepoId = VKey.create(DomainBase.class, parent.getName(), parent);
+      return this;
+    }
+  }
+}

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

@@ -18,16 +18,14 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.googlecode.objectify.annotation.Embed;
-import com.googlecode.objectify.annotation.Ignore;
-import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
-import javax.persistence.Column;
-import javax.persistence.GeneratedValue;
-import javax.persistence.GenerationType;
+import javax.persistence.Entity;
+import javax.persistence.Index;
+import javax.persistence.Table;
 import org.joda.time.DateTime;
 
 /**
@@ -37,75 +35,13 @@ import org.joda.time.DateTime;
  * the resource is loaded from Datastore.
  */
 @Embed
-@javax.persistence.Entity
-public class GracePeriod extends ImmutableObject {
-
-  @javax.persistence.Id
-  @GeneratedValue(strategy = GenerationType.IDENTITY)
-  @Ignore
-  /** Unique id required for hibernate representation. */
-  long id;
-
-  /** The type of grace period. */
-  GracePeriodStatus type;
-
-  /** When the grace period ends. */
-  DateTime expirationTime;
-
-  /** The registrar to bill. */
-  @Column(name = "registrarId")
-  String clientId;
-
-  /**
-   * The one-time billing event corresponding to the action that triggered this grace period, or
-   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
-   * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
-   */
-  // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  VKey<BillingEvent.OneTime> billingEventOneTime = null;
-
-  /**
-   * The recurring billing event corresponding to the action that triggered this grace period, if
-   * applicable - i.e. if the action was an autorenew - or null in all other cases.
-   */
-  // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  VKey<BillingEvent.Recurring> billingEventRecurring = null;
-
-  public GracePeriodStatus getType() {
-    return type;
-  }
-
-  public DateTime getExpirationTime() {
-    return expirationTime;
-  }
-
-  public String getClientId() {
-    return clientId;
-  }
-
-  /** Returns true if this GracePeriod has an associated BillingEvent; i.e. if it's refundable. */
-  public boolean hasBillingEvent() {
-    return billingEventOneTime != null || billingEventRecurring != null;
-  }
-
-  /**
-   * Returns the one time billing event. The value will only be non-null if the type of this grace
-   * period is not AUTO_RENEW.
-   */
-  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
-    return billingEventOneTime;
-  }
-
-  /**
-   * Returns the recurring billing event. The value will only be non-null if the type of this grace
-   * period is AUTO_RENEW.
-   */
-  public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
-    return billingEventRecurring;
-  }
+@Entity
+@Table(indexes = @Index(columnList = "domainRepoId"))
+public class GracePeriod extends GracePeriodBase {
 
   private static GracePeriod createInternal(
       GracePeriodStatus type,
+      String domainRepoId,
       DateTime expirationTime,
       String clientId,
       @Nullable VKey<BillingEvent.OneTime> billingEventOneTime,
@@ -117,6 +53,7 @@ public class GracePeriod extends ImmutableObject {
         "Recurring billing events must be present on (and only on) autorenew grace periods");
     GracePeriod instance = new GracePeriod();
     instance.type = checkArgumentNotNull(type);
+    instance.domainRepoId = checkArgumentNotNull(domainRepoId);
     instance.expirationTime = checkArgumentNotNull(expirationTime);
     instance.clientId = checkArgumentNotNull(clientId);
     instance.billingEventOneTime = billingEventOneTime;
@@ -133,32 +70,51 @@ public class GracePeriod extends ImmutableObject {
    */
   public static GracePeriod create(
       GracePeriodStatus type,
+      String domainRepoId,
       DateTime expirationTime,
       String clientId,
       @Nullable VKey<BillingEvent.OneTime> billingEventOneTime) {
-    return createInternal(type, expirationTime, clientId, billingEventOneTime, null);
+    return createInternal(type, domainRepoId, expirationTime, clientId, billingEventOneTime, null);
   }
 
   /** Creates a GracePeriod for a Recurring billing event. */
   public static GracePeriod createForRecurring(
       GracePeriodStatus type,
+      String domainRepoId,
       DateTime expirationTime,
       String clientId,
       VKey<Recurring> billingEventRecurring) {
     checkArgumentNotNull(billingEventRecurring, "billingEventRecurring cannot be null");
-    return createInternal(type, expirationTime, clientId, null, billingEventRecurring);
+    return createInternal(
+        type, domainRepoId, expirationTime, clientId, null, billingEventRecurring);
   }
 
   /** Creates a GracePeriod with no billing event. */
   public static GracePeriod createWithoutBillingEvent(
-      GracePeriodStatus type, DateTime expirationTime, String clientId) {
-    return createInternal(type, expirationTime, clientId, null, null);
+      GracePeriodStatus type, String domainRepoId, DateTime expirationTime, String clientId) {
+    return createInternal(type, domainRepoId, expirationTime, clientId, null, null);
   }
 
   /** Constructs a GracePeriod of the given type from the provided one-time BillingEvent. */
   public static GracePeriod forBillingEvent(
-      GracePeriodStatus type, BillingEvent.OneTime billingEvent) {
+      GracePeriodStatus type, String domainRepoId, BillingEvent.OneTime billingEvent) {
     return create(
-        type, billingEvent.getBillingTime(), billingEvent.getClientId(), billingEvent.createVKey());
+        type,
+        domainRepoId,
+        billingEvent.getBillingTime(),
+        billingEvent.getClientId(),
+        billingEvent.createVKey());
   }
+
+  /**
+   * Returns a clone of this {@link GracePeriod} with {@link #domainRepoId} set to the given value.
+   *
+   * <p>TODO(b/162739503): Remove this function after fully migrating to Cloud SQL.
+   */
+  public GracePeriod cloneWithDomainRepoId(String domainRepoId) {
+    GracePeriod clone = clone(this);
+    clone.domainRepoId = checkArgumentNotNull(domainRepoId);
+    return clone;
+  }
+
 }

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

@@ -0,0 +1,114 @@
+// 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.model.domain;
+
+import com.googlecode.objectify.annotation.Embed;
+import com.googlecode.objectify.annotation.Ignore;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.OneTime;
+import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.persistence.VKey;
+import javax.persistence.Column;
+import javax.persistence.EnumType;
+import javax.persistence.Enumerated;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.MappedSuperclass;
+import org.joda.time.DateTime;
+
+/** Base class containing common fields and methods for {@link GracePeriod}. */
+@Embed
+@MappedSuperclass
+public class GracePeriodBase extends ImmutableObject {
+
+  /** Unique id required for hibernate representation. */
+  @javax.persistence.Id
+  @GeneratedValue(strategy = GenerationType.IDENTITY)
+  @Ignore
+  Long id;
+
+  /** Repository id for the domain which this grace period belongs to. */
+  @Ignore
+  @Column(nullable = false)
+  String domainRepoId;
+
+  /** The type of grace period. */
+  @Column(nullable = false)
+  @Enumerated(EnumType.STRING)
+  GracePeriodStatus type;
+
+  /** When the grace period ends. */
+  @Column(nullable = false)
+  DateTime expirationTime;
+
+  /** The registrar to bill. */
+  @Column(name = "registrarId", nullable = false)
+  String clientId;
+
+  /**
+   * The one-time billing event corresponding to the action that triggered this grace period, or
+   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
+   * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
+   */
+  // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
+  @Column(name = "billing_event_id")
+  VKey<OneTime> billingEventOneTime = null;
+
+  /**
+   * The recurring billing event corresponding to the action that triggered this grace period, if
+   * applicable - i.e. if the action was an autorenew - or null in all other cases.
+   */
+  // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
+  @Column(name = "billing_recurrence_id")
+  VKey<BillingEvent.Recurring> billingEventRecurring = null;
+
+  public GracePeriodStatus getType() {
+    return type;
+  }
+
+  public String getDomainRepoId() {
+    return domainRepoId;
+  }
+
+  public DateTime getExpirationTime() {
+    return expirationTime;
+  }
+
+  public String getClientId() {
+    return clientId;
+  }
+
+  /** Returns true if this GracePeriod has an associated BillingEvent; i.e. if it's refundable. */
+  public boolean hasBillingEvent() {
+    return billingEventOneTime != null || billingEventRecurring != null;
+  }
+
+  /**
+   * Returns the one time billing event. The value will only be non-null if the type of this grace
+   * period is not AUTO_RENEW.
+   */
+  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
+    return billingEventOneTime;
+  }
+
+  /**
+   * Returns the recurring billing event. The value will only be non-null if the type of this grace
+   * period is AUTO_RENEW.
+   */
+  public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
+    return billingEventRecurring;
+  }
+}

core/src/main/java/google/registry/model/domain/token/AllocationToken.java

@@ -29,12 +29,14 @@ import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Range;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Index;
 import com.googlecode.objectify.annotation.Mapify;
+import com.googlecode.objectify.annotation.OnLoad;
 import google.registry.flows.EppException;
 import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.BackupGroupRoot;
@@ -108,9 +110,22 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
    */
   double discountFraction;
 
+  /** Whether the discount fraction (if any) also applies to premium names. Defaults to false. */
+  boolean discountPremiums;
+
+  /** Up to how many years of initial creation receive the discount (if any). Defaults to 1. */
+  int discountYears = 1;
+
   /** The type of the token, either single-use or unlimited-use. */
-  // TODO(b/130301183): this should not be nullable, we can remove this once we're sure it isn't
-  @Nullable TokenType tokenType;
+  TokenType tokenType;
+
+  // TODO: Remove onLoad once all allocation tokens are migrated to have a discountYears of 1.
+  @OnLoad
+  void onLoad() {
+    if (discountYears == 0) {
+      discountYears = 1;
+    }
+  }
 
   /**
    * Promotional token validity periods.
@@ -146,8 +161,8 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return token;
   }
 
-  public Key<HistoryEntry> getRedemptionHistoryEntry() {
-    return redemptionHistoryEntry;
+  public Optional<Key<HistoryEntry>> getRedemptionHistoryEntry() {
+    return Optional.ofNullable(redemptionHistoryEntry);
   }
 
   public boolean isRedeemed() {
@@ -174,6 +189,16 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return discountFraction;
   }
 
+  public boolean shouldDiscountPremiums() {
+    return discountPremiums;
+  }
+
+  public int getDiscountYears() {
+    // Allocation tokens created prior to the addition of the discountYears field will have a value
+    // of 0 for it, but it should be the default value of 1 to retain the previous behavior.
+    return Math.max(1, discountYears);
+  }
+
   public TokenType getTokenType() {
     return tokenType;
   }
@@ -193,6 +218,7 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
 
   /** A builder for constructing {@link AllocationToken} objects, since they are immutable. */
   public static class Builder extends Buildable.Builder<AllocationToken> {
+
     public Builder() {}
 
     private Builder(AllocationToken instance) {
@@ -210,6 +236,12 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
           getInstance().redemptionHistoryEntry == null
               || TokenType.SINGLE_USE.equals(getInstance().tokenType),
           "Redemption history entry can only be specified for SINGLE_USE tokens");
+      checkArgument(
+          getInstance().discountFraction > 0 || !getInstance().discountPremiums,
+          "Discount premiums can only be specified along with a discount fraction");
+      checkArgument(
+          getInstance().discountFraction > 0 || getInstance().discountYears == 1,
+          "Discount years can only be specified along with a discount fraction");
       if (getInstance().domainName != null) {
         try {
           DomainFlowUtils.validateDomainName(getInstance().domainName);
@@ -258,10 +290,26 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     }
 
     public Builder setDiscountFraction(double discountFraction) {
+      checkArgument(
+          Range.closed(0.0d, 1.0d).contains(discountFraction),
+          "Discount fraction must be between 0 and 1 inclusive");
       getInstance().discountFraction = discountFraction;
       return this;
     }
 
+    public Builder setDiscountPremiums(boolean discountPremiums) {
+      getInstance().discountPremiums = discountPremiums;
+      return this;
+    }
+
+    public Builder setDiscountYears(int discountYears) {
+      checkArgument(
+          Range.closed(1, 10).contains(discountYears),
+          "Discount years must be between 1 and 10 inclusive");
+      getInstance().discountYears = discountYears;
+      return this;
+    }
+
     public Builder setTokenType(TokenType tokenType) {
       checkState(getInstance().tokenType == null, "Token type can only be set once");
       getInstance().tokenType = tokenType;

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

@@ -155,15 +155,7 @@ public abstract class PollMessage extends ImmutableObject
   public abstract VKey<? extends PollMessage> createVKey();
 
   public static VKey<PollMessage> createVKey(Key<PollMessage> key) {
-    // TODO(b/159207551): As it stands, the SQL key generated here doesn't mesh with the primary key
-    // type for the table, which is a single long integer.  Also note that the class id is not
-    // correct here and as such the resulting key will not be loadable from SQL.
-    if (key == null) {
-      return null;
-    }
-    String path =
-        key.getParent().getParent().getName() + "/" + key.getParent().getId() + key.getId();
-    return VKey.create(PollMessage.class, path, key);
+    return VKey.create(PollMessage.class, key.getId(), key);
   }
 
   /** Override Buildable.asBuilder() to give this method stronger typing. */

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

@@ -140,6 +140,16 @@ public class Spec11ThreatMatch extends ImmutableObject implements Buildable, Sql
       return super.build();
     }
 
+    /**
+     * Manually set the ID for testing or other special circumstances.
+     *
+     * <p>In general the ID is generated by SQL and there should be no need to set it manually.
+     */
+    public Builder setId(Long id) {
+      getInstance().id = id;
+      return this;
+    }
+
     public Builder setDomainName(String domainName) {
       getInstance().domainName = domainName;
       getInstance().tld = DomainNameUtils.getTldFromDomainName(domainName);

core/src/main/java/google/registry/persistence/JpaRetries.java

@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableSet;
 import java.sql.SQLException;
 import java.util.function.Predicate;
 import javax.persistence.OptimisticLockException;
+import org.hibernate.exception.JDBCConnectionException;
 
 /** Helpers for identifying retriable database operations. */
 public final class JpaRetries {
@@ -40,6 +41,13 @@ public final class JpaRetries {
               e instanceof SQLException
                   && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
 
+  private static final Predicate<Throwable> RETRIABLE_QUERY_PREDICATE =
+      Predicates.or(
+          JDBCConnectionException.class::isInstance,
+          e ->
+              e instanceof SQLException
+                  && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
+
   public static boolean isFailedTxnRetriable(Throwable throwable) {
     Throwable t = throwable;
     while (t != null) {
@@ -53,6 +61,13 @@ public final class JpaRetries {
 
   public static boolean isFailedQueryRetriable(Throwable throwable) {
     // TODO(weiminyu): check for more error codes.
-    return isFailedTxnRetriable(throwable);
+    Throwable t = throwable;
+    while (t != null) {
+      if (RETRIABLE_QUERY_PREDICATE.test(t)) {
+        return true;
+      }
+      t = t.getCause();
+    }
+    return false;
   }
 }

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

@@ -15,6 +15,7 @@
 package google.registry.persistence.transaction;
 
 import google.registry.persistence.VKey;
+import java.util.function.Supplier;
 import javax.persistence.EntityManager;
 
 /** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
@@ -23,6 +24,12 @@ public interface JpaTransactionManager extends TransactionManager {
   /** Returns the {@link EntityManager} for the current request. */
   EntityManager getEntityManager();
 
+  /** Executes the work in a transaction with no retries and returns the result. */
+  <T> T transactNoRetry(Supplier<T> work);
+
+  /** Executes the work in a transaction with no retries. */
+  void transactNoRetry(Runnable work);
+
   /** Deletes the entity by its id, throws exception if the entity is not deleted. */
   public abstract <T> void assertDelete(VKey<T> key);
 

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

@@ -27,8 +27,11 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
+import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
+import google.registry.util.Retrier;
+import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
@@ -49,6 +52,7 @@ import org.joda.time.DateTime;
 public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private static final Retrier retrier = new Retrier(new SystemSleeper(), 3);
 
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
@@ -94,6 +98,40 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> T transact(Supplier<T> work) {
     // TODO(shicong): Investigate removing transactNew functionality after migration as it may
     //  be same as this one.
+    return retrier.callWithRetry(
+        () -> {
+          if (inTransaction()) {
+            return work.get();
+          }
+          TransactionInfo txnInfo = transactionInfo.get();
+          txnInfo.entityManager = emf.createEntityManager();
+          EntityTransaction txn = txnInfo.entityManager.getTransaction();
+          try {
+            txn.begin();
+            txnInfo.start(clock);
+            T result = work.get();
+            txnInfo.recordTransaction();
+            txn.commit();
+            return result;
+          } catch (RuntimeException | Error e) {
+            // Error is unchecked!
+            try {
+              txn.rollback();
+              logger.atWarning().log("Error during transaction; transaction rolled back");
+            } catch (Throwable rollbackException) {
+              logger.atSevere().withCause(rollbackException).log(
+                  "Rollback failed; suppressing error");
+            }
+            throw e;
+          } finally {
+            txnInfo.clear();
+          }
+        },
+        JpaRetries::isFailedTxnRetriable);
+  }
+
+  @Override
+  public <T> T transactNoRetry(Supplier<T> work) {
     if (inTransaction()) {
       return work.get();
     }
@@ -130,6 +168,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         });
   }
 
+  @Override
+  public void transactNoRetry(Runnable work) {
+    transactNoRetry(
+        () -> {
+          work.run();
+          return null;
+        });
+  }
+
   @Override
   public <T> T transactNew(Supplier<T> work) {
     return transact(work);
@@ -142,11 +189,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transactNewReadOnly(Supplier<T> work) {
-    return transact(
-        () -> {
-          getEntityManager().createNativeQuery("SET TRANSACTION READ ONLY").executeUpdate();
-          return work.get();
-        });
+    return retrier.callWithRetry(
+        () ->
+            transact(
+                () -> {
+                  getEntityManager().createNativeQuery("SET TRANSACTION READ ONLY").executeUpdate();
+                  return work.get();
+                }),
+        JpaRetries::isFailedQueryRetriable);
   }
 
   @Override
@@ -160,7 +210,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T doTransactionless(Supplier<T> work) {
-    return transact(work);
+    return retrier.callWithRetry(() -> transact(work), JpaRetries::isFailedQueryRetriable);
   }
 
   @Override

core/src/main/java/google/registry/tools/CreateOrUpdateDomainCommand.java

@@ -15,20 +15,11 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.util.CollectionUtils.findDuplicates;
 
-import com.beust.jcommander.IStringConverter;
 import com.beust.jcommander.Parameter;
-import com.google.auto.value.AutoValue;
-import com.google.common.base.Ascii;
-import com.google.common.base.CharMatcher;
 import com.google.common.base.Joiner;
-import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.io.BaseEncoding;
-import com.google.template.soy.data.SoyListData;
-import com.google.template.soy.data.SoyMapData;
 import google.registry.tools.params.NameserversParameter;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -80,77 +71,15 @@ abstract class CreateOrUpdateDomainCommand extends MutatingEppToolCommand {
   String password;
 
   @Parameter(
-    names = "--ds_records",
-    description =
-        "Comma-separated list of DS records. Each DS record is given as "
-        + "<keyTag> <alg> <digestType> <digest>, in order, as it appears in the Zonefile.",
-    converter = DsRecordConverter.class
-  )
+      names = "--ds_records",
+      description =
+          "Comma-separated list of DS records. Each DS record is given as "
+              + "<keyTag> <alg> <digestType> <digest>, in order, as it appears in the Zonefile.",
+      converter = DsRecord.Converter.class)
   List<DsRecord> dsRecords = new ArrayList<>();
 
   Set<String> domains;
 
-  @AutoValue
-  abstract static class DsRecord {
-    private static final Splitter SPLITTER =
-        Splitter.on(CharMatcher.whitespace()).omitEmptyStrings();
-
-    public abstract int keyTag();
-    public abstract int alg();
-    public abstract int digestType();
-    public abstract String digest();
-
-    private static DsRecord create(int keyTag, int alg, int digestType, String digest) {
-      digest = Ascii.toUpperCase(digest);
-      checkArgument(
-          BaseEncoding.base16().canDecode(digest),
-          "digest should be even-lengthed hex, but is %s (length %s)",
-          digest,
-          digest.length());
-      return new AutoValue_CreateOrUpdateDomainCommand_DsRecord(keyTag, alg, digestType, digest);
-    }
-
-    /**
-     * Parses a string representation of the DS record.
-     *
-     * <p>The string format accepted is "[keyTag] [alg] [digestType] [digest]" (i.e., the 4
-     * arguments separated by any number of spaces, as it appears in the Zone file)
-     */
-    public static DsRecord parse(String dsRecord) {
-      List<String> elements = SPLITTER.splitToList(dsRecord);
-      checkArgument(
-          elements.size() == 4,
-          "dsRecord %s should have 4 parts, but has %s",
-          dsRecord,
-          elements.size());
-      return DsRecord.create(
-          Integer.parseUnsignedInt(elements.get(0)),
-          Integer.parseUnsignedInt(elements.get(1)),
-          Integer.parseUnsignedInt(elements.get(2)),
-          elements.get(3));
-    }
-
-    public SoyMapData toSoyData() {
-      return new SoyMapData(
-          "keyTag", keyTag(),
-          "alg", alg(),
-          "digestType", digestType(),
-          "digest", digest());
-    }
-
-    public static SoyListData convertToSoy(List<DsRecord> dsRecords) {
-      return new SoyListData(
-          dsRecords.stream().map(DsRecord::toSoyData).collect(toImmutableList()));
-    }
-  }
-
-  public static class DsRecordConverter implements IStringConverter<DsRecord> {
-    @Override
-    public DsRecord convert(String dsRecord) {
-      return DsRecord.parse(dsRecord);
-    }
-  }
-
   @Override
   protected void initEppToolCommand() throws Exception {
     checkArgument(nameservers.size() <= 13, "There can be at most 13 nameservers.");

core/src/main/java/google/registry/tools/DeploySpec11PipelineCommand.java

@@ -14,7 +14,10 @@
 
 package google.registry.tools;
 
+import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.JpaSupplierFactory;
 import google.registry.beam.spec11.Spec11Pipeline;
 import google.registry.config.CredentialModule.LocalCredential;
 import google.registry.config.RegistryConfig.Config;
@@ -31,6 +34,12 @@ public class DeploySpec11PipelineCommand implements Command {
   @Config("projectId")
   String projectId;
 
+  @Parameter(
+      names = {"-p", "--project"},
+      description = "Cloud KMS project ID",
+      required = true)
+  String cloudKmsProjectId;
+
   @Inject
   @Config("beamStagingUrl")
   String beamStagingUrl;
@@ -53,12 +62,19 @@ public class DeploySpec11PipelineCommand implements Command {
 
   @Override
   public void run() {
+    JpaSupplierFactory jpaSupplierFactory =
+        new JpaSupplierFactory(
+            sqlAccessInfoFile,
+            cloudKmsProjectId,
+            JpaTransactionManagerComponent::cloudSqlJpaTransactionManager);
+
     Spec11Pipeline pipeline =
         new Spec11Pipeline(
             projectId,
             beamStagingUrl,
             spec11TemplateUrl,
             reportingBucketUrl,
+            jpaSupplierFactory,
             googleCredentialsBundle,
             retrier);
     pipeline.deploy();

core/src/main/java/google/registry/tools/DsRecord.java

@@ -0,0 +1,90 @@
+// Copyright 2017 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.tools;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+import com.beust.jcommander.IStringConverter;
+import com.google.auto.value.AutoValue;
+import com.google.common.base.Ascii;
+import com.google.common.base.CharMatcher;
+import com.google.common.base.Splitter;
+import com.google.common.io.BaseEncoding;
+import com.google.template.soy.data.SoyListData;
+import com.google.template.soy.data.SoyMapData;
+import java.util.List;
+
+@AutoValue
+abstract class DsRecord {
+  private static final Splitter SPLITTER = Splitter.on(CharMatcher.whitespace()).omitEmptyStrings();
+
+  public abstract int keyTag();
+
+  public abstract int alg();
+
+  public abstract int digestType();
+
+  public abstract String digest();
+
+  private static DsRecord create(int keyTag, int alg, int digestType, String digest) {
+    digest = Ascii.toUpperCase(digest);
+    checkArgument(
+        BaseEncoding.base16().canDecode(digest),
+        "digest should be even-lengthed hex, but is %s (length %s)",
+        digest,
+        digest.length());
+    return new AutoValue_DsRecord(keyTag, alg, digestType, digest);
+  }
+
+  /**
+   * Parses a string representation of the DS record.
+   *
+   * <p>The string format accepted is "[keyTag] [alg] [digestType] [digest]" (i.e., the 4 arguments
+   * separated by any number of spaces, as it appears in the Zone file)
+   */
+  public static DsRecord parse(String dsRecord) {
+    List<String> elements = SPLITTER.splitToList(dsRecord);
+    checkArgument(
+        elements.size() == 4,
+        "dsRecord %s should have 4 parts, but has %s",
+        dsRecord,
+        elements.size());
+    return DsRecord.create(
+        Integer.parseUnsignedInt(elements.get(0)),
+        Integer.parseUnsignedInt(elements.get(1)),
+        Integer.parseUnsignedInt(elements.get(2)),
+        elements.get(3));
+  }
+
+  public SoyMapData toSoyData() {
+    return new SoyMapData(
+        "keyTag", keyTag(),
+        "alg", alg(),
+        "digestType", digestType(),
+        "digest", digest());
+  }
+
+  public static SoyListData convertToSoy(List<DsRecord> dsRecords) {
+    return new SoyListData(dsRecords.stream().map(DsRecord::toSoyData).collect(toImmutableList()));
+  }
+
+  public static class Converter implements IStringConverter<DsRecord> {
+    @Override
+    public DsRecord convert(String dsRecord) {
+      return DsRecord.parse(dsRecord);
+    }
+  }
+}

core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -115,7 +115,20 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
       description =
           "A discount off the base price for the first year between 0.0 and 1.0. Default is 0.0,"
               + " i.e. no discount.")
-  private double discountFraction;
+  private Double discountFraction;
+
+  @Parameter(
+      names = {"--discount_premiums"},
+      description =
+          "Whether the discount is valid for premium names in addition to standard ones. Default"
+              + " is false.",
+      arity = 1)
+  private Boolean discountPremiums;
+
+  @Parameter(
+      names = {"--discount_years"},
+      description = "The number of years the discount applies for. Default is 1, max value is 10.")
+  private Integer discountYears;
 
   @Parameter(
       names = "--token_status_transitions",
@@ -170,8 +183,10 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
                             .setToken(t)
                             .setTokenType(tokenType == null ? SINGLE_USE : tokenType)
                             .setAllowedClientIds(ImmutableSet.copyOf(nullToEmpty(allowedClientIds)))
-                            .setAllowedTlds(ImmutableSet.copyOf(nullToEmpty(allowedTlds)))
-                            .setDiscountFraction(discountFraction);
+                            .setAllowedTlds(ImmutableSet.copyOf(nullToEmpty(allowedTlds)));
+                    Optional.ofNullable(discountFraction).ifPresent(token::setDiscountFraction);
+                    Optional.ofNullable(discountPremiums).ifPresent(token::setDiscountPremiums);
+                    Optional.ofNullable(discountYears).ifPresent(token::setDiscountYears);
                     Optional.ofNullable(tokenStatusTransitions)
                         .ifPresent(token::setTokenStatusTransitions);
                     Optional.ofNullable(domainNames)

core/src/main/java/google/registry/tools/GetAllocationTokenCommand.java

@@ -27,7 +27,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
 import java.util.Collection;
 import java.util.List;
-import java.util.Objects;
+import java.util.Optional;
 
 /** Command to show allocation tokens. */
 @Parameters(separators = " =", commandDescription = "Show allocation token(s)")
@@ -55,11 +55,11 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
       if (loadedTokens.containsKey(token)) {
         AllocationToken loadedToken = loadedTokens.get(token);
         System.out.println(loadedToken.toString());
-        if (loadedToken.getRedemptionHistoryEntry() == null) {
+        if (!loadedToken.getRedemptionHistoryEntry().isPresent()) {
           System.out.printf("Token %s was not redeemed.\n", token);
         } else {
           DomainBase domain =
-              domains.get(loadedToken.getRedemptionHistoryEntry().<DomainBase>getParent());
+              domains.get(loadedToken.getRedemptionHistoryEntry().get().<DomainBase>getParent());
           if (domain == null) {
             System.out.printf("ERROR: Token %s was redeemed but domain can't be loaded.\n", token);
           } else {
@@ -80,7 +80,8 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
     ImmutableList<Key<DomainBase>> domainKeys =
         tokens.stream()
             .map(AllocationToken::getRedemptionHistoryEntry)
-            .filter(Objects::nonNull)
+            .filter(Optional::isPresent)
+            .map(Optional::get)
             .map(Key::<DomainBase>getParent)
             .collect(toImmutableList());
     ImmutableMap.Builder<Key<DomainBase>, DomainBase> domainsBuilder = new ImmutableMap.Builder<>();

core/src/main/java/google/registry/tools/MutatingCommand.java

@@ -143,10 +143,14 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
   protected String execute() throws Exception {
     for (final List<EntityChange> batch : getCollatedEntityChangeBatches()) {
       tm().transact(() -> batch.forEach(this::executeChange));
+      postBatchExecute();
     }
     return String.format("Updated %d entities.\n", changedEntitiesMap.size());
   }
 
+  /** Performs any execution step after each batch. */
+  protected void postBatchExecute() {}
+
   private void executeChange(EntityChange change) {
     // Load the key of the entity to mutate and double-check that it hasn't been
     // modified from the version that existed when the change was prepared.

core/src/main/java/google/registry/tools/RegistryTool.java

@@ -99,6 +99,7 @@ public final class RegistryTool {
           .put("remove_ip_address", RemoveIpAddressCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
           .put("resave_entities", ResaveEntitiesCommand.class)
+          .put("resave_entities_with_unique_id", ResaveEntitiesWithUniqueIdCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
           .put("resave_epp_resource", ResaveEppResourceCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)

core/src/main/java/google/registry/tools/ResaveEntitiesWithUniqueIdCommand.java

@@ -0,0 +1,196 @@
+// 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.tools;
+
+import static com.google.common.base.Preconditions.checkState;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.appengine.api.datastore.KeyFactory;
+import com.google.common.base.Splitter;
+import com.google.common.io.CharStreams;
+import com.google.common.io.Files;
+import com.googlecode.objectify.Key;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.domain.DomainBase;
+import google.registry.util.NonFinalForTesting;
+import java.io.File;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.util.List;
+
+/**
+ * Command to resave entities with a unique id.
+ *
+ * <p>This command is used to address the duplicate id issue we found for certain {@link
+ * BillingEvent.OneTime} entities. The command reassigns an application wide unique id to the
+ * problematic entity and resaves it, it also resaves the entity having reference to the problematic
+ * entity with the updated id.
+ *
+ * <p>To use this command, you will need to provide the path to a file containing a list of strings
+ * representing the literal of Objectify key for the problematic entities. An example key literal
+ * is:
+ *
+ * <pre>
+ * "DomainBase", "111111-TEST", "HistoryEntry", 2222222, "OneTime", 3333333
+ * </pre>
+ *
+ * <p>Note that the double quotes are part of the key literal. The key literal can be retrieved from
+ * the column <code>__key__.path</code> in BigQuery.
+ */
+@Parameters(separators = " =", commandDescription = "Resave entities with a unique id.")
+public class ResaveEntitiesWithUniqueIdCommand extends MutatingCommand {
+
+  @Parameter(
+      names = "--key_paths_file",
+      description =
+          "Key paths file name, each line in the file should be a key literal. An example key"
+              + " literal is: \"DomainBase\", \"111111-TEST\", \"HistoryEntry\", 2222222,"
+              + " \"OneTime\", 3333333")
+  File keyPathsFile;
+
+  @NonFinalForTesting private static InputStream stdin = System.in;
+
+  private String keyChangeMessage;
+
+  @Override
+  protected void init() throws Exception {
+    List<String> keyPaths =
+        keyPathsFile == null
+            ? CharStreams.readLines(new InputStreamReader(stdin, UTF_8))
+            : Files.readLines(keyPathsFile, UTF_8);
+    for (String keyPath : keyPaths) {
+      Key<?> untypedKey = parseKeyPath(keyPath);
+      Object entity = ofy().load().key(untypedKey).now();
+      if (entity == null) {
+        System.err.println(
+            String.format(
+                "Entity %s read from %s doesn't exist in Datastore! Skipping.",
+                untypedKey,
+                keyPathsFile == null ? "STDIN" : "File " + keyPathsFile.getAbsolutePath()));
+        continue;
+      }
+      if (entity instanceof BillingEvent.OneTime) {
+        resaveBillingEvent((BillingEvent.OneTime) entity);
+      } else {
+        throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
+      }
+      flushTransaction();
+    }
+  }
+
+  @Override
+  protected void postBatchExecute() {
+    System.out.println(keyChangeMessage);
+  }
+
+  private void deleteOldAndSaveNewEntity(ImmutableObject oldEntity, ImmutableObject newEntity) {
+    stageEntityChange(oldEntity, null);
+    stageEntityChange(null, newEntity);
+  }
+
+  private void resaveBillingEvent(BillingEvent.OneTime billingEvent) {
+    Key<BillingEvent> key = Key.create(billingEvent);
+    Key<DomainBase> domainKey = getGrandParentAsDomain(key);
+    DomainBase domain = ofy().load().key(domainKey).now();
+
+    // The BillingEvent.OneTime entity to be resaved should be the billing event created a few
+    // years ago, so they should not be referenced from TransferData and GracePeriod in the domain.
+    assertNotInDomainTransferData(domain, key);
+    domain
+        .getGracePeriods()
+        .forEach(
+            gracePeriod ->
+                checkState(
+                    !gracePeriod.getOneTimeBillingEvent().getOfyKey().equals(key),
+                    "Entity %s is referenced by a grace period in domain %s",
+                    key,
+                    domainKey));
+
+    // By setting id to 0L, Buildable.build() will assign an application wide unique id to it.
+    BillingEvent.OneTime uniqIdBillingEvent = billingEvent.asBuilder().setId(0L).build();
+    deleteOldAndSaveNewEntity(billingEvent, uniqIdBillingEvent);
+    keyChangeMessage =
+        String.format("Old Entity Key: %s New Entity Key: %s", key, Key.create(uniqIdBillingEvent));
+  }
+
+  private static boolean isKind(Key<?> key, Class<?> clazz) {
+    return key.getKind().equals(Key.getKind(clazz));
+  }
+
+  static Key<?> parseKeyPath(String keyPath) {
+    List<String> keyComponents = Splitter.on(',').splitToList(keyPath);
+    checkState(
+        keyComponents.size() > 0 && keyComponents.size() % 2 == 0,
+        "Invalid number of key components");
+    com.google.appengine.api.datastore.Key rawKey = null;
+    for (int i = 0, j = 1; j < keyComponents.size(); i += 2, j += 2) {
+      String kindLiteral = keyComponents.get(i).trim();
+      String idOrNameLiteral = keyComponents.get(j).trim();
+      rawKey = createDatastoreKey(rawKey, kindLiteral, idOrNameLiteral);
+    }
+    return Key.create(rawKey);
+  }
+
+  private static com.google.appengine.api.datastore.Key createDatastoreKey(
+      com.google.appengine.api.datastore.Key parent, String kindLiteral, String idOrNameLiteral) {
+    if (isLiteralString(idOrNameLiteral)) {
+      return KeyFactory.createKey(parent, removeQuotes(kindLiteral), removeQuotes(idOrNameLiteral));
+    } else {
+      return KeyFactory.createKey(
+          parent, removeQuotes(kindLiteral), Long.parseLong(idOrNameLiteral));
+    }
+  }
+
+  private static boolean isLiteralString(String raw) {
+    return raw.charAt(0) == '"' && raw.charAt(raw.length() - 1) == '"';
+  }
+
+  private static String removeQuotes(String literal) {
+    return literal.substring(1, literal.length() - 1);
+  }
+
+  private static Key<DomainBase> getGrandParentAsDomain(Key<?> key) {
+    Key<?> grandParent;
+    try {
+      grandParent = key.getParent().getParent();
+    } catch (Throwable e) {
+      throw new IllegalArgumentException("Error retrieving grand parent key", e);
+    }
+    if (!isKind(grandParent, DomainBase.class)) {
+      throw new IllegalArgumentException(
+          String.format("Expected a Key<DomainBase> but got %s", grandParent));
+    }
+    return (Key<DomainBase>) grandParent;
+  }
+
+  private static void assertNotInDomainTransferData(DomainBase domainBase, Key<?> key) {
+    if (!domainBase.getTransferData().isEmpty()) {
+      domainBase
+          .getTransferData()
+          .getServerApproveEntities()
+          .forEach(
+              entityKey ->
+                  checkState(
+                      !entityKey.getOfyKey().equals(key),
+                      "Entity %s is referenced by the transfer data in domain %s",
+                      key,
+                      domainBase.createVKey().getOfyKey()));
+    }
+  }
+}

core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java

@@ -15,7 +15,7 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Strings.nullToEmpty;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Sets.difference;
 import static google.registry.model.EppResourceUtils.checkResourcesExist;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
@@ -26,9 +26,11 @@ import static org.joda.time.DateTimeZone.UTC;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
+import com.google.template.soy.data.SoyListData;
 import com.google.template.soy.data.SoyMapData;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.secdns.DelegationSignerData;
@@ -37,14 +39,10 @@ import google.registry.model.host.HostResource;
 import google.registry.tools.soy.UniformRapidSuspensionSoyInfo;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import javax.xml.bind.annotation.adapters.HexBinaryAdapter;
 import org.joda.time.DateTime;
-import org.json.simple.JSONArray;
-import org.json.simple.JSONValue;
-import org.json.simple.parser.ParseException;
 
 /** A command to suspend a domain for the Uniform Rapid Suspension process. */
 @Parameters(separators = " =",
@@ -59,9 +57,6 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
   /** Client id that made this change. Only recorded in the history entry. **/
   private static final String CLIENT_ID = "CharlestonRoad";
 
-  private static final ImmutableSet<String> DSDATA_FIELDS =
-      ImmutableSet.of("keyTag", "alg", "digestType", "digest");
-
   @Parameter(
       names = {"-n", "--domain_name"},
       description = "Domain to suspend.",
@@ -76,10 +71,12 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
 
   @Parameter(
       names = {"-s", "--dsdata"},
-      description = "Comma-delimited set of dsdata to replace the current dsdata on the domain, "
-          + "where each dsdata is represented as a JSON object with fields 'keyTag', 'alg', "
-          + "'digestType' and 'digest'.")
-  private String newDsData;
+      description =
+          "Comma-delimited set of dsdata to replace the current dsdata on the domain, "
+              + "Each DS record is given as <keyTag> <alg> <digestType> <digest>, in order, as it "
+              + "appears in the Zonefile.",
+      converter = DsRecord.Converter.class)
+  private List<DsRecord> newDsData;
 
   @Parameter(
       names = {"-p", "--locks_to_preserve"},
@@ -88,6 +85,13 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
               + "locks: serverDeleteProhibited, serverTransferProhibited, serverUpdateProhibited")
   private List<String> locksToPreserve = new ArrayList<>();
 
+  @Parameter(
+      names = {"--restore_client_hold"},
+      description =
+          "Restores a CLIENT_HOLD status that was previously removed for a URS suspension (only "
+              + "valid with --undo).")
+  private boolean restoreClientHold;
+
   @Parameter(
       names = {"--undo"},
       description = "Flag indicating that is is an undo command, which removes locks.")
@@ -100,28 +104,16 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
   ImmutableSortedSet<String> existingNameservers;
 
   /** Set of existing dsdata jsons that need to be restored during undo, sorted for nicer output. */
-  ImmutableSortedSet<String> existingDsData;
+  ImmutableList<ImmutableMap<String, Object>> existingDsData;
+
+  /** Set of status values to remove. */
+  ImmutableSet<String> removeStatuses;
 
   @Override
   protected void initMutatingEppToolCommand() {
     superuser = true;
     DateTime now = DateTime.now(UTC);
     ImmutableSet<String> newHostsSet = ImmutableSet.copyOf(newHosts);
-    ImmutableSet.Builder<Map<String, Object>> newDsDataBuilder = new ImmutableSet.Builder<>();
-    try {
-      // Add brackets around newDsData to convert it to a parsable JSON array.
-      String jsonArrayString = String.format("[%s]", nullToEmpty(newDsData));
-      for (Object dsData : (JSONArray) JSONValue.parseWithException(jsonArrayString)) {
-        @SuppressWarnings("unchecked")
-        Map<String, Object> dsDataJson = (Map<String, Object>) dsData;
-        checkArgument(
-            dsDataJson.keySet().equals(DSDATA_FIELDS),
-            "Incorrect fields on --dsdata JSON: " + JSONValue.toJSONString(dsDataJson));
-        newDsDataBuilder.add(dsDataJson);
-      }
-    } catch (ClassCastException | ParseException e) {
-      throw new IllegalArgumentException("Invalid --dsdata JSON", e);
-    }
     Optional<DomainBase> domain = loadByForeignKey(DomainBase.class, domainName, now);
     checkArgumentPresent(domain, "Domain '%s' does not exist or is deleted", domainName);
     Set<String> missingHosts =
@@ -133,18 +125,39 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
     existingNameservers = getExistingNameservers(domain.get());
     existingLocks = getExistingLocks(domain.get());
     existingDsData = getExistingDsData(domain.get());
+    removeStatuses =
+        (hasClientHold(domain.get()) && !undo)
+            ? ImmutableSet.of(StatusValue.CLIENT_HOLD.getXmlName())
+            : ImmutableSet.of();
+    ImmutableSet<String> statusesToApply;
+    if (undo) {
+      statusesToApply =
+          restoreClientHold
+              ? ImmutableSet.of(StatusValue.CLIENT_HOLD.getXmlName())
+              : ImmutableSet.of();
+    } else {
+      statusesToApply = URS_LOCKS;
+    }
     setSoyTemplate(
         UniformRapidSuspensionSoyInfo.getInstance(),
         UniformRapidSuspensionSoyInfo.UNIFORMRAPIDSUSPENSION);
-    addSoyRecord(CLIENT_ID, new SoyMapData(
-        "domainName", domainName,
-        "hostsToAdd", difference(newHostsSet, existingNameservers),
-        "hostsToRemove", difference(existingNameservers, newHostsSet),
-        "locksToApply", undo ? ImmutableSet.of() : URS_LOCKS,
-        "locksToRemove",
-            undo ? difference(URS_LOCKS, ImmutableSet.copyOf(locksToPreserve)) : ImmutableSet.of(),
-        "newDsData", newDsDataBuilder.build(),
-        "reason", (undo ? "Undo " : "") + "Uniform Rapid Suspension"));
+    addSoyRecord(
+        CLIENT_ID,
+        new SoyMapData(
+            "domainName",
+            domainName,
+            "hostsToAdd",
+            difference(newHostsSet, existingNameservers),
+            "hostsToRemove",
+            difference(existingNameservers, newHostsSet),
+            "statusesToApply",
+            statusesToApply,
+            "statusesToRemove",
+            undo ? difference(URS_LOCKS, ImmutableSet.copyOf(locksToPreserve)) : removeStatuses,
+            "newDsData",
+            newDsData != null ? DsRecord.convertToSoy(newDsData) : new SoyListData(),
+            "reason",
+            (undo ? "Undo " : "") + "Uniform Rapid Suspension"));
   }
 
   private ImmutableSortedSet<String> getExistingNameservers(DomainBase domain) {
@@ -165,15 +178,25 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
     return locks.build();
   }
 
-  private ImmutableSortedSet<String> getExistingDsData(DomainBase domain) {
-    ImmutableSortedSet.Builder<String> dsDataJsons = ImmutableSortedSet.naturalOrder();
+  private boolean hasClientHold(DomainBase domain) {
+    for (StatusValue status : domain.getStatusValues()) {
+      if (status == StatusValue.CLIENT_HOLD) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private ImmutableList<ImmutableMap<String, Object>> getExistingDsData(DomainBase domain) {
+    ImmutableList.Builder<ImmutableMap<String, Object>> dsDataJsons = new ImmutableList.Builder();
     HexBinaryAdapter hexBinaryAdapter = new HexBinaryAdapter();
     for (DelegationSignerData dsData : domain.getDsData()) {
-      dsDataJsons.add(JSONValue.toJSONString(ImmutableMap.of(
-          "keyTag", dsData.getKeyTag(),
-          "algorithm", dsData.getAlgorithm(),
-          "digestType", dsData.getDigestType(),
-          "digest", hexBinaryAdapter.marshal(dsData.getDigest()))));
+      dsDataJsons.add(
+          ImmutableMap.of(
+              "keyTag", dsData.getKeyTag(),
+              "algorithm", dsData.getAlgorithm(),
+              "digestType", dsData.getDigestType(),
+              "digest", hexBinaryAdapter.marshal(dsData.getDigest())));
     }
     return dsDataJsons.build();
   }
@@ -194,8 +217,23 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
     if (!existingLocks.isEmpty()) {
       undoBuilder.append(" --locks_to_preserve ").append(Joiner.on(',').join(existingLocks));
     }
+    if (removeStatuses.contains(StatusValue.CLIENT_HOLD.getXmlName())) {
+      undoBuilder.append(" --restore_client_hold");
+    }
     if (!existingDsData.isEmpty()) {
-      undoBuilder.append(" --dsdata ").append(Joiner.on(',').join(existingDsData));
+      ImmutableList<String> formattedDsRecords =
+          existingDsData.stream()
+              .map(
+                  rec ->
+                      String.format(
+                          "%s %s %s %s",
+                          rec.get("keyTag"),
+                          rec.get("algorithm"),
+                          rec.get("digestType"),
+                          rec.get("digest")))
+              .sorted()
+              .collect(toImmutableList());
+      undoBuilder.append(" --dsdata ").append(Joiner.on(',').join(formattedDsRecords));
     }
     return undoBuilder.toString();
   }

core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java

@@ -66,6 +66,19 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
               + "i.e. no discount.")
   private Double discountFraction;
 
+  @Parameter(
+      names = {"--discount_premiums"},
+      description =
+          "Whether the discount is valid for premium names in addition to standard ones. Default"
+              + " is false.",
+      arity = 1)
+  private Boolean discountPremiums;
+
+  @Parameter(
+      names = {"--discount_years"},
+      description = "The number of years the discount applies for. Default is 1, max value is 10.")
+  private Integer discountYears;
+
   @Parameter(
       names = "--token_status_transitions",
       converter = TokenStatusTransitions.class,
@@ -122,6 +135,8 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     Optional.ofNullable(allowedTlds)
         .ifPresent(tlds -> builder.setAllowedTlds(ImmutableSet.copyOf(tlds)));
     Optional.ofNullable(discountFraction).ifPresent(builder::setDiscountFraction);
+    Optional.ofNullable(discountPremiums).ifPresent(builder::setDiscountPremiums);
+    Optional.ofNullable(discountYears).ifPresent(builder::setDiscountYears);
     Optional.ofNullable(tokenStatusTransitions).ifPresent(builder::setTokenStatusTransitions);
     return builder.build();
   }

core/src/main/java/google/registry/tools/UpdateDomainCommand.java

@@ -76,10 +76,10 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
   private List<String> addStatuses = new ArrayList<>();
 
   @Parameter(
-    names = "--add_ds_records",
-    description = "DS records to add. Cannot be set if --ds_records or --clear_ds_records is set.",
-    converter = DsRecordConverter.class
-  )
+      names = "--add_ds_records",
+      description =
+          "DS records to add. Cannot be set if --ds_records or --clear_ds_records is set.",
+      converter = DsRecord.Converter.class)
   private List<DsRecord> addDsRecords = new ArrayList<>();
 
   @Parameter(
@@ -110,11 +110,10 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
   private List<String> removeStatuses = new ArrayList<>();
 
   @Parameter(
-    names = "--remove_ds_records",
-    description =
-        "DS records to remove. Cannot be set if --ds_records or --clear_ds_records is set.",
-    converter = DsRecordConverter.class
-  )
+      names = "--remove_ds_records",
+      description =
+          "DS records to remove. Cannot be set if --ds_records or --clear_ds_records is set.",
+      converter = DsRecord.Converter.class)
   private List<DsRecord> removeDsRecords = new ArrayList<>();
 
   @Parameter(

core/src/main/resources/META-INF/persistence.xml

@@ -25,6 +25,7 @@
     <class>google.registry.model.contact.ContactHistory</class>
     <class>google.registry.model.contact.ContactResource</class>
     <class>google.registry.model.domain.DomainBase</class>
+    <class>google.registry.model.domain.DomainHistory</class>
     <class>google.registry.model.host.HostHistory</class>
     <class>google.registry.model.host.HostResource</class>
     <class>google.registry.model.registrar.Registrar</class>

core/src/main/resources/google/registry/tools/soy/UniformRapidSuspension.soy

@@ -21,8 +21,8 @@
 {@param domainName: string}
 {@param hostsToAdd: list<string>}
 {@param hostsToRemove: list<string>}
-{@param locksToApply: list<string>}
-{@param locksToRemove: list<string>}
+{@param statusesToApply: list<string>}
+{@param statusesToRemove: list<string>}
 {@param newDsData: list<[keyTag:int, alg:int, digestType:int, digest:string]>}
 {@param reason: string}
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
@@ -39,7 +39,7 @@
           {/for}
           </domain:ns>
         {/if}
-        {for $la in $locksToApply}
+        {for $la in $statusesToApply}
           <domain:status s="{$la}" />
         {/for}
         </domain:add>
@@ -51,7 +51,7 @@
           {/for}
           </domain:ns>
         {/if}
-        {for $lr in $locksToRemove}
+        {for $lr in $statusesToRemove}
           <domain:status s="{$lr}" />
         {/for}
         </domain:rem>

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

@@ -117,9 +117,10 @@ public class ResaveEntityActionTest {
 
   @Test
   void test_domainPendingDeletion_isResavedAndReenqueued() {
+    DomainBase newDomain = newDomainBase("domain.tld");
     DomainBase domain =
         persistResource(
-            newDomainBase("domain.tld")
+            newDomain
                 .asBuilder()
                 .setDeletionTime(clock.nowUtc().plusDays(35))
                 .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
@@ -127,6 +128,7 @@ public class ResaveEntityActionTest {
                     ImmutableSet.of(
                         GracePeriod.createWithoutBillingEvent(
                             GracePeriodStatus.REDEMPTION,
+                            newDomain.getRepoId(),
                             clock.nowUtc().plusDays(30),
                             "TheRegistrar")))
                 .build());

core/src/test/java/google/registry/beam/initsql/DomainBaseUtilTest.java

@@ -161,6 +161,7 @@ public class DomainBaseUtilTest {
                     .addGracePeriod(
                         GracePeriod.create(
                             GracePeriodStatus.ADD,
+                            "4-COM",
                             fakeClock.nowUtc().plusDays(1),
                             "registrar",
                             null))

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -214,7 +214,11 @@ class InitSqlPipelineTest {
                   .setSmdId("smdid")
                   .addGracePeriod(
                       GracePeriod.create(
-                          GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(1), "registrar", null))
+                          GracePeriodStatus.ADD,
+                          "4-COM",
+                          fakeClock.nowUtc().plusDays(1),
+                          "registrar",
+                          null))
                   .build());
       exportDir = store.export(exportRootDir.getAbsolutePath(), ALL_KINDS, ImmutableSet.of());
       commitLogDir = Files.createDirectory(tmpDir.resolve("commits")).toFile();

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -17,15 +17,22 @@ package google.registry.beam.spec11;
 import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.withSettings;
 
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.io.CharStreams;
 import google.registry.beam.TestPipelineExtension;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
+import google.registry.model.reporting.Spec11ThreatMatch;
+import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
 import google.registry.util.GoogleCredentialsBundle;
@@ -55,22 +62,42 @@ import org.apache.http.entity.BasicHttpEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.message.BasicStatusLine;
 import org.joda.time.DateTime;
+import org.joda.time.LocalDate;
+import org.joda.time.format.ISODateTimeFormat;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.api.io.TempDir;
+import org.mockito.Mock;
 import org.mockito.invocation.InvocationOnMock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
 /** Unit tests for {@link Spec11Pipeline}. */
+@ExtendWith(MockitoExtension.class)
+@MockitoSettings(strictness = Strictness.LENIENT)
 class Spec11PipelineTest {
+  private static class SaveNewThreatMatchAnswer implements Answer<Void>, Serializable {
+    @Override
+    public Void answer(InvocationOnMock invocation) {
+      Runnable runnable = invocation.getArgument(0, Runnable.class);
+      runnable.run();
+      return null;
+    }
+  }
 
   private static PipelineOptions pipelineOptions;
 
+  @Mock(serializable = true)
+  private static JpaTransactionManager mockJpaTm;
+
   @BeforeAll
   static void beforeAll() {
     pipelineOptions = PipelineOptionsFactory.create();
@@ -93,20 +120,23 @@ class Spec11PipelineTest {
   void beforeEach() throws IOException {
     String beamTempFolder =
         Files.createDirectory(tmpDir.resolve("beam_temp")).toAbsolutePath().toString();
+
     spec11Pipeline =
         new Spec11Pipeline(
             "test-project",
             beamTempFolder + "/staging",
             beamTempFolder + "/templates/invoicing",
             tmpDir.toAbsolutePath().toString(),
+            () -> mockJpaTm,
             GoogleCredentialsBundle.create(GoogleCredentials.create(null)),
             retrier);
   }
 
   private static final ImmutableList<String> BAD_DOMAINS =
-      ImmutableList.of("111.com", "222.com", "444.com", "no-email.com");
+      ImmutableList.of(
+          "111.com", "222.com", "444.com", "no-email.com", "testThreatMatchToSqlBad.com");
 
-  private ImmutableList<Subdomain> getInputDomains() {
+  private ImmutableList<Subdomain> getInputDomainsJson() {
     ImmutableList.Builder<Subdomain> subdomainsBuilder = new ImmutableList.Builder<>();
     // Put in at least 2 batches worth (x > 490) to guarantee multiple executions.
     // Put in half for theRegistrar and half for someRegistrar
@@ -134,17 +164,18 @@ class Spec11PipelineTest {
   @SuppressWarnings("unchecked")
   void testEndToEndPipeline_generatesExpectedFiles() throws Exception {
     // Establish mocks for testing
-    ImmutableList<Subdomain> inputRows = getInputDomains();
-    CloseableHttpClient httpClient = mock(CloseableHttpClient.class, withSettings().serializable());
+    ImmutableList<Subdomain> inputRows = getInputDomainsJson();
+    CloseableHttpClient mockHttpClient =
+        mock(CloseableHttpClient.class, withSettings().serializable());
 
     // Return a mock HttpResponse that returns a JSON response based on the request.
-    when(httpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder());
+    when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder());
 
     EvaluateSafeBrowsingFn evalFn =
         new EvaluateSafeBrowsingFn(
             StaticValueProvider.of("apikey"),
             new Retrier(new FakeSleeper(new FakeClock()), 3),
-            (Serializable & Supplier) () -> httpClient);
+            (Serializable & Supplier) () -> mockHttpClient);
 
     // Apply input and evaluation transforms
     PCollection<Subdomain> input = testPipeline.apply(Create.of(inputRows));
@@ -207,6 +238,56 @@ class Spec11PipelineTest {
                 .toString());
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testSpec11ThreatMatchToSql() throws Exception {
+    doAnswer(new SaveNewThreatMatchAnswer()).when(mockJpaTm).transact(any(Runnable.class));
+
+    // Create one bad and one good Subdomain to test with evaluateUrlHealth. Only the bad one should
+    // be detected and persisted.
+    Subdomain badDomain =
+        Subdomain.create(
+            "testThreatMatchToSqlBad.com", "theDomain", "theRegistrar", "fake@theRegistrar.com");
+    Subdomain goodDomain =
+        Subdomain.create(
+            "testThreatMatchToSqlGood.com",
+            "someDomain",
+            "someRegistrar",
+            "fake@someRegistrar.com");
+
+    // Establish a mock HttpResponse that returns a JSON response based on the request.
+    CloseableHttpClient mockHttpClient =
+        mock(CloseableHttpClient.class, withSettings().serializable());
+    when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder());
+
+    EvaluateSafeBrowsingFn evalFn =
+        new EvaluateSafeBrowsingFn(
+            StaticValueProvider.of("apikey"),
+            new Retrier(new FakeSleeper(new FakeClock()), 3),
+            (Serializable & Supplier) () -> mockHttpClient);
+
+    // Apply input and evaluation transforms
+    PCollection<Subdomain> input = testPipeline.apply(Create.of(badDomain, goodDomain));
+    spec11Pipeline.evaluateUrlHealth(input, evalFn, StaticValueProvider.of("2020-06-10"));
+    testPipeline.run();
+
+    // Verify that the expected threat created from the bad Subdomain and the persisted
+    // Spec11TThreatMatch are equal.
+    Spec11ThreatMatch expected =
+        new Spec11ThreatMatch()
+            .asBuilder()
+            .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
+            .setCheckDate(LocalDate.parse("2020-06-10", ISODateTimeFormat.date()))
+            .setDomainName(badDomain.domainName())
+            .setDomainRepoId(badDomain.domainRepoId())
+            .setRegistrarId(badDomain.registrarId())
+            .build();
+
+    verify(mockJpaTm).transact(any(Runnable.class));
+    verify(mockJpaTm).saveNew(expected);
+    verifyNoMoreInteractions(mockJpaTm);
+  }
+
   /**
    * A serializable {@link Answer} that returns a mock HTTP response based on the HTTP request's
    * content.

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

@@ -16,15 +16,12 @@ package google.registry.flows;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.eppcommon.EppXmlTransformer.marshal;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.testing.DatastoreHelper.POLL_MESSAGE_ID_STRIPPER;
-import static google.registry.testing.DatastoreHelper.getPollMessages;
 import static google.registry.testing.DatastoreHelper.stripBillingEventId;
 import static google.registry.xml.XmlTestUtils.assertXmlEquals;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -34,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.common.collect.ObjectArrays;
-import com.google.common.collect.Streams;
 import google.registry.config.RegistryConfig.ConfigModule.TmchCaMode;
 import google.registry.flows.EppTestComponent.FakesAndMocksModule;
 import google.registry.flows.picker.FlowPicker;
@@ -44,7 +40,6 @@ import google.registry.model.eppcommon.ProtocolDefinition;
 import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.ofy.Ofy;
-import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
 import google.registry.monitoring.whitebox.EppMetric;
@@ -185,6 +180,7 @@ public abstract class FlowTestCase<F extends Flow> {
       builder.put(
           GracePeriod.create(
               entry.getKey().getType(),
+              entry.getKey().getDomainRepoId(),
               entry.getKey().getExpirationTime(),
               entry.getKey().getClientId(),
               null),
@@ -213,24 +209,6 @@ public abstract class FlowTestCase<F extends Flow> {
         .isEqualTo(canonicalizeGracePeriods(expected));
   }
 
-  protected void assertPollMessages(String clientId, PollMessage... expected) {
-    assertPollMessagesHelper(getPollMessages(clientId), expected);
-  }
-
-  protected void assertPollMessages(PollMessage... expected) {
-    assertPollMessagesHelper(getPollMessages(), expected);
-  }
-
-  /** Assert that the list matches all the poll messages in the fake Datastore. */
-  private void assertPollMessagesHelper(
-      Iterable<PollMessage> pollMessages, PollMessage... expected) {
-    // Ordering is irrelevant but duplicates should be considered independently.
-    assertThat(
-            Streams.stream(pollMessages).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
-  }
-
   private EppOutput runFlowInternal(CommitMode commitMode, UserPrivileges userPrivileges)
       throws Exception {
     eppMetricBuilder = EppMetric.builderForRequest(clock);

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

@@ -282,13 +282,14 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
   }
 
   @Test
-  void testSuccess_allocationTokenPromotion() throws Exception {
+  void testSuccess_allocationTokenPromotion_singleYear() throws Exception {
     createTld("example");
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
             .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
+            .setDiscountYears(2)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                     .put(START_OF_TIME, TokenStatus.NOT_STARTED)
@@ -300,6 +301,69 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
     runFlowAssertResponse(loadFile("domain_check_allocationtoken_fee_response.xml"));
   }
 
+  @Test
+  void testSuccess_allocationTokenPromotion_multiYearAndPremiums() throws Exception {
+    createTld("example");
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.9)
+            .setDiscountYears(3)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().minusDays(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    setEppInput(
+        "domain_check_allocationtoken_promotion.xml", ImmutableMap.of("DOMAIN", "rich.example"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_check_allocationtoken_promotion_response.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "rich.example")
+                .put("COST_1YR", "10.00")
+                .put("COST_2YR", "20.00")
+                .put("COST_5YR", "230.00")
+                .put("FEE_CLASS", "<fee:class>premium</fee:class>")
+                .build()));
+  }
+
+  @Test
+  void testSuccess_allocationTokenPromotion_multiYear() throws Exception {
+    createTld("tld");
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("single.tld")
+            .setDiscountFraction(0.444)
+            .setDiscountYears(2)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().minusDays(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    setEppInput(
+        "domain_check_allocationtoken_promotion.xml", ImmutableMap.of("DOMAIN", "single.tld"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_check_allocationtoken_promotion_response.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "single.tld")
+                .put("COST_1YR", "7.23")
+                .put("COST_2YR", "14.46")
+                .put("COST_5YR", "53.46")
+                .put("FEE_CLASS", "")
+                .build()));
+  }
+
   @Test
   void testSuccess_promotionNotActive() throws Exception {
     createTld("example");

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

@@ -68,6 +68,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig;
 import google.registry.flows.EppException;
@@ -150,12 +151,12 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
-import google.registry.model.domain.token.AllocationToken.TokenType;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
@@ -324,7 +325,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     assertGracePeriods(
         domain.getGracePeriods(),
         ImmutableMap.of(
-            GracePeriod.create(GracePeriodStatus.ADD, billingTime, "TheRegistrar", null),
+            GracePeriod.create(
+                GracePeriodStatus.ADD, domain.getRepoId(), billingTime, "TheRegistrar", null),
             createBillingEvent));
     assertDnsTasksEnqueued(getUniqueIdFromCommand());
     assertEppResourceIndexEntityFor(domain);
@@ -435,7 +437,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testFailure_invalidAllocationToken() {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     EppException thrown = assertThrows(InvalidAllocationTokenException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
@@ -445,7 +449,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   void testFailure_reservedDomainCreate_allocationTokenIsForADifferentDomain() {
     // Try to register a reserved domain name with an allocation token valid for a different domain
     // name.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -464,7 +469,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   void testFailure_nonreservedDomainCreate_allocationTokenIsForADifferentDomain() {
     // Try to register a non-reserved domain name with an allocation token valid for a different
     // domain name.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -481,7 +488,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testFailure_alreadyRedemeedAllocationToken() {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -497,7 +506,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testSuccess_validAllocationToken_isRedeemed() throws Exception {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     AllocationToken token =
         persistResource(
@@ -508,12 +519,14 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry =
         ofy().load().type(HistoryEntry.class).ancestor(reloadResourceByForeignKey()).first().now();
     assertThat(ofy().load().entity(token).now().getRedemptionHistoryEntry())
-        .isEqualTo(Key.create(historyEntry));
+        .hasValue(Key.create(historyEntry));
   }
 
   @Test
   void testSuccess_validAllocationToken_multiUse() throws Exception {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     allocationToken =
         persistResource(
@@ -531,7 +544,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     runFlow();
     assertSuccessfulCreate("tld", ImmutableSet.of(), allocationToken);
     clock.advanceOneMilli();
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "otherexample.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "otherexample.tld", "YEARS", "2"));
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "otherexample.tld")));
   }
@@ -543,8 +558,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistContactsAndHosts("foo.tld");
     assertTransactionalFlow(true);
     String expectedResponseXml =
-        loadFile(
-            "domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.foo.tld"));
+        loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.foo.tld"));
     runFlowAssertResponse(CommitMode.LIVE, UserPrivileges.NORMAL, expectedResponseXml);
     assertSuccessfulCreate("foo.tld", ImmutableSet.of());
     assertNoLordn();
@@ -576,8 +590,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         "domain_create_registration_encoded_signed_mark.xml",
         ImmutableMap.of("DOMAIN", "test-validate.tld", "PHASE", "open"));
     persistContactsAndHosts();
-    EppException thrown =
-        assertThrows(SignedMarksOnlyDuringSunriseException.class, this::runFlow);
+    EppException thrown = assertThrows(SignedMarksOnlyDuringSunriseException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
@@ -1184,13 +1197,14 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         Registry.get("tld")
             .asBuilder()
             .setTldStateTransitions(
-                ImmutableSortedMap.of(
-                    START_OF_TIME, PREDELEGATION,
-                    DateTime.parse("1999-01-01T00:00:00Z"), QUIET_PERIOD,
-                    // The anchor tenant is created here, on 1999-04-03
-                    DateTime.parse("1999-07-01T00:00:00Z"), START_DATE_SUNRISE,
-                    DateTime.parse("2000-01-01T00:00:00Z"), GENERAL_AVAILABILITY))
+                new ImmutableSortedMap.Builder<DateTime, TldState>(Ordering.natural())
+                    .put(START_OF_TIME, PREDELEGATION)
+                    .put(DateTime.parse("1999-01-01T00:00:00Z"), QUIET_PERIOD)
+                    .put(DateTime.parse("1999-07-01T00:00:00Z"), START_DATE_SUNRISE)
+                    .put(DateTime.parse("2000-01-01T00:00:00Z"), GENERAL_AVAILABILITY)
+                    .build())
             .build());
+    // The anchor tenant is created during the quiet period, on 1999-04-03.
     setEppInput("domain_create_anchor_allocationtoken.xml");
     persistContactsAndHosts();
     runFlowAssertResponse(loadFile("domain_create_anchor_response.xml"));
@@ -1210,7 +1224,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                 .build());
     // Despite the domain being FULLY_BLOCKED, the non-superuser create succeeds the domain is also
     // RESERVED_FOR_SPECIFIC_USE and the correct allocation token is passed.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "resdom.tld")));
@@ -1233,7 +1248,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                 .setTokenType(SINGLE_USE)
                 .setDomainName("resdom.tld")
                 .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "resdom.tld")));
@@ -1247,7 +1263,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         ofy().load().key(Key.create(AllocationToken.class, token)).now();
     assertThat(reloadedToken.isRedeemed()).isTrue();
     assertThat(reloadedToken.getRedemptionHistoryEntry())
-        .isEqualTo(Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0)));
+        .hasValue(Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0)));
   }
 
   private void assertAllocationTokenWasNotRedeemed(String token) {
@@ -1264,7 +1280,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1274,7 +1290,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .build())
             .build());
     clock.advanceOneMilli();
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.tld")));
     BillingEvent.OneTime billingEvent =
@@ -1284,14 +1302,131 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   }
 
   @Test
-  void testSuccess_promotionDoesNotApplyToPremiumPrice() {
-    // At the moment, discounts cannot apply to premium domains
+  void testSuccess_allocationToken_multiYearDiscount_maxesAtTokenDiscountYears() throws Exception {
+    // 2yrs @ $13 + 3yrs @ $13 * (1 - 0.73) = $36.53
+    runTest_allocationToken_multiYearDiscount(false, 0.73, 3, Money.of(USD, 36.53));
+  }
+
+  @Test
+  void testSuccess_allocationToken_multiYearDiscount_maxesAtNumRegistrationYears()
+      throws Exception {
+    // 5yrs @ $13 * (1 - 0.276) = $47.06
+    runTest_allocationToken_multiYearDiscount(false, 0.276, 10, Money.of(USD, 47.06));
+  }
+
+  void runTest_allocationToken_multiYearDiscount(
+      boolean discountPremiums, double discountFraction, int discountYears, Money expectedPrice)
+      throws Exception {
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("example.tld")
+            .setDiscountFraction(discountFraction)
+            .setDiscountYears(discountYears)
+            .setDiscountPremiums(discountPremiums)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "5"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_wildcard.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "example.tld")
+                .put("CRDATE", "1999-04-03T22:00:00.0Z")
+                .put("EXDATE", "2004-04-03T22:00:00.0Z")
+                .build()));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
+    assertThat(billingEvent.getCost()).isEqualTo(expectedPrice);
+  }
+
+  @Test
+  void testSuccess_allocationToken_multiYearDiscount_worksForPremiums() throws Exception {
     createTld("example");
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.98)
+            .setDiscountYears(2)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "3", "FEE", "104.00"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2002-04-03T22:00:00.0Z", "FEE", "104.00")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("rich.example");
+    // 1yr @ $100 + 2yrs @ $100 * (1 - 0.98) = $104
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 104.00));
+  }
+
+  @Test
+  void testSuccess_allocationToken_singleYearDiscount_worksForPremiums() throws Exception {
+    createTld("example");
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.95555)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "3", "FEE", "204.44"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2002-04-03T22:00:00.0Z", "FEE", "204.44")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("rich.example");
+    // 2yrs @ $100 + 1yr @ $100 * (1 - 0.95555) = $204.44
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 204.44));
+  }
+
+  @Test
+  void testSuccess_promotionDoesNotApplyToPremiumPrice() {
+    // Discounts only apply to premium domains if the token is explicitly configured to allow it.
+    createTld("example");
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1301,7 +1436,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .build())
             .build());
     clock.advanceOneMilli();
-    setEppInput("domain_create_premium_allocationtoken.xml");
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "2", "FEE", "193.50"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenInvalidForPremiumNameException.class, this::runFlow))
         .marshalsToXml();
@@ -1313,7 +1450,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1322,7 +1459,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(60), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotInPromotionException.class, this::runFlow))
         .marshalsToXml();
@@ -1334,7 +1473,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setAllowedTlds(ImmutableSet.of("example"))
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
@@ -1344,7 +1483,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotValidForTldException.class, this::runFlow))
         .marshalsToXml();
@@ -1356,7 +1497,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setAllowedClientIds(ImmutableSet.of("someClientId"))
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
@@ -1366,7 +1507,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotValidForRegistrarException.class, this::runFlow))
         .marshalsToXml();
@@ -1562,7 +1705,11 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     // Modify the Registrar to block premium names.
     persistResource(loadRegistrar("TheRegistrar").asBuilder().setBlockPremiumNames(true).build());
     runFlowAssertResponse(
-        CommitMode.LIVE, SUPERUSER, loadFile("domain_create_response_premium.xml"));
+        CommitMode.LIVE,
+        SUPERUSER,
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2001-04-03T22:00:00.0Z", "FEE", "200.00")));
     assertSuccessfulCreate("example", ImmutableSet.of());
   }
 

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

@@ -36,6 +36,7 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.getOnlyPollMessage;
@@ -195,6 +196,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
                     ImmutableSet.of(
                         GracePeriod.createForRecurring(
                             GracePeriodStatus.AUTO_RENEW,
+                            domain.getRepoId(),
                             A_MONTH_AGO.plusDays(45),
                             "TheRegistrar",
                             autorenewBillingEvent.createVKey())))
@@ -290,7 +292,8 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
   void testDryRun() throws Exception {
     setUpSuccessfulTest();
     setUpGracePeriods(
-        GracePeriod.create(GracePeriodStatus.ADD, TIME_BEFORE_FLOW.plusDays(1), "foo", null));
+        GracePeriod.create(
+            GracePeriodStatus.ADD, domain.getRepoId(), TIME_BEFORE_FLOW.plusDays(1), "foo", null));
     dryRunFlowAssertResponse(loadFile("generic_success_response.xml"));
   }
 
@@ -314,7 +317,8 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     setUpSuccessfulTest();
     BillingEvent.OneTime graceBillingEvent =
         persistResource(createBillingEvent(Reason.CREATE, Money.of(USD, 123)));
-    setUpGracePeriods(GracePeriod.forBillingEvent(gracePeriodStatus, graceBillingEvent));
+    setUpGracePeriods(
+        GracePeriod.forBillingEvent(gracePeriodStatus, domain.getRepoId(), graceBillingEvent));
     // We should see exactly one poll message, which is for the autorenew 1 month in the future.
     assertPollMessages(createAutorenewPollMessage("TheRegistrar").build());
     clock.advanceOneMilli();
@@ -388,9 +392,14 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     BillingEvent.OneTime renewBillingEvent =
         persistResource(createBillingEvent(Reason.RENEW, Money.of(USD, 456)));
     setUpGracePeriods(
-        GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, renewBillingEvent),
+        GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, domain.getRepoId(), renewBillingEvent),
         // This grace period has no associated billing event, so it won't cause a cancellation.
-        GracePeriod.create(GracePeriodStatus.TRANSFER, TIME_BEFORE_FLOW.plusDays(1), "foo", null));
+        GracePeriod.create(
+            GracePeriodStatus.TRANSFER,
+            domain.getRepoId(),
+            TIME_BEFORE_FLOW.plusDays(1),
+            "foo",
+            null));
     // We should see exactly one poll message, which is for the autorenew 1 month in the future.
     assertPollMessages(createAutorenewPollMessage("TheRegistrar").build());
     DateTime expectedExpirationTime = domain.getRegistrationExpirationTime().minusYears(2);
@@ -424,6 +433,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         .containsExactly(
             GracePeriod.create(
                 GracePeriodStatus.REDEMPTION,
+                domain.getRepoId(),
                 clock.nowUtc().plus(Registry.get("tld").getRedemptionGracePeriodLength()),
                 "TheRegistrar",
                 null));
@@ -624,6 +634,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         .containsExactly(
             GracePeriod.create(
                 GracePeriodStatus.REDEMPTION,
+                domain.getRepoId(),
                 clock.nowUtc().plus(Registry.get("tld").getRedemptionGracePeriodLength()),
                 "TheRegistrar",
                 null));
@@ -696,7 +707,8 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     BillingEvent.OneTime graceBillingEvent =
         persistResource(createBillingEvent(Reason.CREATE, Money.of(USD, 123)));
     // Use a grace period so that the delete is immediate, simplifying the assertions below.
-    setUpGracePeriods(GracePeriod.forBillingEvent(GracePeriodStatus.ADD, graceBillingEvent));
+    setUpGracePeriods(
+        GracePeriod.forBillingEvent(GracePeriodStatus.ADD, domain.getRepoId(), graceBillingEvent));
     // Add a nameserver.
     HostResource host = persistResource(newHostResource("ns1.example.tld"));
     persistResource(
@@ -1004,7 +1016,11 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     setUpSuccessfulTest();
     setUpGracePeriods(
         GracePeriod.create(
-            GracePeriodStatus.ADD, TIME_BEFORE_FLOW.plusDays(1), "TheRegistrar", null));
+            GracePeriodStatus.ADD,
+            domain.getRepoId(),
+            TIME_BEFORE_FLOW.plusDays(1),
+            "TheRegistrar",
+            null));
     setUpGracePeriodDurations();
     clock.advanceOneMilli();
     earlierHistoryEntry =
@@ -1027,7 +1043,11 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     setUpSuccessfulTest();
     setUpGracePeriods(
         GracePeriod.create(
-            GracePeriodStatus.ADD, TIME_BEFORE_FLOW.plusDays(1), "TheRegistrar", null));
+            GracePeriodStatus.ADD,
+            domain.getRepoId(),
+            TIME_BEFORE_FLOW.plusDays(1),
+            "TheRegistrar",
+            null));
     setUpGracePeriodDurations();
     clock.advanceOneMilli();
     earlierHistoryEntry =
@@ -1081,6 +1101,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         .containsExactly(
             GracePeriod.create(
                 GracePeriodStatus.REDEMPTION,
+                domain.getRepoId(),
                 clock.nowUtc().plus(standardDays(15)),
                 "TheRegistrar",
                 null));
@@ -1127,6 +1148,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         .containsExactly(
             GracePeriod.create(
                 GracePeriodStatus.REDEMPTION,
+                domain.getRepoId(),
                 clock.nowUtc().plus(standardDays(15)),
                 "TheRegistrar",
                 null));

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

@@ -308,7 +308,8 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
         domain
             .asBuilder()
             .addGracePeriod(
-                GracePeriod.create(gracePeriodStatus, clock.nowUtc().plusDays(1), "foo", null))
+                GracePeriod.create(
+                    gracePeriodStatus, domain.getRepoId(), clock.nowUtc().plusDays(1), "foo", null))
             .setCreationClientId("NewRegistrar")
             .setCreationTimeForTest(DateTime.parse("2003-11-26T22:00:00.0Z"))
             .setRegistrationExpirationTime(DateTime.parse("2005-11-26T22:00:00.0Z"))
@@ -337,7 +338,11 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.createForRecurring(
-                    GracePeriodStatus.AUTO_RENEW, clock.nowUtc().plusDays(1), "foo", recurringVKey))
+                    GracePeriodStatus.AUTO_RENEW,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    recurringVKey))
             .build());
     doSuccessfulTest("domain_info_response_autorenewperiod.xml", false);
   }
@@ -352,7 +357,11 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.REDEMPTION, clock.nowUtc().plusDays(1), "foo", null))
+                    GracePeriodStatus.REDEMPTION,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
             .build());
     doSuccessfulTest("domain_info_response_redemptionperiod.xml", false);
@@ -367,7 +376,11 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.RENEW, clock.nowUtc().plusDays(1), "foo", null))
+                    GracePeriodStatus.RENEW,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .build());
     doSuccessfulTest("domain_info_response_renewperiod.xml", false);
   }
@@ -381,10 +394,18 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.RENEW, clock.nowUtc().plusDays(1), "foo", null))
+                    GracePeriodStatus.RENEW,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.RENEW, clock.nowUtc().plusDays(2), "foo", null))
+                    GracePeriodStatus.RENEW,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(2),
+                    "foo",
+                    null))
             .build());
     doSuccessfulTest("domain_info_response_renewperiod.xml", false);
   }
@@ -398,7 +419,11 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.TRANSFER, clock.nowUtc().plusDays(1), "foo", null))
+                    GracePeriodStatus.TRANSFER,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .build());
     doSuccessfulTest("domain_info_response_transferperiod.xml", false);
   }
@@ -421,10 +446,19 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
         domain
             .asBuilder()
             .addGracePeriod(
-                GracePeriod.create(GracePeriodStatus.ADD, clock.nowUtc().plusDays(1), "foo", null))
+                GracePeriod.create(
+                    GracePeriodStatus.ADD,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.RENEW, clock.nowUtc().plusDays(2), "foo", null))
+                    GracePeriodStatus.RENEW,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(2),
+                    "foo",
+                    null))
             .build());
     doSuccessfulTest("domain_info_response_stackedaddrenewperiod.xml", false);
   }
@@ -437,7 +471,12 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
         domain
             .asBuilder()
             .addGracePeriod(
-                GracePeriod.create(GracePeriodStatus.ADD, clock.nowUtc().plusDays(1), "foo", null))
+                GracePeriod.create(
+                    GracePeriodStatus.ADD,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .setDsData(
                 ImmutableSet.of(
                     DelegationSignerData.create(

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.flows.domain.DomainTransferFlowTestCase.persistWithPendingTransfer;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.loadRegistrar;
@@ -236,6 +237,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
         ImmutableMap.of(
             GracePeriod.create(
                 GracePeriodStatus.RENEW,
+                domain.getRepoId(),
                 clock.nowUtc().plus(Registry.get("tld").getRenewGracePeriodLength()),
                 renewalClientId,
                 null),

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

@@ -17,6 +17,7 @@ package google.registry.flows.domain;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.getPollMessages;
@@ -112,7 +113,11 @@ class DomainRestoreRequestFlowTest
             .setDeletionTime(clock.nowUtc().plusDays(35))
             .addGracePeriod(
                 GracePeriod.create(
-                    GracePeriodStatus.REDEMPTION, clock.nowUtc().plusDays(1), "foo", null))
+                    GracePeriodStatus.REDEMPTION,
+                    domain.getRepoId(),
+                    clock.nowUtc().plusDays(1),
+                    "foo",
+                    null))
             .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
             .setDeletePollMessage(
                 persistResource(

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

@@ -288,6 +288,7 @@ class DomainTransferApproveFlowTest
         ImmutableMap.of(
             GracePeriod.create(
                 GracePeriodStatus.TRANSFER,
+                domain.getRepoId(),
                 clock.nowUtc().plus(registry.getTransferGracePeriodLength()),
                 "NewRegistrar",
                 null),
@@ -649,6 +650,7 @@ class DomainTransferApproveFlowTest
                 .addGracePeriod(
                     GracePeriod.createForRecurring(
                         GracePeriodStatus.AUTO_RENEW,
+                        domain.getRepoId(),
                         autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()),
                         "TheRegistrar",
                         existingAutorenewEvent))

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

@@ -21,6 +21,7 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createPollMessageForImplicitTransfer;
 import static google.registry.testing.DatastoreHelper.deleteResource;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;

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

@@ -332,6 +332,7 @@ class DomainTransferRequestFlowTest
           ImmutableMap.of(
               GracePeriod.create(
                   GracePeriodStatus.TRANSFER,
+                  domain.getRepoId(),
                   implicitTransferTime.plus(registry.getTransferGracePeriodLength()),
                   "NewRegistrar",
                   null),
@@ -962,6 +963,7 @@ class DomainTransferRequestFlowTest
                 .addGracePeriod(
                     GracePeriod.createForRecurring(
                         GracePeriodStatus.AUTO_RENEW,
+                        domain.getRepoId(),
                         autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()),
                         "TheRegistrar",
                         existingAutorenewEvent))
@@ -1088,6 +1090,7 @@ class DomainTransferRequestFlowTest
                 .addGracePeriod(
                     GracePeriod.createForRecurring(
                         GracePeriodStatus.AUTO_RENEW,
+                        domain.getRepoId(),
                         autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()),
                         "TheRegistrar",
                         domain.getAutorenewBillingEvent()))
@@ -1117,6 +1120,7 @@ class DomainTransferRequestFlowTest
                 .addGracePeriod(
                     GracePeriod.createForRecurring(
                         GracePeriodStatus.AUTO_RENEW,
+                        domain.getRepoId(),
                         autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()),
                         "TheRegistrar",
                         existingAutorenewEvent))

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

@@ -385,10 +385,11 @@ public class BillingEventTest extends EntityTestCase {
 
   @Test
   void testSuccess_cancellation_forGracePeriod_withOneTime() {
-    BillingEvent.Cancellation newCancellation = BillingEvent.Cancellation.forGracePeriod(
-        GracePeriod.forBillingEvent(GracePeriodStatus.ADD, oneTime),
-        historyEntry2,
-        "foo.tld");
+    BillingEvent.Cancellation newCancellation =
+        BillingEvent.Cancellation.forGracePeriod(
+            GracePeriod.forBillingEvent(GracePeriodStatus.ADD, domain.getRepoId(), oneTime),
+            historyEntry2,
+            "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);
@@ -400,6 +401,7 @@ public class BillingEventTest extends EntityTestCase {
         BillingEvent.Cancellation.forGracePeriod(
             GracePeriod.createForRecurring(
                 GracePeriodStatus.AUTO_RENEW,
+                domain.getRepoId(),
                 now.plusYears(1).plusDays(45),
                 "a registrar",
                 recurring.createVKey()),
@@ -418,7 +420,10 @@ public class BillingEventTest extends EntityTestCase {
             () ->
                 BillingEvent.Cancellation.forGracePeriod(
                     GracePeriod.createWithoutBillingEvent(
-                        GracePeriodStatus.REDEMPTION, now.plusDays(1), "a registrar"),
+                        GracePeriodStatus.REDEMPTION,
+                        domain.getRepoId(),
+                        now.plusDays(1),
+                        "a registrar"),
                     historyEntry,
                     "foo.tld"));
     assertThat(thrown).hasMessageThat().contains("grace period without billing event");

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

@@ -18,6 +18,7 @@ import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableO
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.SqlHelper.assertThrowForeignKeyViolation;
 import static google.registry.testing.SqlHelper.saveRegistrar;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static org.joda.time.DateTimeZone.UTC;
 
@@ -25,6 +26,7 @@ import com.google.common.collect.ImmutableSet;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DesignatedContact.Type;
 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;
@@ -100,6 +102,8 @@ public class DomainBaseSqlTest {
             .setLaunchNotice(
                 LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
             .setSmdId("smdid")
+            .addGracePeriod(
+                GracePeriod.create(GracePeriodStatus.ADD, "4-COM", END_OF_TIME, "registrar1", null))
             .build();
 
     host =

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

@@ -154,6 +154,7 @@ public class DomainBaseTest extends EntityTestCase {
                     .addGracePeriod(
                         GracePeriod.create(
                             GracePeriodStatus.ADD,
+                            "4-COM",
                             fakeClock.nowUtc().plusDays(1),
                             "registrar",
                             null))
@@ -371,7 +372,11 @@ public class DomainBaseTest extends EntityTestCase {
                 // Okay for billing event to be null since the point of this grace period is just
                 // to check that the transfer will clear all existing grace periods.
                 GracePeriod.create(
-                    GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(100), "foo", null))
+                    GracePeriodStatus.ADD,
+                    domain.getRepoId(),
+                    fakeClock.nowUtc().plusDays(100),
+                    "foo",
+                    null))
             .build();
     DomainBase afterTransfer = domain.cloneProjectedAtTime(fakeClock.nowUtc().plusDays(1));
     DateTime newExpirationTime = oldExpirationTime.plusYears(1);
@@ -382,6 +387,7 @@ public class DomainBaseTest extends EntityTestCase {
         .containsExactly(
             GracePeriod.create(
                 GracePeriodStatus.TRANSFER,
+                domain.getRepoId(),
                 fakeClock
                     .nowUtc()
                     .plusDays(1)
@@ -499,9 +505,24 @@ public class DomainBaseTest extends EntityTestCase {
   void testStackedGracePeriods() {
     ImmutableList<GracePeriod> gracePeriods =
         ImmutableList.of(
-            GracePeriod.create(GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(3), "foo", null),
-            GracePeriod.create(GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(2), "bar", null),
-            GracePeriod.create(GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(1), "baz", null));
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(3),
+                "foo",
+                null),
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(2),
+                "bar",
+                null),
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(1),
+                "baz",
+                null));
     domain = domain.asBuilder().setGracePeriods(ImmutableSet.copyOf(gracePeriods)).build();
     for (int i = 1; i < 3; ++i) {
       assertThat(domain.cloneProjectedAtTime(fakeClock.nowUtc().plusDays(i)).getGracePeriods())
@@ -513,14 +534,32 @@ public class DomainBaseTest extends EntityTestCase {
   void testGracePeriodsByType() {
     ImmutableSet<GracePeriod> addGracePeriods =
         ImmutableSet.of(
-            GracePeriod.create(GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(3), "foo", null),
-            GracePeriod.create(GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(1), "baz", null));
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(3),
+                "foo",
+                null),
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(1),
+                "baz",
+                null));
     ImmutableSet<GracePeriod> renewGracePeriods =
         ImmutableSet.of(
             GracePeriod.create(
-                GracePeriodStatus.RENEW, fakeClock.nowUtc().plusDays(3), "foo", null),
+                GracePeriodStatus.RENEW,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(3),
+                "foo",
+                null),
             GracePeriod.create(
-                GracePeriodStatus.RENEW, fakeClock.nowUtc().plusDays(1), "baz", null));
+                GracePeriodStatus.RENEW,
+                domain.getRepoId(),
+                fakeClock.nowUtc().plusDays(1),
+                "baz",
+                null));
     domain =
         domain
             .asBuilder()
@@ -589,6 +628,7 @@ public class DomainBaseTest extends EntityTestCase {
         .containsExactly(
             GracePeriod.createForRecurring(
                 GracePeriodStatus.AUTO_RENEW,
+                domain.getRepoId(),
                 oldExpirationTime
                     .plusYears(2)
                     .plus(Registry.get("com").getAutoRenewGracePeriodLength()),
@@ -757,6 +797,7 @@ public class DomainBaseTest extends EntityTestCase {
                     ImmutableSet.of(
                         GracePeriod.createForRecurring(
                             GracePeriodStatus.AUTO_RENEW,
+                            domain.getRepoId(),
                             now.plusDays(1),
                             "NewRegistrar",
                             recurringBillKey)))

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

@@ -63,8 +63,9 @@ public class GracePeriodTest {
 
   @Test
   void testSuccess_forBillingEvent() {
-    GracePeriod gracePeriod = GracePeriod.forBillingEvent(GracePeriodStatus.ADD, onetime);
+    GracePeriod gracePeriod = GracePeriod.forBillingEvent(GracePeriodStatus.ADD, "1-TEST", onetime);
     assertThat(gracePeriod.getType()).isEqualTo(GracePeriodStatus.ADD);
+    assertThat(gracePeriod.getDomainRepoId()).isEqualTo("1-TEST");
     assertThat(gracePeriod.getOneTimeBillingEvent()).isEqualTo(onetime.createVKey());
     assertThat(gracePeriod.getRecurringBillingEvent()).isNull();
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
@@ -74,9 +75,11 @@ public class GracePeriodTest {
 
   @Test
   void testSuccess_createWithoutBillingEvent() {
-    GracePeriod gracePeriod = GracePeriod.createWithoutBillingEvent(
-        GracePeriodStatus.REDEMPTION, now, "TheRegistrar");
+    GracePeriod gracePeriod =
+        GracePeriod.createWithoutBillingEvent(
+            GracePeriodStatus.REDEMPTION, "1-TEST", now, "TheRegistrar");
     assertThat(gracePeriod.getType()).isEqualTo(GracePeriodStatus.REDEMPTION);
+    assertThat(gracePeriod.getDomainRepoId()).isEqualTo("1-TEST");
     assertThat(gracePeriod.getOneTimeBillingEvent()).isNull();
     assertThat(gracePeriod.getRecurringBillingEvent()).isNull();
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
@@ -89,7 +92,7 @@ public class GracePeriodTest {
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () -> GracePeriod.forBillingEvent(GracePeriodStatus.AUTO_RENEW, onetime));
+            () -> GracePeriod.forBillingEvent(GracePeriodStatus.AUTO_RENEW, "1-TEST", onetime));
     assertThat(thrown).hasMessageThat().contains("autorenew");
   }
 
@@ -101,6 +104,7 @@ public class GracePeriodTest {
             () ->
                 GracePeriod.createForRecurring(
                     GracePeriodStatus.RENEW,
+                    "1-TEST",
                     now.plusDays(1),
                     "TheRegistrar",
                     VKey.create(Recurring.class, 12345)));

core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java

@@ -44,7 +44,7 @@ import org.junit.jupiter.api.Test;
 class AllocationTokenTest extends EntityTestCase {
 
   @BeforeEach
-  void setup() {
+  void beforeEach() {
     createTld("foo");
   }
 
@@ -59,6 +59,8 @@ class AllocationTokenTest extends EntityTestCase {
                 .setAllowedTlds(ImmutableSet.of("dev", "app"))
                 .setAllowedClientIds(ImmutableSet.of("TheRegistrar, NewRegistrar"))
                 .setDiscountFraction(0.5)
+                .setDiscountPremiums(true)
+                .setDiscountYears(3)
                 .setTokenStatusTransitions(
                     ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                         .put(START_OF_TIME, NOT_STARTED)
@@ -155,7 +157,7 @@ class AllocationTokenTest extends EntityTestCase {
   }
 
   @Test
-  void testBuild_invalidTLD() {
+  void testBuild_invalidTld() {
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
@@ -268,6 +270,50 @@ class AllocationTokenTest extends EntityTestCase {
     assertTerminal(CANCELLED);
   }
 
+  @Test
+  void testSetDiscountFractionTooHigh() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountFraction(1.1));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
+  }
+
+  @Test
+  void testSetDiscountFractionTooLow() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountFraction(-.0001));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
+  }
+
+  @Test
+  void testSetDiscountYearsTooHigh() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountYears(11));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years must be between 1 and 10 inclusive");
+  }
+
+  @Test
+  void testSetDiscountYearsTooLow() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountYears(0));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years must be between 1 and 10 inclusive");
+  }
+
   @Test
   void testBuild_noTokenType() {
     IllegalArgumentException thrown =
@@ -295,6 +341,38 @@ class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token must not be blank");
   }
 
+  @Test
+  void testBuild_discountPremiumsRequiresDiscountFraction() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new AllocationToken.Builder()
+                    .setToken("abc")
+                    .setTokenType(SINGLE_USE)
+                    .setDiscountPremiums(true)
+                    .build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount premiums can only be specified along with a discount fraction");
+  }
+
+  @Test
+  void testBuild_discountYearsRequiresDiscountFraction() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new AllocationToken.Builder()
+                    .setToken("abc")
+                    .setTokenType(SINGLE_USE)
+                    .setDiscountYears(2)
+                    .build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years can only be specified along with a discount fraction");
+  }
+
   private void assertBadInitialTransition(TokenStatus status) {
     assertBadTransition(
         ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()

core/src/test/java/google/registry/model/history/ContactHistoryTest.java

@@ -1,4 +1,4 @@
-// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+// 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.
@@ -15,7 +15,9 @@
 package google.registry.model.history;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatastoreHelper.newContactResourceWithRoid;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
@@ -24,7 +26,6 @@ import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.model.transfer.ContactTransferData;
 import google.registry.persistence.VKey;
 import org.junit.jupiter.api.Test;
 
@@ -37,26 +38,18 @@ public class ContactHistoryTest extends EntityTestCase {
 
   @Test
   void testPersistence() {
-    saveRegistrar("registrar1");
-
-    ContactResource contact =
-        new ContactResource.Builder()
-            .setRepoId("contact1")
-            .setContactId("contactId")
-            .setCreationClientId("registrar1")
-            .setPersistedCurrentSponsorClientId("registrar1")
-            .setTransferData(new ContactTransferData.Builder().build())
-            .build();
+    saveRegistrar("TheRegistrar");
 
+    ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
     jpaTm().transact(() -> jpaTm().saveNew(contact));
-    VKey<ContactResource> contactVKey = VKey.createSql(ContactResource.class, "contact1");
+    VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().load(contactVKey));
     ContactHistory contactHistory =
         new ContactHistory.Builder()
             .setType(HistoryEntry.Type.HOST_CREATE)
             .setXmlBytes("<xml></xml>".getBytes(UTF_8))
             .setModificationTime(fakeClock.nowUtc())
-            .setClientId("registrar1")
+            .setClientId("TheRegistrar")
             .setTrid(Trid.create("ABC-123", "server-trid"))
             .setBySuperuser(false)
             .setReason("reason")
@@ -70,18 +63,17 @@ public class ContactHistoryTest extends EntityTestCase {
             () -> {
               ContactHistory fromDatabase = jpaTm().load(VKey.createSql(ContactHistory.class, 1L));
               assertContactHistoriesEqual(fromDatabase, contactHistory);
+              assertThat(fromDatabase.getContactRepoId().getSqlKey())
+                  .isEqualTo(contactHistory.getContactRepoId().getSqlKey());
             });
   }
 
-  private void assertContactHistoriesEqual(ContactHistory one, ContactHistory two) {
-    // enough of the fields get changed during serialization that we can't depend on .equals()
-    assertThat(one.getClientId()).isEqualTo(two.getClientId());
-    assertThat(one.getContactRepoId()).isEqualTo(two.getContactRepoId());
-    assertThat(one.getBySuperuser()).isEqualTo(two.getBySuperuser());
-    assertThat(one.getRequestedByRegistrar()).isEqualTo(two.getRequestedByRegistrar());
-    assertThat(one.getReason()).isEqualTo(two.getReason());
-    assertThat(one.getTrid()).isEqualTo(two.getTrid());
-    assertThat(one.getType()).isEqualTo(two.getType());
-    assertThat(one.getContactBase().getContactId()).isEqualTo(two.getContactBase().getContactId());
+  static void assertContactHistoriesEqual(ContactHistory one, ContactHistory two) {
+    assertAboutImmutableObjects()
+        .that(one)
+        .isEqualExceptFields(two, "contactBase", "contactRepoId", "parent");
+    assertAboutImmutableObjects()
+        .that(one.getContactBase())
+        .isEqualExceptFields(two.getContactBase(), "repoId");
   }
 }

core/src/test/java/google/registry/model/history/DomainHistoryTest.java

@@ -0,0 +1,90 @@
+// 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.model.history;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatastoreHelper.newContactResourceWithRoid;
+import static google.registry.testing.DatastoreHelper.newDomainBase;
+import static google.registry.testing.DatastoreHelper.newHostResourceWithRoid;
+import static google.registry.testing.SqlHelper.saveRegistrar;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import google.registry.model.EntityTestCase;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostResource;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import org.junit.jupiter.api.Test;
+
+/** Tests for {@link DomainHistory}. */
+public class DomainHistoryTest extends EntityTestCase {
+
+  public DomainHistoryTest() {
+    super(JpaEntityCoverageCheck.ENABLED);
+  }
+
+  @Test
+  public void testPersistence() {
+    saveRegistrar("TheRegistrar");
+
+    HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
+    jpaTm().transact(() -> jpaTm().saveNew(host));
+    ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
+    jpaTm().transact(() -> jpaTm().saveNew(contact));
+
+    DomainBase domain =
+        newDomainBase("example.tld", "domainRepoId", contact)
+            .asBuilder()
+            .setNameservers(host.createVKey())
+            .build();
+    jpaTm().transact(() -> jpaTm().saveNew(domain));
+
+    DomainHistory domainHistory =
+        new DomainHistory.Builder()
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("TheRegistrar")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(true)
+            .setDomainContent(domain)
+            .setDomainRepoId(domain.createVKey())
+            .build();
+    jpaTm().transact(() -> jpaTm().saveNew(domainHistory));
+
+    jpaTm()
+        .transact(
+            () -> {
+              DomainHistory fromDatabase =
+                  jpaTm().load(VKey.createSql(DomainHistory.class, domainHistory.getId()));
+              assertDomainHistoriesEqual(fromDatabase, domainHistory);
+              assertThat(fromDatabase.getDomainRepoId().getSqlKey())
+                  .isEqualTo(domainHistory.getDomainRepoId().getSqlKey());
+            });
+  }
+
+  static void assertDomainHistoriesEqual(DomainHistory one, DomainHistory two) {
+    assertAboutImmutableObjects()
+        .that(one)
+        .isEqualExceptFields(two, "domainContent", "domainRepoId", "parent");
+  }
+}

core/src/test/java/google/registry/model/history/HostHistoryTest.java

@@ -14,12 +14,13 @@
 
 package google.registry.model.history;
 
+import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatastoreHelper.newHostResourceWithRoid;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
-import com.google.common.collect.ImmutableSet;
 import google.registry.model.EntityTestCase;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.host.HostHistory;
@@ -37,16 +38,9 @@ public class HostHistoryTest extends EntityTestCase {
 
   @Test
   void testPersistence() {
-    saveRegistrar("registrar1");
+    saveRegistrar("TheRegistrar");
 
-    HostResource host =
-        new HostResource.Builder()
-            .setRepoId("host1")
-            .setHostName("ns1.example.com")
-            .setCreationClientId("TheRegistrar")
-            .setPersistedCurrentSponsorClientId("TheRegistrar")
-            .setInetAddresses(ImmutableSet.of())
-            .build();
+    HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
     jpaTm().transact(() -> jpaTm().saveNew(host));
     VKey<HostResource> hostVKey = VKey.createSql(HostResource.class, "host1");
     HostResource hostFromDb = jpaTm().transact(() -> jpaTm().load(hostVKey));
@@ -55,7 +49,7 @@ public class HostHistoryTest extends EntityTestCase {
             .setType(HistoryEntry.Type.HOST_CREATE)
             .setXmlBytes("<xml></xml>".getBytes(UTF_8))
             .setModificationTime(fakeClock.nowUtc())
-            .setClientId("registrar1")
+            .setClientId("TheRegistrar")
             .setTrid(Trid.create("ABC-123", "server-trid"))
             .setBySuperuser(false)
             .setReason("reason")
@@ -70,6 +64,8 @@ public class HostHistoryTest extends EntityTestCase {
               HostHistory fromDatabase =
                   jpaTm().load(VKey.createSql(HostHistory.class, hostHistory.getId()));
               assertHostHistoriesEqual(fromDatabase, hostHistory);
+              assertThat(fromDatabase.getHostRepoId().getSqlKey())
+                  .isEqualTo(hostHistory.getHostRepoId().getSqlKey());
             });
   }
 

core/src/test/java/google/registry/model/reporting/Spec11ThreatMatchTest.java

@@ -33,6 +33,7 @@ import google.registry.persistence.VKey;
 import org.joda.time.LocalDate;
 import org.joda.time.format.ISODateTimeFormat;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link Spec11ThreatMatch}. */
@@ -114,11 +115,14 @@ public class Spec11ThreatMatchTest extends EntityTestCase {
 
     VKey<Spec11ThreatMatch> threatVKey = VKey.createSql(Spec11ThreatMatch.class, threat.getId());
     Spec11ThreatMatch persistedThreat = jpaTm().transact(() -> jpaTm().load(threatVKey));
+
+    // Threat object saved for the first time doesn't have an ID; it is generated by SQL
     threat.id = persistedThreat.id;
     assertThat(threat).isEqualTo(persistedThreat);
   }
 
   @Test
+  @Disabled("We can't rely on foreign keys until we've migrated to SQL")
   void testThreatForeignKeyConstraints() {
     assertThrowForeignKeyViolation(
         () -> {

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

@@ -18,6 +18,11 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
@@ -26,12 +31,16 @@ import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension
 import google.registry.testing.FakeClock;
 import java.io.Serializable;
 import java.math.BigInteger;
+import java.sql.SQLException;
 import java.util.NoSuchElementException;
+import java.util.function.Supplier;
 import javax.persistence.Entity;
 import javax.persistence.EntityManager;
 import javax.persistence.Id;
 import javax.persistence.IdClass;
+import javax.persistence.OptimisticLockException;
 import javax.persistence.RollbackException;
+import org.hibernate.exception.JDBCConnectionException;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -130,6 +139,117 @@ class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey))).isEqualTo(theEntity);
   }
 
+  @Test
+  public void transact_retriesOptimisticLockExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(OptimisticLockException.class).when(spyJpaTm).delete(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        OptimisticLockException.class,
+        () -> spyJpaTm.transact(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(3)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(OptimisticLockException.class, () -> spyJpaTm.transact(supplier));
+    verify(spyJpaTm, times(6)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transactNoRetry_doesNotRetryOptimisticLockException() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(OptimisticLockException.class).when(spyJpaTm).delete(any(VKey.class));
+    spyJpaTm.transactNoRetry(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        OptimisticLockException.class,
+        () -> spyJpaTm.transactNoRetry(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(1)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(OptimisticLockException.class, () -> spyJpaTm.transactNoRetry(supplier));
+    verify(spyJpaTm, times(2)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transact_retriesNestedOptimisticLockExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(new RuntimeException().initCause(new OptimisticLockException()))
+        .when(spyJpaTm)
+        .delete(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class, () -> spyJpaTm.transact(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(3)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(RuntimeException.class, () -> spyJpaTm.transact(supplier));
+    verify(spyJpaTm, times(6)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transactNewReadOnly_retriesJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(JDBCConnectionException.class).when(spyJpaTm).load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        JDBCConnectionException.class,
+        () -> spyJpaTm.transactNewReadOnly(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.load(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(JDBCConnectionException.class, () -> spyJpaTm.transactNewReadOnly(supplier));
+    verify(spyJpaTm, times(6)).load(theEntityKey);
+  }
+
+  @Test
+  public void transactNewReadOnly_retriesNestedJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(
+            new RuntimeException()
+                .initCause(new JDBCConnectionException("connection exception", new SQLException())))
+        .when(spyJpaTm)
+        .load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class,
+        () -> spyJpaTm.transactNewReadOnly(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.load(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(RuntimeException.class, () -> spyJpaTm.transactNewReadOnly(supplier));
+    verify(spyJpaTm, times(6)).load(theEntityKey);
+  }
+
+  @Test
+  public void doTransactionless_retriesJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(JDBCConnectionException.class).when(spyJpaTm).load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class,
+        () -> spyJpaTm.doTransactionless(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+  }
+
   @Test
   void saveNew_throwsExceptionIfEntityExists() {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();

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

@@ -278,6 +278,7 @@ public class DomainBaseToXjcConverterTest {
                 ImmutableSet.of(
                     GracePeriod.forBillingEvent(
                         GracePeriodStatus.RENEW,
+                        domain.getRepoId(),
                         persistResource(
                             new BillingEvent.OneTime.Builder()
                                 .setReason(Reason.RENEW)
@@ -291,6 +292,7 @@ public class DomainBaseToXjcConverterTest {
                                 .build())),
                     GracePeriod.create(
                         GracePeriodStatus.TRANSFER,
+                        domain.getRepoId(),
                         DateTime.parse("1920-01-01T00:00:00Z"),
                         "foo",
                         null)))

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

@@ -122,6 +122,7 @@ final class RdeFixtures {
                 ImmutableSet.of(
                     GracePeriod.forBillingEvent(
                         GracePeriodStatus.RENEW,
+                        domain.getRepoId(),
                         persistResource(
                             new BillingEvent.OneTime.Builder()
                                 .setReason(Reason.RENEW)
@@ -135,6 +136,7 @@ final class RdeFixtures {
                                 .build())),
                     GracePeriod.create(
                         GracePeriodStatus.TRANSFER,
+                        domain.getRepoId(),
                         DateTime.parse("1992-01-01T00:00:00Z"),
                         "foo",
                         null)))

core/src/test/java/google/registry/schema/integration/SqlIntegrationTestSuite.java

@@ -20,6 +20,7 @@ import google.registry.model.billing.BillingEventTest;
 import google.registry.model.contact.ContactResourceTest;
 import google.registry.model.domain.DomainBaseSqlTest;
 import google.registry.model.history.ContactHistoryTest;
+import google.registry.model.history.DomainHistoryTest;
 import google.registry.model.history.HostHistoryTest;
 import google.registry.model.poll.PollMessageTest;
 import google.registry.model.registry.RegistryLockDaoTest;
@@ -77,6 +78,7 @@ import org.junit.runner.RunWith;
   ContactResourceTest.class,
   CursorDaoTest.class,
   DomainBaseSqlTest.class,
+  DomainHistoryTest.class,
   HostHistoryTest.class,
   LockDaoTest.class,
   PollMessageTest.class,

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

@@ -17,6 +17,7 @@ package google.registry.testing;
 import static com.google.common.io.Files.asCharSink;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -28,12 +29,14 @@ import com.google.common.collect.Multimaps;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import io.github.classgraph.ClassGraph;
 import io.github.classgraph.ScanResult;
 import java.io.File;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Map;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -60,53 +63,72 @@ class AppEngineExtensionTest {
   private static final String UNDECLARED_INDEX =
       DECLARED_INDEX.replace("ContactResource", "NoSuchResource");
 
-  private final AppEngineExtension appEngineRule =
+  /**
+   * Sets up test AppEngine instance.
+   *
+   * <p>Not registered as extension since this instance's afterEach() method is also under test. All
+   * methods should call {@link AppEngineExtension#afterEach appEngine.afterEach} explicitly.
+   */
+  private final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
+  private JpaTransactionManager originalJpa;
+
   @RegisterExtension
   final ContextCapturingMetaExtension context = new ContextCapturingMetaExtension();
 
   @BeforeEach
   void beforeEach() throws Exception {
-    appEngineRule.beforeEach(context.getContext());
+    originalJpa = jpaTm();
+    appEngine.beforeEach(context.getContext());
+  }
+
+  @AfterEach
+  void afterEach() {
+    // Note: cannot use isSameInstanceAs() because DummyTransactionManager would throw on any
+    // access.
+    assertWithMessage("Original state not restore. Is appEngine.afterEach not called by this test?")
+        .that(originalJpa == jpaTm())
+        .isTrue();
   }
 
   @Test
   void testTeardown_successNoAutoIndexFile() throws Exception {
-    appEngineRule.afterEach(context.getContext());
+    appEngine.afterEach(context.getContext());
   }
 
   @Test
   void testTeardown_successEmptyAutoIndexFile() throws Exception {
     writeAutoIndexFile("");
-    appEngineRule.afterEach(context.getContext());
+    appEngine.afterEach(context.getContext());
   }
 
   @Test
   void testTeardown_successWhiteSpacesOnlyAutoIndexFile() throws Exception {
     writeAutoIndexFile("  ");
-    appEngineRule.afterEach(context.getContext());
+    appEngine.afterEach(context.getContext());
   }
 
   @Test
   void testTeardown_successOnlyDeclaredIndexesUsed() throws Exception {
     writeAutoIndexFile(DECLARED_INDEX);
-    appEngineRule.afterEach(context.getContext());
+    appEngine.afterEach(context.getContext());
   }
 
   @Test
   void testTeardown_failureUndeclaredIndexesUsed() throws Exception {
     writeAutoIndexFile(UNDECLARED_INDEX);
-    assertThrows(AssertionError.class, () -> appEngineRule.afterEach(context.getContext()));
+    assertThrows(AssertionError.class, () -> appEngine.afterEach(context.getContext()));
   }
 
   @Test
-  void testRegisterOfyEntities_duplicateEntitiesWithSameName_fails() {
+  void testRegisterOfyEntities_duplicateEntitiesWithSameName_fails() throws Exception {
     AppEngineExtension appEngineRule =
         AppEngineExtension.builder()
             .withDatastoreAndCloudSql()
             .withOfyTestEntities(google.registry.testing.TestObject.class, TestObject.class)
             .build();
+    // Thrown before JPA is set up, therefore no need to call afterEach.
     IllegalStateException thrown =
         assertThrows(
             IllegalStateException.class, () -> appEngineRule.beforeEach(context.getContext()));
@@ -118,10 +140,12 @@ class AppEngineExtensionTest {
                 TestObject.class.getName(),
                 "TestObject",
                 google.registry.testing.TestObject.class.getName()));
+    // The class level extension.
+    appEngine.afterEach(context.getContext());
   }
 
   @Test
-  void testOfyEntities_uniqueKinds() {
+  void testOfyEntities_uniqueKinds() throws Exception {
     try (ScanResult scanResult =
         new ClassGraph()
             .enableAnnotationInfo()
@@ -147,10 +171,11 @@ class AppEngineExtensionTest {
           .that(conflictingKinds)
           .isEmpty();
     }
+    appEngine.afterEach(context.getContext());
   }
 
   private void writeAutoIndexFile(String content) throws IOException {
-    asCharSink(new File(appEngineRule.tmpDir, "datastore-indexes-auto.xml"), UTF_8).write(content);
+    asCharSink(new File(appEngine.tmpDir, "datastore-indexes-auto.xml"), UTF_8).write(content);
   }
 
   @Entity

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

@@ -26,6 +26,8 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
@@ -51,7 +53,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Streams;
 import com.google.common.net.InetAddresses;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.cmd.Saver;
@@ -71,6 +72,7 @@ import google.registry.model.domain.DomainAuthInfo;
 import google.registry.model.domain.DomainBase;
 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;
@@ -100,7 +102,6 @@ import google.registry.tmch.LordnTaskUtils;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
@@ -119,12 +120,16 @@ public class DatastoreHelper {
                   String.class));
 
   public static HostResource newHostResource(String hostName) {
+    return newHostResourceWithRoid(hostName, generateNewContactHostRoid());
+  }
+
+  public static HostResource newHostResourceWithRoid(String hostName, String repoId) {
     return new HostResource.Builder()
         .setHostName(hostName)
         .setCreationClientId("TheRegistrar")
         .setPersistedCurrentSponsorClientId("TheRegistrar")
         .setCreationTimeForTest(START_OF_TIME)
-        .setRepoId(generateNewContactHostRoid())
+        .setRepoId(repoId)
         .build();
   }
 
@@ -498,9 +503,10 @@ public class DatastoreHelper {
       DateTime creationTime,
       DateTime expirationTime) {
     String domainName = String.format("%s.%s", label, tld);
+    String repoId = generateNewDomainRoid(tld);
     DomainBase domain =
         new DomainBase.Builder()
-            .setRepoId(generateNewDomainRoid(tld))
+            .setRepoId(repoId)
             .setDomainName(domainName)
             .setPersistedCurrentSponsorClientId("TheRegistrar")
             .setCreationClientId("TheRegistrar")
@@ -513,7 +519,7 @@ public class DatastoreHelper {
                     DesignatedContact.create(Type.TECH, contact.createVKey())))
             .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("fooBAR")))
             .addGracePeriod(
-                GracePeriod.create(GracePeriodStatus.ADD, now.plusDays(10), "foo", null))
+                GracePeriod.create(GracePeriodStatus.ADD, repoId, now.plusDays(10), "foo", null))
             .build();
     HistoryEntry historyEntryDomainCreate =
         persistResource(
@@ -681,20 +687,15 @@ public class DatastoreHelper {
 
   /** Assert that the actual billing event matches the expected one, ignoring IDs. */
   public static void assertBillingEventsEqual(BillingEvent actual, BillingEvent expected) {
-    assertThat(stripBillingEventId(actual)).isEqualTo(stripBillingEventId(expected));
+    assertAboutImmutableObjects().that(actual).isEqualExceptFields(expected, "id");
   }
 
   /** Assert that the actual billing events match the expected ones, ignoring IDs and order. */
   public static void assertBillingEventsEqual(
       Iterable<BillingEvent> actual, Iterable<BillingEvent> expected) {
-    assertThat(
-            Streams.stream(actual)
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Streams.stream(expected)
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()));
+    assertThat(actual)
+        .comparingElementsUsing(immutableObjectCorrespondence("id"))
+        .containsExactlyElementsIn(expected);
   }
 
   /** Assert that the expected billing events are exactly the ones found in the fake Datastore. */
@@ -712,14 +713,7 @@ public class DatastoreHelper {
    */
   public static void assertBillingEventsForResource(
       EppResource resource, BillingEvent... expected) {
-    assertThat(
-            Streams.stream(getBillingEvents(resource))
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected)
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()));
+    assertBillingEventsEqual(getBillingEvents(resource), Arrays.asList(expected));
   }
 
   /** Assert that there are no billing events. */
@@ -727,74 +721,71 @@ public class DatastoreHelper {
     assertThat(getBillingEvents()).isEmpty();
   }
 
-  /** Strips the billing event ID (really, sets it to a constant value) to facilitate comparison. */
+  /**
+   * Strips the billing event ID (really, sets it to a constant value) to facilitate comparison.
+   *
+   * <p>Note: Prefer {@link #assertPollMessagesEqual} when that is suitable.
+   */
   public static BillingEvent stripBillingEventId(BillingEvent billingEvent) {
     return billingEvent.asBuilder().setId(1L).build();
   }
 
   /** Assert that the actual poll message matches the expected one, ignoring IDs. */
   public static void assertPollMessagesEqual(PollMessage actual, PollMessage expected) {
-    assertThat(POLL_MESSAGE_ID_STRIPPER.apply(actual))
-        .isEqualTo(POLL_MESSAGE_ID_STRIPPER.apply(expected));
+    assertAboutImmutableObjects().that(actual).isEqualExceptFields(expected, "id");
   }
 
   /** Assert that the actual poll messages match the expected ones, ignoring IDs and order. */
   public static void assertPollMessagesEqual(
       Iterable<PollMessage> actual, Iterable<PollMessage> expected) {
-    assertThat(Streams.stream(actual).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Streams.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
+    assertThat(actual)
+        .comparingElementsUsing(immutableObjectCorrespondence("id"))
+        .containsExactlyElementsIn(expected);
+  }
+
+  public static void assertPollMessages(String clientId, PollMessage... expected) {
+    assertPollMessagesEqual(getPollMessages(clientId), Arrays.asList(expected));
+  }
+
+  public static void assertPollMessages(PollMessage... expected) {
+    assertPollMessagesEqual(getPollMessages(), Arrays.asList(expected));
   }
 
   public static void assertPollMessagesForResource(EppResource resource, PollMessage... expected) {
-    assertThat(
-            getPollMessages(resource)
-                .stream()
-                .map(POLL_MESSAGE_ID_STRIPPER)
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
+    assertPollMessagesEqual(getPollMessages(resource), Arrays.asList(expected));
   }
 
-  /** Helper to effectively erase the poll message ID to facilitate comparison. */
-  public static final Function<PollMessage, PollMessage> POLL_MESSAGE_ID_STRIPPER =
-      pollMessage -> pollMessage.asBuilder().setId(1L).build();
-
   public static ImmutableList<PollMessage> getPollMessages() {
-    return Streams.stream(ofy().load().type(PollMessage.class)).collect(toImmutableList());
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId) {
-    return Streams.stream(ofy().load().type(PollMessage.class).filter("clientId", clientId))
-        .collect(toImmutableList());
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).filter("clientId", clientId));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(EppResource resource) {
-    return Streams.stream(ofy().load().type(PollMessage.class).ancestor(resource))
-        .collect(toImmutableList());
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).ancestor(resource));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId, DateTime now) {
-    return Streams.stream(
-            ofy()
-                .load()
-                .type(PollMessage.class)
-                .filter("clientId", clientId)
-                .filter("eventTime <=", now.toDate()))
-        .collect(toImmutableList());
+    return ImmutableList.copyOf(
+        ofy()
+            .load()
+            .type(PollMessage.class)
+            .filter("clientId", clientId)
+            .filter("eventTime <=", now.toDate()));
   }
 
   /** Gets all PollMessages associated with the given EppResource. */
   public static ImmutableList<PollMessage> getPollMessages(
       EppResource resource, String clientId, DateTime now) {
-    return Streams.stream(
-            ofy()
-                .load()
-                .type(PollMessage.class)
-                .ancestor(resource)
-                .filter("clientId", clientId)
-                .filter("eventTime <=", now.toDate()))
-        .collect(toImmutableList());
+    return ImmutableList.copyOf(
+        ofy()
+            .load()
+            .type(PollMessage.class)
+            .ancestor(resource)
+            .filter("clientId", clientId)
+            .filter("eventTime <=", now.toDate()));
   }
 
   public static PollMessage getOnlyPollMessage(String clientId) {
@@ -828,6 +819,12 @@ public class DatastoreHelper {
         .collect(onlyElement());
   }
 
+  public static void assertAllocationTokens(AllocationToken... expectedTokens) {
+    assertThat(ofy().load().type(AllocationToken.class).list())
+        .comparingElementsUsing(immutableObjectCorrespondence("updateTimestamp", "creationTime"))
+        .containsExactlyElementsIn(expectedTokens);
+  }
+
   /** Returns a newly allocated, globally unique domain repoId of the format HEX-TLD. */
   public static String generateNewDomainRoid(String tld) {
     return createDomainRepoId(ObjectifyService.allocateId(), tld);

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatastoreHelper.assertAllocationTokens;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -32,11 +33,9 @@ import static org.mockito.Mockito.verify;
 import com.beust.jcommander.ParameterException;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.base.Joiner;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Maps;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -160,6 +159,8 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         "--allowed_client_ids", "TheRegistrar,NewRegistrar",
         "--allowed_tlds", "tld,example",
         "--discount_fraction", "0.5",
+        "--discount_premiums", "true",
+        "--discount_years", "6",
         "--token_status_transitions",
             String.format(
                 "\"%s=NOT_STARTED,%s=VALID,%s=ENDED\"", START_OF_TIME, promoStart, promoEnd));
@@ -170,6 +171,8 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
             .setAllowedClientIds(ImmutableSet.of("TheRegistrar", "NewRegistrar"))
             .setAllowedTlds(ImmutableSet.of("tld", "example"))
             .setDiscountFraction(0.5)
+            .setDiscountPremiums(true)
+            .setDiscountYears(6)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                     .put(START_OF_TIME, TokenStatus.NOT_STARTED)
@@ -309,26 +312,6 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("For UNLIMITED_USE tokens, must specify --token_status_transitions");
   }
 
-  private void assertAllocationTokens(AllocationToken... expectedTokens) {
-    // Using ImmutableObject comparison here is tricky because the creation/updated timestamps are
-    // neither easy nor valuable to test here.
-    ImmutableMap<String, AllocationToken> actualTokens =
-        Maps.uniqueIndex(ofy().load().type(AllocationToken.class), AllocationToken::getToken);
-    assertThat(actualTokens).hasSize(expectedTokens.length);
-    for (AllocationToken expectedToken : expectedTokens) {
-      AllocationToken match = actualTokens.get(expectedToken.getToken());
-      assertThat(match).isNotNull();
-      assertThat(match.getRedemptionHistoryEntry())
-          .isEqualTo(expectedToken.getRedemptionHistoryEntry());
-      assertThat(match.getAllowedClientIds()).isEqualTo(expectedToken.getAllowedClientIds());
-      assertThat(match.getAllowedTlds()).isEqualTo(expectedToken.getAllowedTlds());
-      assertThat(match.getDiscountFraction()).isEqualTo(expectedToken.getDiscountFraction());
-      assertThat(match.getTokenStatusTransitions())
-          .isEqualTo(expectedToken.getTokenStatusTransitions());
-      assertThat(match.getTokenType()).isEqualTo(expectedToken.getTokenType());
-    }
-  }
-
   private AllocationToken createToken(
       String token,
       @Nullable Key<HistoryEntry> redemptionHistoryEntry,

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

@@ -0,0 +1,146 @@
+// 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.persistActiveDomain;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.joda.money.CurrencyUnit.USD;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.model.EppResource;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Reason;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.Period;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.transfer.DomainTransferData;
+import org.joda.money.Money;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit tests for {@link ResaveEntitiesWithUniqueIdCommand}. */
+class ResaveEntitiesWithUniqueIdCommandTest
+    extends CommandTestCase<ResaveEntitiesWithUniqueIdCommand> {
+
+  DomainBase domain;
+  HistoryEntry historyEntry;
+  PollMessage.Autorenew autorenewToResave;
+  BillingEvent.OneTime billingEventToResave;
+
+  @BeforeEach
+  void setUp() {
+    createTld("foobar");
+    domain = persistActiveDomain("foo.foobar");
+    historyEntry = persistHistoryEntry(domain);
+    autorenewToResave = persistAutorenewPollMessage(historyEntry);
+    billingEventToResave = persistBillingEvent(historyEntry);
+  }
+
+  @Test
+  void resaveBillingEvent_succeeds() throws Exception {
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(billingEventToResave)));
+
+    int count = 0;
+    for (BillingEvent.OneTime billingEvent :
+        ofy().load().type(BillingEvent.OneTime.class).ancestor(historyEntry)) {
+      count++;
+      assertThat(billingEvent.getId()).isNotEqualTo(billingEventToResave.getId());
+      assertThat(billingEvent.asBuilder().setId(billingEventToResave.getId()).build())
+          .isEqualTo(billingEventToResave);
+    }
+    assertThat(count).isEqualTo(1);
+  }
+
+  @Test
+  void resaveBillingEvent_failsWhenReferredByDomain() throws Exception {
+    persistResource(
+        domain
+            .asBuilder()
+            .setTransferData(
+                new DomainTransferData.Builder()
+                    .setServerApproveEntities(ImmutableSet.of(billingEventToResave.createVKey()))
+                    .build())
+            .build());
+
+    assertThrows(
+        IllegalStateException.class,
+        () ->
+            runCommand(
+                "--force",
+                "--key_paths_file",
+                writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(billingEventToResave))));
+  }
+
+  private PollMessage.Autorenew persistAutorenewPollMessage(HistoryEntry historyEntry) {
+    return persistResource(
+        new PollMessage.Autorenew.Builder()
+            .setClientId("TheRegistrar")
+            .setEventTime(fakeClock.nowUtc())
+            .setMsg("Test poll message")
+            .setParent(historyEntry)
+            .setAutorenewEndTime(fakeClock.nowUtc().plusDays(365))
+            .setTargetId("foobar.foo")
+            .build());
+  }
+
+  private BillingEvent.OneTime persistBillingEvent(HistoryEntry historyEntry) {
+    return persistResource(
+        new BillingEvent.OneTime.Builder()
+            .setClientId("a registrar")
+            .setTargetId("foo.tld")
+            .setParent(historyEntry)
+            .setReason(Reason.CREATE)
+            .setFlags(ImmutableSet.of(BillingEvent.Flag.ANCHOR_TENANT))
+            .setPeriodYears(2)
+            .setCost(Money.of(USD, 1))
+            .setEventTime(fakeClock.nowUtc())
+            .setBillingTime(fakeClock.nowUtc().plusDays(5))
+            .build());
+  }
+
+  private HistoryEntry persistHistoryEntry(EppResource parent) {
+    return persistResource(
+        new HistoryEntry.Builder()
+            .setParent(parent)
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setPeriod(Period.create(1, Period.Unit.YEARS))
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("foo")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(false)
+            .build());
+  }
+
+  private static String getKeyPathLiteral(Object entity) {
+    Key<?> key = Key.create(entity);
+    return String.format(
+        "\"DomainBase\", \"%s\", \"HistoryEntry\", %s, \"%s\", %s",
+        key.getParent().getParent().getName(), key.getParent().getId(), key.getKind(), key.getId());
+  }
+}

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

@@ -71,7 +71,7 @@ class UniformRapidSuspensionCommandTest
     runCommandForced(
         "--domain_name=evil.tld",
         "--hosts=urs1.example.com,urs2.example.com",
-        "--dsdata={\"keyTag\":1,\"alg\":1,\"digestType\":1,\"digest\":\"abc\"}");
+        "--dsdata=1 1 1 abcd");
     eppVerifier
         .expectClientId("CharlestonRoad")
         .expectSuperuser()
@@ -79,10 +79,9 @@ class UniformRapidSuspensionCommandTest
     assertInStdout("uniform_rapid_suspension --undo");
     assertInStdout("--domain_name evil.tld");
     assertInStdout("--hosts ns1.example.com,ns2.example.com");
-    assertInStdout("--dsdata "
-        + "{\"keyTag\":1,\"algorithm\":2,\"digestType\":3,\"digest\":\"DEAD\"},"
-        + "{\"keyTag\":4,\"algorithm\":5,\"digestType\":6,\"digest\":\"BEEF\"}");
+    assertInStdout("--dsdata 1 2 3 DEAD,4 5 6 BEEF");
     assertNotInStdout("--locks_to_preserve");
+    assertNotInStdout("--restore_client_hold");
   }
 
   @Test
@@ -122,6 +121,29 @@ class UniformRapidSuspensionCommandTest
     assertInStdout("--locks_to_preserve serverDeleteProhibited");
   }
 
+  @Test
+  void testCommand_removeClientHold() throws Exception {
+    persistResource(
+        newDomainBase("evil.tld")
+            .asBuilder()
+            .addStatusValue(StatusValue.CLIENT_HOLD)
+            .addNameserver(ns1.createVKey())
+            .addNameserver(ns2.createVKey())
+            .build());
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--hosts=urs1.example.com,urs2.example.com",
+        "--dsdata=1 1 1 abcd");
+    eppVerifier
+        .expectClientId("CharlestonRoad")
+        .expectSuperuser()
+        .verifySent("uniform_rapid_suspension_with_client_hold.xml");
+    assertInStdout("uniform_rapid_suspension --undo");
+    assertInStdout("--domain_name evil.tld");
+    assertInStdout("--hosts ns1.example.com,ns2.example.com");
+    assertInStdout("--restore_client_hold");
+  }
+
   @Test
   void testUndo_removesLocksReplacesHostsAndDsData() throws Exception {
     persistDomainWithHosts(urs1, urs2);
@@ -149,6 +171,21 @@ class UniformRapidSuspensionCommandTest
     assertNotInStdout("--undo");  // Undo shouldn't print a new undo command.
   }
 
+  @Test
+  void testUndo_restoresClientHolds() throws Exception {
+    persistDomainWithHosts(urs1, urs2);
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--undo",
+        "--hosts=ns1.example.com,ns2.example.com",
+        "--restore_client_hold");
+    eppVerifier
+        .expectClientId("CharlestonRoad")
+        .expectSuperuser()
+        .verifySent("uniform_rapid_suspension_undo_client_hold.xml");
+    assertNotInStdout("--undo"); // Undo shouldn't print a new undo command.
+  }
+
   @Test
   void testFailure_locksToPreserveWithoutUndo() {
     persistActiveDomain("evil.tld");
@@ -177,11 +214,10 @@ class UniformRapidSuspensionCommandTest
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () ->
-                runCommandForced(
-                    "--domain_name=evil.tld",
-                    "--dsdata={\"keyTag\":1,\"alg\":1,\"digestType\":1,\"digest\":\"abc\",\"foo\":1}"));
-    assertThat(thrown).hasMessageThat().contains("Incorrect fields on --dsdata JSON");
+            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1 1 1 abc 1"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("dsRecord 1 1 1 abc 1 should have 4 parts, but has 5");
   }
 
   @Test
@@ -190,11 +226,8 @@ class UniformRapidSuspensionCommandTest
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () ->
-                runCommandForced(
-                    "--domain_name=evil.tld",
-                    "--dsdata={\"keyTag\":1,\"alg\":1,\"digestType\":1}"));
-    assertThat(thrown).hasMessageThat().contains("Incorrect fields on --dsdata JSON");
+            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1 1 1"));
+    assertThat(thrown).hasMessageThat().contains("dsRecord 1 1 1 should have 4 parts, but has 3");
   }
 
   @Test
@@ -203,7 +236,7 @@ class UniformRapidSuspensionCommandTest
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=[1,2,3]"));
-    assertThat(thrown).hasMessageThat().contains("Invalid --dsdata JSON");
+            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1,2,3"));
+    assertThat(thrown).hasMessageThat().contains("dsRecord 1 should have 4 parts, but has 1");
   }
 }

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

@@ -77,6 +77,24 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getDiscountFraction()).isEqualTo(0.15);
   }
 
+  @Test
+  void testUpdateDiscountPremiums() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setDiscountFraction(0.5).setDiscountPremiums(false).build());
+    runCommandForced("--prefix", "token", "--discount_premiums", "true");
+    assertThat(reloadResource(token).shouldDiscountPremiums()).isTrue();
+    runCommandForced("--prefix", "token", "--discount_premiums", "false");
+    assertThat(reloadResource(token).shouldDiscountPremiums()).isFalse();
+  }
+
+  @Test
+  void testUpdateDiscountYears() throws Exception {
+    AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--discount_years", "4");
+    assertThat(reloadResource(token).getDiscountYears()).isEqualTo(4);
+  }
+
   @Test
   void testUpdateStatusTransitions() throws Exception {
     DateTime now = DateTime.now(UTC);

core/src/test/java/google/registry/webdriver/RegistrarConsoleScreenshotTest.java

@@ -238,6 +238,9 @@ class RegistrarConsoleScreenshotTest extends WebDriverTestCase {
     Thread.sleep(1000);
     driver.waitForElement(By.tagName("h1"));
     driver.waitForElement(By.id("reg-app-btn-add")).click();
+    // Attempt to fix flaky tests. The going theory is that the click button CSS animation needs to
+    // finish before the screenshot is captured.
+    Thread.sleep(250);
     driver.diffPage("page");
   }
 

core/src/test/java/google/registry/webdriver/WebDriverTestCase.java

@@ -18,7 +18,7 @@ import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Base class for tests that needs a {@link WebDriverPlusScreenDifferExtension}. */
-@Timeout(60)
+@Timeout(120)
 class WebDriverTestCase {
 
   @RegisterExtension

core/src/test/java/google/registry/whois/DomainWhoisResponseTest.java

@@ -228,11 +228,12 @@ class DomainWhoisResponseTest {
     VKey<ContactResource> adminResourceKey = adminContact.createVKey();
     VKey<ContactResource> techResourceKey = techContact.createVKey();
 
+    String repoId = "3-TLD";
     domainBase =
         persistResource(
             new DomainBase.Builder()
                 .setDomainName("example.tld")
-                .setRepoId("3-TLD")
+                .setRepoId(repoId)
                 .setLastEppUpdateTime(DateTime.parse("2009-05-29T20:13:00Z"))
                 .setCreationTimeForTest(DateTime.parse("2000-10-08T00:45:00Z"))
                 .setRegistrationExpirationTime(DateTime.parse("2010-10-08T00:44:59Z"))
@@ -252,8 +253,9 @@ class DomainWhoisResponseTest {
                 .setDsData(ImmutableSet.of(DelegationSignerData.create(1, 2, 3, "deadface")))
                 .setGracePeriods(
                     ImmutableSet.of(
-                        GracePeriod.create(GracePeriodStatus.ADD, END_OF_TIME, "", null),
-                        GracePeriod.create(GracePeriodStatus.TRANSFER, END_OF_TIME, "", null)))
+                        GracePeriod.create(GracePeriodStatus.ADD, repoId, END_OF_TIME, "", null),
+                        GracePeriod.create(
+                            GracePeriodStatus.TRANSFER, repoId, END_OF_TIME, "", null)))
                 .build());
   }
 

core/src/test/resources/google/registry/flows/domain/domain_check_allocationtoken_promotion.xml

@@ -0,0 +1,39 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <check>
+      <domain:check
+          xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>%DOMAIN%</domain:name>
+      </domain:check>
+    </check>
+    <extension>
+      <allocationToken:allocationToken
+          xmlns:allocationToken=
+              "urn:ietf:params:xml:ns:allocationToken-1.0">
+        abc123
+      </allocationToken:allocationToken>
+      <fee:check
+          xmlns:fee="urn:ietf:params:xml:ns:fee-0.6">
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+        </fee:domain>
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">2</fee:period>
+        </fee:domain>
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">5</fee:period>
+        </fee:domain>
+      </fee:check>
+    </extension>
+    <clTRID>ABC-12345</clTRID>
+  </command>
+</epp>

core/src/test/resources/google/registry/flows/domain/domain_check_allocationtoken_promotion_response.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns:launch="urn:ietf:params:xml:ns:launch-1.0" xmlns:secDNS="urn:ietf:params:xml:ns:secDNS-1.1" xmlns:host="urn:ietf:params:xml:ns:host-1.0" xmlns:fee11="urn:ietf:params:xml:ns:fee-0.11" xmlns:fee12="urn:ietf:params:xml:ns:fee-0.12" xmlns:fee="urn:ietf:params:xml:ns:fee-0.6" xmlns:rgp="urn:ietf:params:xml:ns:rgp-1.0" xmlns="urn:ietf:params:xml:ns:epp-1.0" xmlns:domain="urn:ietf:params:xml:ns:domain-1.0" xmlns:contact="urn:ietf:params:xml:ns:contact-1.0">
+  <response>
+    <result code="1000">
+      <msg>Command completed successfully</msg>
+    </result>
+    <resData>
+      <domain:chkData>
+        <domain:cd>
+          <domain:name avail="true">%DOMAIN%</domain:name>
+        </domain:cd>
+      </domain:chkData>
+    </resData>
+    <extension>
+      <fee:chkData>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+          <fee:fee description="create">%COST_1YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">2</fee:period>
+          <fee:fee description="create">%COST_2YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">5</fee:period>
+          <fee:fee description="create">%COST_5YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+      </fee:chkData>
+    </extension>
+    <trID>
+      <clTRID>ABC-12345</clTRID>
+      <svTRID>server-trid</svTRID>
+    </trID>
+  </response>
+</epp>

core/src/test/resources/google/registry/flows/domain/domain_create_allocationtoken.xml

@@ -4,7 +4,7 @@
       <domain:create
        xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>%DOMAIN%</domain:name>
-        <domain:period unit="y">2</domain:period>
+        <domain:period unit="y">%YEARS%</domain:period>
         <domain:ns>
           <domain:hostObj>ns1.example.net</domain:hostObj>
           <domain:hostObj>ns2.example.net</domain:hostObj>

core/src/test/resources/google/registry/flows/domain/domain_create_premium_allocationtoken.xml

@@ -4,7 +4,7 @@
       <domain:create
           xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>rich.example</domain:name>
-        <domain:period unit="y">2</domain:period>
+        <domain:period unit="y">%YEARS%</domain:period>
         <domain:ns>
           <domain:hostObj>ns1.example.net</domain:hostObj>
           <domain:hostObj>ns2.example.net</domain:hostObj>
@@ -25,7 +25,7 @@
       </allocationToken:allocationToken>
       <fee:create xmlns:fee="urn:ietf:params:xml:ns:fee-0.12">
         <fee:currency>USD</fee:currency>
-        <fee:fee>193.5</fee:fee>
+        <fee:fee>%FEE%</fee:fee>
       </fee:create>
     </extension>
     <clTRID>ABC-12345</clTRID>

core/src/test/resources/google/registry/flows/domain/domain_create_response_premium.xml

@@ -8,13 +8,13 @@
        xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>rich.example</domain:name>
         <domain:crDate>1999-04-03T22:00:00.0Z</domain:crDate>
-        <domain:exDate>2001-04-03T22:00:00.0Z</domain:exDate>
+        <domain:exDate>%EXDATE%</domain:exDate>
       </domain:creData>
     </resData>
     <extension>
       <fee:creData xmlns:fee="urn:ietf:params:xml:ns:fee-0.12">
         <fee:currency>USD</fee:currency>
-        <fee:fee description="create">200.00</fee:fee>
+        <fee:fee description="create">%FEE%</fee:fee>
       </fee:creData>
     </extension>
     <trID>

core/src/test/resources/google/registry/flows/domain/domain_create_response_wildcard.xml

@@ -0,0 +1,19 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <response>
+    <result code="1000">
+      <msg>Command completed successfully</msg>
+    </result>
+    <resData>
+      <domain:creData
+       xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>%DOMAIN%</domain:name>
+        <domain:crDate>%CRDATE%</domain:crDate>
+        <domain:exDate>%EXDATE%</domain:exDate>
+      </domain:creData>
+    </resData>
+    <trID>
+      <clTRID>ABC-12345</clTRID>
+      <svTRID>server-trid</svTRID>
+    </trID>
+  </response>
+</epp>

core/src/test/resources/google/registry/model/schema.txt

@@ -231,12 +231,14 @@ class google.registry.model.domain.secdns.DelegationSignerData {
 }
 class google.registry.model.domain.token.AllocationToken {
   @Id java.lang.String token;
+  boolean discountPremiums;
   com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> redemptionHistoryEntry;
   double discountFraction;
   google.registry.model.CreateAutoTimestamp creationTime;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;
   google.registry.model.common.TimedTransitionProperty<google.registry.model.domain.token.AllocationToken$TokenStatus, google.registry.model.domain.token.AllocationToken$TokenStatusTransition> tokenStatusTransitions;
   google.registry.model.domain.token.AllocationToken$TokenType tokenType;
+  int discountYears;
   java.lang.String domainName;
   java.util.Set<java.lang.String> allowedClientIds;
   java.util.Set<java.lang.String> allowedTlds;

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension.xml

@@ -31,7 +31,7 @@
             <secDNS:keyTag>1</secDNS:keyTag>
             <secDNS:alg>1</secDNS:alg>
             <secDNS:digestType>1</secDNS:digestType>
-            <secDNS:digest>abc</secDNS:digest>
+            <secDNS:digest>ABCD</secDNS:digest>
           </secDNS:dsData>
         </secDNS:add>
       </secDNS:update>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_undo_client_hold.xml

@@ -0,0 +1,38 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <update>
+      <domain:update xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>evil.tld</domain:name>
+        <domain:add>
+          <domain:ns>
+            <domain:hostObj>ns1.example.com</domain:hostObj>
+            <domain:hostObj>ns2.example.com</domain:hostObj>
+          </domain:ns>
+          <domain:status s="clientHold" />
+        </domain:add>
+        <domain:rem>
+          <domain:ns>
+            <domain:hostObj>urs1.example.com</domain:hostObj>
+            <domain:hostObj>urs2.example.com</domain:hostObj>
+          </domain:ns>
+          <domain:status s="serverDeleteProhibited" />
+          <domain:status s="serverTransferProhibited" />
+          <domain:status s="serverUpdateProhibited" />
+        </domain:rem>
+      </domain:update>
+    </update>
+    <extension>
+      <secDNS:update xmlns:secDNS="urn:ietf:params:xml:ns:secDNS-1.1">
+        <secDNS:rem>
+          <secDNS:all>true</secDNS:all>
+        </secDNS:rem>
+      </secDNS:update>
+      <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
+        <metadata:reason>Undo Uniform Rapid Suspension</metadata:reason>
+        <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>
+      </metadata:metadata>
+    </extension>
+    <clTRID>RegistryTool</clTRID>
+  </command>
+</epp>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_with_client_hold.xml

@@ -0,0 +1,46 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <update>
+      <domain:update xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>evil.tld</domain:name>
+        <domain:add>
+          <domain:ns>
+            <domain:hostObj>urs1.example.com</domain:hostObj>
+            <domain:hostObj>urs2.example.com</domain:hostObj>
+          </domain:ns>
+          <domain:status s="serverDeleteProhibited"/>
+          <domain:status s="serverTransferProhibited"/>
+          <domain:status s="serverUpdateProhibited"/>
+        </domain:add>
+        <domain:rem>
+          <domain:ns>
+            <domain:hostObj>ns2.example.com</domain:hostObj>
+            <domain:hostObj>ns1.example.com</domain:hostObj>
+          </domain:ns>
+          <domain:status s="clientHold"/>
+        </domain:rem>
+      </domain:update>
+    </update>
+    <extension>
+      <secDNS:update xmlns:secDNS="urn:ietf:params:xml:ns:secDNS-1.1">
+        <secDNS:rem>
+          <secDNS:all>true</secDNS:all>
+        </secDNS:rem>
+        <secDNS:add>
+          <secDNS:dsData>
+            <secDNS:keyTag>1</secDNS:keyTag>
+            <secDNS:alg>1</secDNS:alg>
+            <secDNS:digestType>1</secDNS:digestType>
+            <secDNS:digest>ABCD</secDNS:digest>
+          </secDNS:dsData>
+        </secDNS:add>
+      </secDNS:update>
+      <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
+        <metadata:reason>Uniform Rapid Suspension</metadata:reason>
+        <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>
+      </metadata:metadata>
+    </extension>
+    <clTRID>RegistryTool</clTRID>
+  </command>
+</epp>

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

@@ -0,0 +1,102 @@
+-- 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.
+
+CREATE TABLE "DomainHistory" (
+   history_revision_id int8 NOT NULL,
+    history_by_superuser boolean NOT NULL,
+    history_registrar_id text,
+    history_modification_time timestamptz NOT NULL,
+    history_reason text NOT NULL,
+    history_requested_by_registrar boolean NOT NULL,
+    history_client_transaction_id text,
+    history_server_transaction_id text,
+    history_type text NOT NULL,
+    history_xml_bytes bytea NOT NULL,
+    admin_contact text,
+    auth_info_repo_id text,
+    auth_info_value text,
+    billing_recurrence_id int8,
+    autorenew_poll_message_id int8,
+    billing_contact text,
+    deletion_poll_message_id int8,
+    domain_name text,
+    idn_table_name text,
+    last_transfer_time timestamptz,
+    launch_notice_accepted_time timestamptz,
+    launch_notice_expiration_time timestamptz,
+    launch_notice_tcn_id text,
+    launch_notice_validator_id text,
+    registrant_contact text,
+    registration_expiration_time timestamptz,
+    smd_id text,
+    subordinate_hosts text[],
+    tech_contact text,
+    tld text,
+    transfer_billing_cancellation_id int8,
+    transfer_billing_recurrence_id int8,
+    transfer_autorenew_poll_message_id int8,
+    transfer_billing_event_id int8,
+    transfer_renew_period_unit text,
+    transfer_renew_period_value int4,
+    transfer_registration_expiration_time timestamptz,
+    transfer_gaining_poll_message_id int8,
+    transfer_losing_poll_message_id int8,
+    transfer_client_txn_id text,
+    transfer_server_txn_id text,
+    transfer_gaining_registrar_id text,
+    transfer_losing_registrar_id text,
+    transfer_pending_expiration_time timestamptz,
+    transfer_request_time timestamptz,
+    transfer_status text,
+    creation_registrar_id text NOT NULL,
+    creation_time timestamptz NOT NULL,
+    current_sponsor_registrar_id text NOT NULL,
+    deletion_time timestamptz,
+    last_epp_update_registrar_id text,
+    last_epp_update_time timestamptz,
+    statuses text[],
+    update_timestamp timestamptz,
+    domain_repo_id text NOT NULL,
+    PRIMARY KEY (history_revision_id)
+);
+
+CREATE TABLE "DomainHistoryHost" (
+   domain_history_history_revision_id int8 NOT NULL,
+   host_repo_id text
+);
+
+ALTER TABLE IF EXISTS "DomainHost" RENAME ns_hosts TO host_repo_id;
+
+CREATE INDEX IDXrh4xmrot9bd63o382ow9ltfig ON "DomainHistory" (creation_time);
+CREATE INDEX IDXaro1omfuaxjwmotk3vo00trwm ON "DomainHistory" (history_registrar_id);
+CREATE INDEX IDXsu1nam10cjes9keobapn5jvxj ON "DomainHistory" (history_type);
+CREATE INDEX IDX6w3qbtgce93cal2orjg1tw7b7 ON "DomainHistory" (history_modification_time);
+
+ALTER TABLE IF EXISTS "DomainHistory"
+   ADD CONSTRAINT fk_domain_history_registrar_id
+   FOREIGN KEY (history_registrar_id)
+   REFERENCES "Registrar";
+
+ALTER TABLE IF EXISTS "DomainHistory"
+   ADD CONSTRAINT fk_domain_history_domain_repo_id
+   FOREIGN KEY (domain_repo_id)
+   REFERENCES "Domain";
+
+ALTER TABLE ONLY public."DomainHistory" ALTER COLUMN history_revision_id
+   SET DEFAULT nextval('public."history_id_sequence"'::regclass);
+
+ALTER TABLE IF EXISTS "DomainHistoryHost"
+   ADD CONSTRAINT FK6b8eqdxwe3guc56tgpm89atx
+   FOREIGN KEY (domain_history_history_revision_id)
+   REFERENCES "DomainHistory";

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

@@ -0,0 +1,41 @@
+-- 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.
+
+create table "GracePeriod" (
+    id  bigserial not null,
+    billing_event_id int8,
+    billing_recurrence_id int8,
+    registrar_id text not null,
+    domain_repo_id text not null,
+    expiration_time timestamptz not null,
+    type text not null,
+    primary key (id)
+);
+
+alter table if exists "GracePeriod"
+   add constraint FK2mys4hojm6ev2g9tmy5aq6m7g
+   foreign key (domain_repo_id)
+   references "Domain";
+
+alter table if exists "GracePeriod"
+   add constraint fk_grace_period_billing_event_id
+   foreign key (billing_event_id)
+   references "BillingEvent";
+
+alter table if exists "GracePeriod"
+   add constraint fk_grace_period_billing_recurrence_id
+   foreign key (billing_recurrence_id)
+   references "BillingRecurrence";
+
+create index IDXj1mtx98ndgbtb1bkekahms18w on "GracePeriod" (domain_repo_id);

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

@@ -0,0 +1,18 @@
+-- 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.
+
+create index IDXoqd7n4hbx86hvlgkilq75olas on "Contact" (contact_id);
+
+alter table if exists "Contact"
+    drop constraint if exists ukoqd7n4hbx86hvlgkilq75olas cascade;

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

@@ -0,0 +1,16 @@
+-- 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.
+
+-- We want this in general, but not until we've migrated
+ALTER TABLE IF EXISTS "Spec11ThreatMatch" DROP CONSTRAINT "fk_safebrowsing_threat_domain_repo_id";

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

@@ -268,18 +268,83 @@ create sequence history_id_sequence start 1 increment 1;
         primary key (repo_id)
     );
 
+    create table "DomainHistory" (
+       history_revision_id int8 not null,
+        history_by_superuser boolean not null,
+        history_registrar_id text,
+        history_modification_time timestamptz not null,
+        history_reason text not null,
+        history_requested_by_registrar boolean not null,
+        history_client_transaction_id text,
+        history_server_transaction_id text,
+        history_type text not null,
+        history_xml_bytes bytea not null,
+        admin_contact text,
+        auth_info_repo_id text,
+        auth_info_value text,
+        billing_recurrence_id int8,
+        autorenew_poll_message_id int8,
+        billing_contact text,
+        deletion_poll_message_id int8,
+        domain_name text,
+        idn_table_name text,
+        last_transfer_time timestamptz,
+        launch_notice_accepted_time timestamptz,
+        launch_notice_expiration_time timestamptz,
+        launch_notice_tcn_id text,
+        launch_notice_validator_id text,
+        registrant_contact text,
+        registration_expiration_time timestamptz,
+        smd_id text,
+        subordinate_hosts text[],
+        tech_contact text,
+        tld text,
+        transfer_billing_cancellation_id int8,
+        transfer_billing_recurrence_id int8,
+        transfer_autorenew_poll_message_id int8,
+        transfer_billing_event_id int8,
+        transfer_renew_period_unit text,
+        transfer_renew_period_value int4,
+        transfer_registration_expiration_time timestamptz,
+        transfer_gaining_poll_message_id int8,
+        transfer_losing_poll_message_id int8,
+        transfer_client_txn_id text,
+        transfer_server_txn_id text,
+        transfer_gaining_registrar_id text,
+        transfer_losing_registrar_id text,
+        transfer_pending_expiration_time timestamptz,
+        transfer_request_time timestamptz,
+        transfer_status text,
+        creation_registrar_id text not null,
+        creation_time timestamptz not null,
+        current_sponsor_registrar_id text not null,
+        deletion_time timestamptz,
+        last_epp_update_registrar_id text,
+        last_epp_update_time timestamptz,
+        statuses text[],
+        update_timestamp timestamptz,
+        domain_repo_id text not null,
+        primary key (history_revision_id)
+    );
+
+    create table "DomainHistoryHost" (
+       domain_history_history_revision_id int8 not null,
+        host_repo_id text
+    );
+
     create table "DomainHost" (
        domain_repo_id text not null,
-        ns_hosts text
+        host_repo_id text
     );
 
     create table "GracePeriod" (
        id  bigserial not null,
-        billing_event_one_time int8,
-        billing_event_recurring int8,
-        registrar_id text,
-        expiration_time timestamptz,
-        type int4,
+        billing_event_id int8,
+        billing_recurrence_id int8,
+        registrar_id text not null,
+        domain_repo_id text not null,
+        expiration_time timestamptz not null,
+        type text not null,
         primary key (id)
     );
 
@@ -515,10 +580,8 @@ create index IDXjny8wuot75b5e6p38r47wdawu on "BillingRecurrence" (recurrence_tim
 create index IDX3y752kr9uh4kh6uig54vemx0l on "Contact" (creation_time);
 create index IDXtm415d6fe1rr35stm33s5mg18 on "Contact" (current_sponsor_registrar_id);
 create index IDXn1f711wicdnooa2mqb7g1m55o on "Contact" (deletion_time);
+create index IDXoqd7n4hbx86hvlgkilq75olas on "Contact" (contact_id);
 create index IDX1p3esngcwwu6hstyua6itn6ff on "Contact" (search_name);
-
-    alter table if exists "Contact" 
-       add constraint UKoqd7n4hbx86hvlgkilq75olas unique (contact_id);
 create index IDXo1xdtpij2yryh0skxe9v91sep on "ContactHistory" (creation_time);
 create index IDXhp33wybmb6tbpr1bq7ttwk8je on "ContactHistory" (history_registrar_id);
 create index IDX9q53px6r302ftgisqifmc6put on "ContactHistory" (history_type);
@@ -528,6 +591,11 @@ create index IDXhsjqiy2lyobfymplb28nm74lm on "Domain" (current_sponsor_registrar
 create index IDX5mnf0wn20tno4b9do88j61klr on "Domain" (deletion_time);
 create index IDXc5aw4pk1vkd6ymhvkpanmoadv on "Domain" (domain_name);
 create index IDXrwl38wwkli1j7gkvtywi9jokq on "Domain" (tld);
+create index IDXrh4xmrot9bd63o382ow9ltfig on "DomainHistory" (creation_time);
+create index IDXaro1omfuaxjwmotk3vo00trwm on "DomainHistory" (history_registrar_id);
+create index IDXsu1nam10cjes9keobapn5jvxj on "DomainHistory" (history_type);
+create index IDX6w3qbtgce93cal2orjg1tw7b7 on "DomainHistory" (history_modification_time);
+create index IDXj1mtx98ndgbtb1bkekahms18w on "GracePeriod" (domain_repo_id);
 create index IDXfg2nnjlujxo6cb9fha971bq2n on "HostHistory" (creation_time);
 create index IDX1iy7njgb7wjmj9piml4l2g0qi on "HostHistory" (history_registrar_id);
 create index IDXkkwbwcwvrdkkqothkiye4jiff on "HostHistory" (host_name);
@@ -554,11 +622,21 @@ create index spec11threatmatch_check_date_idx on "Spec11ThreatMatch" (check_date
        foreign key (revision_id) 
        references "ClaimsList";
 
+    alter table if exists "DomainHistoryHost" 
+       add constraint FK378h8v3j8qd8xtjn2e0bcmrtj 
+       foreign key (domain_history_history_revision_id) 
+       references "DomainHistory";
+
     alter table if exists "DomainHost" 
        add constraint FKeq1guccbre1yk3oosgp2io554 
        foreign key (domain_repo_id) 
        references "Domain";
 
+    alter table if exists "GracePeriod" 
+       add constraint FK2mys4hojm6ev2g9tmy5aq6m7g 
+       foreign key (domain_repo_id) 
+       references "Domain";
+
     alter table if exists "PremiumEntry" 
        add constraint FKo0gw90lpo1tuee56l0nb6y6g5 
        foreign key (revision_id) 

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

@@ -405,16 +405,123 @@ CREATE TABLE public."Domain" (
 );
 
 
+--
+-- Name: DomainHistory; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."DomainHistory" (
+    history_revision_id bigint DEFAULT nextval('public.history_id_sequence'::regclass) NOT NULL,
+    history_by_superuser boolean NOT NULL,
+    history_registrar_id text,
+    history_modification_time timestamp with time zone NOT NULL,
+    history_reason text NOT NULL,
+    history_requested_by_registrar boolean NOT NULL,
+    history_client_transaction_id text,
+    history_server_transaction_id text,
+    history_type text NOT NULL,
+    history_xml_bytes bytea NOT NULL,
+    admin_contact text,
+    auth_info_repo_id text,
+    auth_info_value text,
+    billing_recurrence_id bigint,
+    autorenew_poll_message_id bigint,
+    billing_contact text,
+    deletion_poll_message_id bigint,
+    domain_name text,
+    idn_table_name text,
+    last_transfer_time timestamp with time zone,
+    launch_notice_accepted_time timestamp with time zone,
+    launch_notice_expiration_time timestamp with time zone,
+    launch_notice_tcn_id text,
+    launch_notice_validator_id text,
+    registrant_contact text,
+    registration_expiration_time timestamp with time zone,
+    smd_id text,
+    subordinate_hosts text[],
+    tech_contact text,
+    tld text,
+    transfer_billing_cancellation_id bigint,
+    transfer_billing_recurrence_id bigint,
+    transfer_autorenew_poll_message_id bigint,
+    transfer_billing_event_id bigint,
+    transfer_renew_period_unit text,
+    transfer_renew_period_value integer,
+    transfer_registration_expiration_time timestamp with time zone,
+    transfer_gaining_poll_message_id bigint,
+    transfer_losing_poll_message_id bigint,
+    transfer_client_txn_id text,
+    transfer_server_txn_id text,
+    transfer_gaining_registrar_id text,
+    transfer_losing_registrar_id text,
+    transfer_pending_expiration_time timestamp with time zone,
+    transfer_request_time timestamp with time zone,
+    transfer_status text,
+    creation_registrar_id text NOT NULL,
+    creation_time timestamp with time zone NOT NULL,
+    current_sponsor_registrar_id text NOT NULL,
+    deletion_time timestamp with time zone,
+    last_epp_update_registrar_id text,
+    last_epp_update_time timestamp with time zone,
+    statuses text[],
+    update_timestamp timestamp with time zone,
+    domain_repo_id text NOT NULL
+);
+
+
+--
+-- Name: DomainHistoryHost; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."DomainHistoryHost" (
+    domain_history_history_revision_id bigint NOT NULL,
+    host_repo_id text
+);
+
+
 --
 -- Name: DomainHost; Type: TABLE; Schema: public; Owner: -
 --
 
 CREATE TABLE public."DomainHost" (
     domain_repo_id text NOT NULL,
-    ns_hosts text
+    host_repo_id text
 );
 
 
+--
+-- Name: GracePeriod; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."GracePeriod" (
+    id bigint NOT NULL,
+    billing_event_id bigint,
+    billing_recurrence_id bigint,
+    registrar_id text NOT NULL,
+    domain_repo_id text NOT NULL,
+    expiration_time timestamp with time zone NOT NULL,
+    type text NOT NULL
+);
+
+
+--
+-- Name: GracePeriod_id_seq; Type: SEQUENCE; Schema: public; Owner: -
+--
+
+CREATE SEQUENCE public."GracePeriod_id_seq"
+    START WITH 1
+    INCREMENT BY 1
+    NO MINVALUE
+    NO MAXVALUE
+    CACHE 1;
+
+
+--
+-- Name: GracePeriod_id_seq; Type: SEQUENCE OWNED BY; Schema: public; Owner: -
+--
+
+ALTER SEQUENCE public."GracePeriod_id_seq" OWNED BY public."GracePeriod".id;
+
+
 --
 -- Name: HostHistory; Type: TABLE; Schema: public; Owner: -
 --
@@ -827,6 +934,13 @@ ALTER TABLE ONLY public."BillingRecurrence" ALTER COLUMN billing_recurrence_id S
 ALTER TABLE ONLY public."ClaimsList" ALTER COLUMN revision_id SET DEFAULT nextval('public."ClaimsList_revision_id_seq"'::regclass);
 
 
+--
+-- Name: GracePeriod id; Type: DEFAULT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."GracePeriod" ALTER COLUMN id SET DEFAULT nextval('public."GracePeriod_id_seq"'::regclass);
+
+
 --
 -- Name: PollMessage poll_message_id; Type: DEFAULT; Schema: public; Owner: -
 --
@@ -933,6 +1047,14 @@ ALTER TABLE ONLY public."Cursor"
     ADD CONSTRAINT "Cursor_pkey" PRIMARY KEY (scope, type);
 
 
+--
+-- Name: DomainHistory DomainHistory_pkey; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHistory"
+    ADD CONSTRAINT "DomainHistory_pkey" PRIMARY KEY (history_revision_id);
+
+
 --
 -- Name: Domain Domain_pkey; Type: CONSTRAINT; Schema: public; Owner: -
 --
@@ -941,6 +1063,14 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT "Domain_pkey" PRIMARY KEY (repo_id);
 
 
+--
+-- Name: GracePeriod GracePeriod_pkey; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."GracePeriod"
+    ADD CONSTRAINT "GracePeriod_pkey" PRIMARY KEY (id);
+
+
 --
 -- Name: HostHistory HostHistory_pkey; Type: CONSTRAINT; Schema: public; Owner: -
 --
@@ -1053,14 +1183,6 @@ ALTER TABLE ONLY public."RegistryLock"
     ADD CONSTRAINT idx_registry_lock_repo_id_revision_id UNIQUE (repo_id, revision_id);
 
 
---
--- Name: Contact ukoqd7n4hbx86hvlgkilq75olas; Type: CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Contact"
-    ADD CONSTRAINT ukoqd7n4hbx86hvlgkilq75olas UNIQUE (contact_id);
-
-
 --
 -- Name: idx1iy7njgb7wjmj9piml4l2g0qi; Type: INDEX; Schema: public; Owner: -
 --
@@ -1131,6 +1253,13 @@ CREATE INDEX idx6py6ocrab0ivr76srcd2okpnq ON public."BillingEvent" USING btree (
 CREATE INDEX idx6syykou4nkc7hqa5p8r92cpch ON public."BillingRecurrence" USING btree (event_time);
 
 
+--
+-- Name: idx6w3qbtgce93cal2orjg1tw7b7; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idx6w3qbtgce93cal2orjg1tw7b7 ON public."DomainHistory" USING btree (history_modification_time);
+
+
 --
 -- Name: idx73l103vc5900ig3p4odf0cngt; Type: INDEX; Schema: public; Owner: -
 --
@@ -1166,6 +1295,13 @@ CREATE INDEX idx_registry_lock_registrar_id ON public."RegistryLock" USING btree
 CREATE INDEX idx_registry_lock_verification_code ON public."RegistryLock" USING btree (verification_code);
 
 
+--
+-- Name: idxaro1omfuaxjwmotk3vo00trwm; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idxaro1omfuaxjwmotk3vo00trwm ON public."DomainHistory" USING btree (history_registrar_id);
+
+
 --
 -- Name: idxaydgox62uno9qx8cjlj5lauye; Type: INDEX; Schema: public; Owner: -
 --
@@ -1215,6 +1351,13 @@ CREATE INDEX idxhmv411mdqo5ibn4vy7ykxpmlv ON public."BillingEvent" USING btree (
 CREATE INDEX idxhp33wybmb6tbpr1bq7ttwk8je ON public."ContactHistory" USING btree (history_registrar_id);
 
 
+--
+-- Name: idxj1mtx98ndgbtb1bkekahms18w; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idxj1mtx98ndgbtb1bkekahms18w ON public."GracePeriod" USING btree (domain_repo_id);
+
+
 --
 -- Name: idxj77pfwhui9f0i7wjq6lmibovj; Type: INDEX; Schema: public; Owner: -
 --
@@ -1264,6 +1407,13 @@ CREATE INDEX idxn898pb9mwcg359cdwvolb11ck ON public."BillingRecurrence" USING bt
 CREATE INDEX idxo1xdtpij2yryh0skxe9v91sep ON public."ContactHistory" USING btree (creation_time);
 
 
+--
+-- Name: idxoqd7n4hbx86hvlgkilq75olas; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idxoqd7n4hbx86hvlgkilq75olas ON public."Contact" USING btree (contact_id);
+
+
 --
 -- Name: idxp3usbtvk0v1m14i5tdp4xnxgc; Type: INDEX; Schema: public; Owner: -
 --
@@ -1285,6 +1435,13 @@ CREATE INDEX idxplxf9v56p0wg8ws6qsvd082hk ON public."BillingEvent" USING btree (
 CREATE INDEX idxqa3g92jc17e8dtiaviy4fet4x ON public."BillingCancellation" USING btree (billing_time);
 
 
+--
+-- Name: idxrh4xmrot9bd63o382ow9ltfig; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idxrh4xmrot9bd63o382ow9ltfig ON public."DomainHistory" USING btree (creation_time);
+
+
 --
 -- Name: idxrwl38wwkli1j7gkvtywi9jokq; Type: INDEX; Schema: public; Owner: -
 --
@@ -1292,6 +1449,13 @@ CREATE INDEX idxqa3g92jc17e8dtiaviy4fet4x ON public."BillingCancellation" USING
 CREATE INDEX idxrwl38wwkli1j7gkvtywi9jokq ON public."Domain" USING btree (tld);
 
 
+--
+-- Name: idxsu1nam10cjes9keobapn5jvxj; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE INDEX idxsu1nam10cjes9keobapn5jvxj ON public."DomainHistory" USING btree (history_type);
+
+
 --
 -- Name: idxsudwswtwqnfnx2o1hx4s0k0g5; Type: INDEX; Schema: public; Owner: -
 --
@@ -1379,6 +1543,14 @@ ALTER TABLE ONLY public."RegistryLock"
     ADD CONSTRAINT fk2lhcwpxlnqijr96irylrh1707 FOREIGN KEY (relock_revision_id) REFERENCES public."RegistryLock"(revision_id);
 
 
+--
+-- Name: GracePeriod fk2mys4hojm6ev2g9tmy5aq6m7g; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."GracePeriod"
+    ADD CONSTRAINT fk2mys4hojm6ev2g9tmy5aq6m7g FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
+
+
 --
 -- Name: Domain fk2u3srsfbei272093m3b3xwj23; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1395,6 +1567,14 @@ ALTER TABLE ONLY public."HostHistory"
     ADD CONSTRAINT fk3d09knnmxrt6iniwnp8j2ykga FOREIGN KEY (history_registrar_id) REFERENCES public."Registrar"(registrar_id);
 
 
+--
+-- Name: DomainHistoryHost fk6b8eqdxwe3guc56tgpm89atx; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHistoryHost"
+    ADD CONSTRAINT fk6b8eqdxwe3guc56tgpm89atx FOREIGN KEY (domain_history_history_revision_id) REFERENCES public."DomainHistory"(history_revision_id);
+
+
 --
 -- Name: ClaimsEntry fk6sc6at5hedffc0nhdcab6ivuq; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1531,6 +1711,22 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT fk_domain_deletion_poll_message_id FOREIGN KEY (deletion_poll_message_id) REFERENCES public."PollMessage"(poll_message_id);
 
 
+--
+-- Name: DomainHistory fk_domain_history_domain_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHistory"
+    ADD CONSTRAINT fk_domain_history_domain_repo_id FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
+
+
+--
+-- Name: DomainHistory fk_domain_history_registrar_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHistory"
+    ADD CONSTRAINT fk_domain_history_registrar_id FOREIGN KEY (history_registrar_id) REFERENCES public."Registrar"(registrar_id);
+
+
 --
 -- Name: Domain fk_domain_registrant_contact; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1592,7 +1788,23 @@ ALTER TABLE ONLY public."Domain"
 --
 
 ALTER TABLE ONLY public."DomainHost"
-    ADD CONSTRAINT fk_domainhost_host_valid FOREIGN KEY (ns_hosts) REFERENCES public."HostResource"(repo_id);
+    ADD CONSTRAINT fk_domainhost_host_valid FOREIGN KEY (host_repo_id) REFERENCES public."HostResource"(repo_id);
+
+
+--
+-- Name: GracePeriod fk_grace_period_billing_event_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."GracePeriod"
+    ADD CONSTRAINT fk_grace_period_billing_event_id FOREIGN KEY (billing_event_id) REFERENCES public."BillingEvent"(billing_event_id);
+
+
+--
+-- Name: GracePeriod fk_grace_period_billing_recurrence_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."GracePeriod"
+    ADD CONSTRAINT fk_grace_period_billing_recurrence_id FOREIGN KEY (billing_recurrence_id) REFERENCES public."BillingRecurrence"(billing_recurrence_id);
 
 
 --
@@ -1659,14 +1871,6 @@ ALTER TABLE ONLY public."PollMessage"
     ADD CONSTRAINT fk_poll_message_transfer_response_losing_registrar_id FOREIGN KEY (transfer_response_losing_registrar_id) REFERENCES public."Registrar"(registrar_id);
 
 
---
--- Name: Spec11ThreatMatch fk_safebrowsing_threat_domain_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Spec11ThreatMatch"
-    ADD CONSTRAINT fk_safebrowsing_threat_domain_repo_id FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
-
-
 --
 -- Name: DomainHost fkfmi7bdink53swivs390m2btxg; Type: FK CONSTRAINT; Schema: public; Owner: -
 --

java-format/google-java-format-git-diff.sh

@@ -69,7 +69,7 @@ function callGoogleJavaFormatDiff() {
     "format")
       showNoncompliantFiles "$forkPoint" "\033[1mReformatting: "
       callResult=$(git diff -U0 ${forkPoint} | \
-          ${SCRIPT_DIR}/google-java-format-diff.py \ 
+          ${SCRIPT_DIR}/google-java-format-diff.py \
           --google-java-format-jar "${SCRIPT_DIR}/${JAR_NAME}" \
           -p1 -i)
       ;;

release/builder/build.sh

@@ -17,12 +17,15 @@ set -e
 apt-get update -y
 apt-get install locales -y
 locale-gen en_US.UTF-8
-apt-get install apt-utils -y
+apt-get install apt-utils gnupg -y
 apt-get upgrade -y
 # Install Java
 apt-get install openjdk-11-jdk-headless -y
-# Install npm
-apt-get install npm -y
+# Install Python
+apt-get install python -y
+# Install Node
+curl -sL https://deb.nodesource.com/setup_current.x | bash -
+apt-get install -y nodejs
 # Install gcloud
 # Cribbed from https://cloud.google.com/sdk/docs/quickstart-debian-ubuntu
 apt-get install lsb-release -y

release/cloudbuild-deploy.yaml

@@ -38,21 +38,31 @@ steps:
       --ciphertext-file=- --plaintext-file=tool-credential.json \
       --location=global --keyring=nomulus-tool-keyring --key=nomulus-tool-key
 # Deploy the Spec11 pipeline to GCS.
-- name: 'gcr.io/$PROJECT_ID/nomulus-tool:latest'
-  args:
-  - -e
-  - ${_ENV}
-  - --credential
-  - tool-credential.json
-  - deploy_spec11_pipeline
+#- name: 'gcr.io/$PROJECT_ID/builder:latest'
+#  entrypoint: /bin/bash
+#  args:
+#  - -c
+#  - |
+#    set -e
+#    if [ ${_ENV} == production ]; then
+#      project_id="domain-registry"
+#    else
+#      project_id="domain-registry-${_ENV}"
+#    fi
+#    echo "gs://$${project_id}-beam/cloudsql/admin_credential.enc"
+#    gsutil cp gs://$PROJECT_ID-deploy/${TAG_NAME}/nomulus.jar .
+#    java -jar nomulus.jar -e ${_ENV} --credential tool-credential.json \
+#      --sql_access_info \
+#      "gs://$${project_id}-beam/cloudsql/admin_credential.enc" \
+#      deploy_spec11_pipeline --project $${project_id}
 # Deploy the invoicing pipeline to GCS.
-- name: 'gcr.io/$PROJECT_ID/nomulus-tool:latest'
-  args:
-  - -e
-  - ${_ENV}
-  - --credential
-  - tool-credential.json
-  - deploy_invoicing_pipeline
+#- name: 'gcr.io/$PROJECT_ID/nomulus-tool:latest'
+#  args:
+#  - -e
+#  - ${_ENV}
+#  - --credential
+#  - tool-credential.json
+#  - deploy_invoicing_pipeline
 # Save the deployed tag for the current environment on GCS. Because of b/137891685
 # which causes the for-loop in the next step to fail, this may not be the last step.
 # TODO(weiminyu): do this in last step.

release/cloudbuild-nomulus.yaml

@@ -20,16 +20,16 @@ steps:
 - name: 'gcr.io/${PROJECT_ID}/builder:latest'
   args: ['mkdir', 'nomulus']
 # Run tests
-- name: 'gcr.io/${PROJECT_ID}/builder:latest'
-  # Set home for Gradle caches. Must be consistent with last step below
-  # and ./build_nomulus_for_env.sh
-  env: [ 'GRADLE_USER_HOME=./cloudbuild-caches' ]
-  args: ['./gradlew',
-         'test',
-         '-PskipDockerIncompatibleTests=true',
-         '-PmavenUrl=gcs://domain-registry-maven-repository/maven',
-         '-PpluginsUrl=gcs://domain-registry-maven-repository/plugins'
-        ]
+#- name: 'gcr.io/${PROJECT_ID}/builder:latest'
+#  # Set home for Gradle caches. Must be consistent with last step below
+#  # and ./build_nomulus_for_env.sh
+#  env: [ 'GRADLE_USER_HOME=./cloudbuild-caches' ]
+#  args: ['./gradlew',
+#         'test',
+#         '-PskipDockerIncompatibleTests=true',
+#         '-PmavenUrl=gcs://domain-registry-maven-repository/maven',
+#         '-PpluginsUrl=gcs://domain-registry-maven-repository/plugins'
+#        ]
 # Build the tool binary and image.
 - name: 'gcr.io/${PROJECT_ID}/builder:latest'
   args: ['release/build_nomulus_for_env.sh', 'tool', 'output']

util/src/main/java/google/registry/util/CollectionUtils.java

@@ -75,38 +75,39 @@ public class CollectionUtils {
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V> ImmutableSet<V> nullSafeImmutableCopy(Set<V> data) {
+  public static <V> ImmutableSet<V> nullSafeImmutableCopy(@Nullable Set<V> data) {
     return data == null ? null : ImmutableSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link List}. */
-  public static <V> ImmutableList<V> nullSafeImmutableCopy(List<V> data) {
+  public static <V> ImmutableList<V> nullSafeImmutableCopy(@Nullable List<V> data) {
     return data == null ? null : ImmutableList.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V> ImmutableSet<V> nullToEmptyImmutableCopy(Set<V> data) {
+  public static <V> ImmutableSet<V> nullToEmptyImmutableCopy(@Nullable Set<V> data) {
     return data == null ? ImmutableSet.of() : ImmutableSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V extends Comparable<V>>
-      ImmutableSortedSet<V> nullToEmptyImmutableSortedCopy(Set<V> data) {
+  public static <V extends Comparable<V>> ImmutableSortedSet<V> nullToEmptyImmutableSortedCopy(
+      @Nullable Set<V> data) {
     return data == null ? ImmutableSortedSet.of() : ImmutableSortedSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link SortedMap}. */
-  public static <K, V> ImmutableSortedMap<K, V> nullToEmptyImmutableCopy(SortedMap<K, V> data) {
+  public static <K, V> ImmutableSortedMap<K, V> nullToEmptyImmutableCopy(
+      @Nullable SortedMap<K, V> data) {
     return data == null ? ImmutableSortedMap.of() : ImmutableSortedMap.copyOfSorted(data);
   }
 
   /** Defensive copy helper for {@link List}. */
-  public static <V> ImmutableList<V> nullToEmptyImmutableCopy(List<V> data) {
+  public static <V> ImmutableList<V> nullToEmptyImmutableCopy(@Nullable List<V> data) {
     return data == null ? ImmutableList.of() : ImmutableList.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Map}. */
-  public static <K, V> ImmutableMap<K, V> nullToEmptyImmutableCopy(Map<K, V> data) {
+  public static <K, V> ImmutableMap<K, V> nullToEmptyImmutableCopy(@Nullable Map<K, V> data) {
     return data == null ? ImmutableMap.of() : ImmutableMap.copyOf(data);
   }
 

util/src/main/java/google/registry/util/PreconditionsUtils.java

@@ -29,33 +29,36 @@ public class PreconditionsUtils {
    * preferable to throw an IAE instead of an NPE, such as where we want an IAE to indicate that
    * it's just a bad argument/parameter and reserve NPEs for bugs and unexpected null values.
    */
-  public static <T> T checkArgumentNotNull(T reference) {
+  public static <T> T checkArgumentNotNull(@Nullable T reference) {
     checkArgument(reference != null);
     return reference;
   }
 
   /** Checks whether the provided reference is null, throws IAE if it is, and returns it if not. */
-  public static <T> T checkArgumentNotNull(T reference, @Nullable Object errorMessage) {
+  public static <T> T checkArgumentNotNull(@Nullable T reference, @Nullable Object errorMessage) {
     checkArgument(reference != null, errorMessage);
     return reference;
   }
 
   /** Checks whether the provided reference is null, throws IAE if it is, and returns it if not. */
   public static <T> T checkArgumentNotNull(
-      T reference, @Nullable String errorMessageTemplate, @Nullable Object... errorMessageArgs) {
+      @Nullable T reference,
+      @Nullable String errorMessageTemplate,
+      @Nullable Object... errorMessageArgs) {
     checkArgument(reference != null, errorMessageTemplate, errorMessageArgs);
     return reference;
   }
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
-  public static <T> T checkArgumentPresent(Optional<T> reference) {
+  public static <T> T checkArgumentPresent(@Nullable Optional<T> reference) {
     checkArgumentNotNull(reference);
     checkArgument(reference.isPresent());
     return reference.get();
   }
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
-  public static <T> T checkArgumentPresent(Optional<T> reference, @Nullable Object errorMessage) {
+  public static <T> T checkArgumentPresent(
+      @Nullable Optional<T> reference, @Nullable Object errorMessage) {
     checkArgumentNotNull(reference, errorMessage);
     checkArgument(reference.isPresent(), errorMessage);
     return reference.get();
@@ -63,7 +66,7 @@ public class PreconditionsUtils {
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
   public static <T> T checkArgumentPresent(
-      Optional<T> reference,
+      @Nullable Optional<T> reference,
       @Nullable String errorMessageTemplate,
       @Nullable Object... errorMessageArgs) {
     checkArgumentNotNull(reference, errorMessageTemplate, errorMessageArgs);