Refactor LevelDbFileBuilder to accept DS Entity (#599)

* Refactor LevelDbFileBuilder to accept DS Entity

Builder now can directly work with Datastore Entity objects.
No need to wrap data in ComparableEntity.

build.gradle

@@ -197,6 +197,10 @@ task runPresubmits(type: Exec) {
   args('config/presubmits.py')
 }
 
+def javadocSource = []
+def javadocClasspath = []
+def javadocDependentTasks = []
+
 subprojects {
   // Skip no-op project
   if (project.name == 'services') return
@@ -317,6 +321,10 @@ subprojects {
       }
     }
   }
+
+  javadocSource << project.sourceSets.main.allJava
+  javadocClasspath << project.sourceSets.main.compileClasspath
+  javadocDependentTasks << project.tasks.compileJava
 }
 
 // If "-P verboseTestOutput=true" is passed in, configure all subprojects to dump all of their
@@ -445,3 +453,18 @@ task javaIncrementalFormatApply {
 }
 
 tasks.build.dependsOn(tasks.javaIncrementalFormatCheck)
+
+task javadoc(type: Javadoc) {
+  source javadocSource
+  classpath = files(javadocClasspath)
+  destinationDir = file("${buildDir}/docs/javadoc")
+  // In a lot of places we don't write @return so suppress warnings about that.
+  options.addBooleanOption('Xdoclint:all,-missing', true)
+  options.addBooleanOption("-allow-script-in-comments",true)
+  options.tags = ["type:a:Generic Type",
+                  "error:a:Expected Error"]
+}
+
+tasks.build.dependsOn(tasks.javadoc)
+
+javadocDependentTasks.each { tasks.javadoc.shouldRunAfter(it) }

common/src/main/java/google/registry/util/Sleeper.java

@@ -21,7 +21,6 @@ import org.joda.time.ReadableDuration;
  * An object which accepts requests to put the current thread to sleep.
  *
  * @see SystemSleeper
- * @see google.registry.testing.FakeSleeper
  */
 @ThreadSafe
 public interface Sleeper {

core/build.gradle

@@ -880,6 +880,9 @@ task standardTest(type: FilteringTest) {
   // forkEvery 1
 
   // Sets the maximum number of test executors that may exist at the same time.
+  // Also, Gradle executes tests in 1 thread and some of our test infrastructures
+  // depend on that, e.g. DualDatabaseTestInvocationContextProvider injects
+  // different implementation of TransactionManager into TransactionManagerFactory.
   maxParallelForks 5
 
   systemProperty 'test.projectRoot', rootProject.projectRootDir

core/src/main/java/google/registry/batch/ExpandRecurringBillingEventsAction.java

@@ -155,89 +155,100 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
       }
       int numBillingEventsSaved = 0;
       try {
-        numBillingEventsSaved = tm().transactNew(() -> {
-          ImmutableSet.Builder<OneTime> syntheticOneTimesBuilder =
-              new ImmutableSet.Builder<>();
-          final Registry tld = Registry.get(getTldFromDomainName(recurring.getTargetId()));
+        numBillingEventsSaved =
+            tm().transactNew(
+                    () -> {
+                      ImmutableSet.Builder<OneTime> syntheticOneTimesBuilder =
+                          new ImmutableSet.Builder<>();
+                      final Registry tld =
+                          Registry.get(getTldFromDomainName(recurring.getTargetId()));
 
-          // Determine the complete set of times at which this recurring event should occur
-          // (up to and including the runtime of the mapreduce).
-          Iterable<DateTime> eventTimes =
-              recurring.getRecurrenceTimeOfYear().getInstancesInRange(Range.closed(
-                  recurring.getEventTime(),
-                  earliestOf(recurring.getRecurrenceEndTime(), executeTime)));
+                      // Determine the complete set of times at which this recurring event should
+                      // occur (up to and including the runtime of the mapreduce).
+                      Iterable<DateTime> eventTimes =
+                          recurring
+                              .getRecurrenceTimeOfYear()
+                              .getInstancesInRange(
+                                  Range.closed(
+                                      recurring.getEventTime(),
+                                      earliestOf(recurring.getRecurrenceEndTime(), executeTime)));
 
-          // Convert these event times to billing times
-          final ImmutableSet<DateTime> billingTimes =
-              getBillingTimesInScope(eventTimes, cursorTime, executeTime, tld);
+                      // Convert these event times to billing times
+                      final ImmutableSet<DateTime> billingTimes =
+                          getBillingTimesInScope(eventTimes, cursorTime, executeTime, tld);
 
-          Key<? extends EppResource> domainKey = recurring.getParentKey().getParent();
-          Iterable<OneTime> oneTimesForDomain =
-              ofy().load().type(OneTime.class).ancestor(domainKey);
+                      Key<? extends EppResource> domainKey = recurring.getParentKey().getParent();
+                      Iterable<OneTime> oneTimesForDomain =
+                          ofy().load().type(OneTime.class).ancestor(domainKey);
 
-          // Determine the billing times that already have OneTime events persisted.
-          ImmutableSet<DateTime> existingBillingTimes =
-              getExistingBillingTimes(oneTimesForDomain, recurring);
+                      // Determine the billing times that already have OneTime events persisted.
+                      ImmutableSet<DateTime> existingBillingTimes =
+                          getExistingBillingTimes(oneTimesForDomain, recurring);
 
-          ImmutableSet.Builder<HistoryEntry> historyEntriesBuilder =
-              new ImmutableSet.Builder<>();
-          // Create synthetic OneTime events for all billing times that do not yet have an event
-          // persisted.
-          for (DateTime billingTime : difference(billingTimes, existingBillingTimes)) {
-            // Construct a new HistoryEntry that parents over the OneTime
-            HistoryEntry historyEntry = new HistoryEntry.Builder()
-                .setBySuperuser(false)
-                .setClientId(recurring.getClientId())
-                .setModificationTime(tm().getTransactionTime())
-                .setParent(domainKey)
-                .setPeriod(Period.create(1, YEARS))
-                .setReason("Domain autorenewal by ExpandRecurringBillingEventsAction")
-                .setRequestedByRegistrar(false)
-                .setType(DOMAIN_AUTORENEW)
-                // Don't write a domain transaction record if the recurrence was ended prior to the
-                // billing time (i.e. a domain was deleted during the autorenew grace period).
-                .setDomainTransactionRecords(
-                    recurring.getRecurrenceEndTime().isBefore(billingTime)
-                        ? ImmutableSet.of()
-                        : ImmutableSet.of(
-                            DomainTransactionRecord.create(
-                                tld.getTldStr(),
-                                // We report this when the autorenew grace period ends
-                                billingTime,
-                                TransactionReportField.netRenewsFieldFromYears(1),
-                                1)))
-                .build();
-            historyEntriesBuilder.add(historyEntry);
+                      ImmutableSet.Builder<HistoryEntry> historyEntriesBuilder =
+                          new ImmutableSet.Builder<>();
+                      // Create synthetic OneTime events for all billing times that do not yet have
+                      // an event persisted.
+                      for (DateTime billingTime : difference(billingTimes, existingBillingTimes)) {
+                        // Construct a new HistoryEntry that parents over the OneTime
+                        HistoryEntry historyEntry =
+                            new HistoryEntry.Builder()
+                                .setBySuperuser(false)
+                                .setClientId(recurring.getClientId())
+                                .setModificationTime(tm().getTransactionTime())
+                                .setParent(domainKey)
+                                .setPeriod(Period.create(1, YEARS))
+                                .setReason(
+                                    "Domain autorenewal by ExpandRecurringBillingEventsAction")
+                                .setRequestedByRegistrar(false)
+                                .setType(DOMAIN_AUTORENEW)
+                                // Don't write a domain transaction record if the recurrence was
+                                // ended prior to the billing time (i.e. a domain was deleted
+                                // during the autorenew grace period).
+                                .setDomainTransactionRecords(
+                                    recurring.getRecurrenceEndTime().isBefore(billingTime)
+                                        ? ImmutableSet.of()
+                                        : ImmutableSet.of(
+                                            DomainTransactionRecord.create(
+                                                tld.getTldStr(),
+                                                // We report this when the autorenew grace period
+                                                // ends
+                                                billingTime,
+                                                TransactionReportField.netRenewsFieldFromYears(1),
+                                                1)))
+                                .build();
+                        historyEntriesBuilder.add(historyEntry);
 
-            DateTime eventTime = billingTime.minus(tld.getAutoRenewGracePeriodLength());
-            // Determine the cost for a one-year renewal.
-            Money renewCost = getDomainRenewCost(recurring.getTargetId(), eventTime, 1);
-            syntheticOneTimesBuilder.add(new OneTime.Builder()
-                .setBillingTime(billingTime)
-                .setClientId(recurring.getClientId())
-                .setCost(renewCost)
-                .setEventTime(eventTime)
-                .setFlags(union(recurring.getFlags(), Flag.SYNTHETIC))
-                .setParent(historyEntry)
-                .setPeriodYears(1)
-                .setReason(recurring.getReason())
-                .setSyntheticCreationTime(executeTime)
-                .setCancellationMatchingBillingEvent(Key.create(recurring))
-                .setTargetId(recurring.getTargetId())
-                .build());
-          }
-          Set<HistoryEntry> historyEntries = historyEntriesBuilder.build();
-          Set<OneTime> syntheticOneTimes = syntheticOneTimesBuilder.build();
-          if (!isDryRun) {
-            ImmutableSet<ImmutableObject> entitiesToSave =
-                new ImmutableSet.Builder<ImmutableObject>()
-                    .addAll(historyEntries)
-                    .addAll(syntheticOneTimes)
-                    .build();
-            ofy().save().entities(entitiesToSave).now();
-          }
-          return syntheticOneTimes.size();
-        });
+                        DateTime eventTime = billingTime.minus(tld.getAutoRenewGracePeriodLength());
+                        // Determine the cost for a one-year renewal.
+                        Money renewCost = getDomainRenewCost(recurring.getTargetId(), eventTime, 1);
+                        syntheticOneTimesBuilder.add(
+                            new OneTime.Builder()
+                                .setBillingTime(billingTime)
+                                .setClientId(recurring.getClientId())
+                                .setCost(renewCost)
+                                .setEventTime(eventTime)
+                                .setFlags(union(recurring.getFlags(), Flag.SYNTHETIC))
+                                .setParent(historyEntry)
+                                .setPeriodYears(1)
+                                .setReason(recurring.getReason())
+                                .setSyntheticCreationTime(executeTime)
+                                .setCancellationMatchingBillingEvent(recurring.createVKey())
+                                .setTargetId(recurring.getTargetId())
+                                .build());
+                      }
+                      Set<HistoryEntry> historyEntries = historyEntriesBuilder.build();
+                      Set<OneTime> syntheticOneTimes = syntheticOneTimesBuilder.build();
+                      if (!isDryRun) {
+                        ImmutableSet<ImmutableObject> entitiesToSave =
+                            new ImmutableSet.Builder<ImmutableObject>()
+                                .addAll(historyEntries)
+                                .addAll(syntheticOneTimes)
+                                .build();
+                        ofy().save().entities(entitiesToSave).now();
+                      }
+                      return syntheticOneTimes.size();
+                    });
       } catch (Throwable t) {
         getContext().incrementCounter("error: " + t.getClass().getSimpleName());
         getContext().incrementCounter(ERROR_COUNTER);
@@ -279,7 +290,8 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
       return Streams.stream(oneTimesForDomain)
           .filter(
               billingEvent ->
-                  Key.create(recurringEvent)
+                  recurringEvent
+                      .createVKey()
                       .equals(billingEvent.getCancellationMatchingBillingEvent()))
           .map(OneTime::getBillingTime)
           .collect(toImmutableSet());

core/src/main/java/google/registry/bigquery/BigqueryConnection.java

@@ -442,10 +442,8 @@ public class BigqueryConnection implements AutoCloseable {
    * Returns the result of calling queryToLocalTable, but synchronously to avoid spawning new
    * background threads, which App Engine doesn't support.
    *
-   * @see <a href="https://cloud.google.com/appengine/docs/standard/java/runtime#Threads">
-   *   App Engine Runtime</a>
-   *
-   * <p>Returns the results of the query in an ImmutableTable on success.
+   * @see <a href="https://cloud.google.com/appengine/docs/standard/java/runtime#Threads">App Engine
+   *     Runtime</a>
    */
   public ImmutableTable<Integer, TableFieldSchema, Object> queryToLocalTableSync(String querySql) {
     Job job = new Job()
@@ -576,8 +574,6 @@ public class BigqueryConnection implements AutoCloseable {
 
   /**
    * Launch a job, wait for it to complete, but <i>do not</i> check for errors.
-   *
-   * @throws BigqueryJobFailureException
    */
   public Job runJob(Job job, @Nullable AbstractInputStreamContent data) {
     return checkJob(waitForJob(launchJob(job, data)));

core/src/main/java/google/registry/dns/writer/dnsupdate/DnsMessageTransport.java

@@ -84,7 +84,7 @@ public class DnsMessageTransport {
    * @param query a message to send
    * @return the response received from the server
    * @throws IOException if the Socket input/output streams throws one
-   * @throws IllegalArgumentException if the query is too large to be sent (> 65535 bytes)
+   * @throws IllegalArgumentException if the query is too large to be sent (> 65535 bytes)
    */
   public Message send(Message query) throws IOException {
     try (Socket socket = factory.createSocket(InetAddress.getByName(updateHost), DNS_PORT)) {

core/src/main/java/google/registry/export/BigqueryPollJobAction.java

@@ -73,8 +73,10 @@ public class BigqueryPollJobAction implements Runnable {
 
   @Override
   public void run() {
-    checkJobOutcome();  // Throws a NotModifiedException if the job hasn't completed.
-    if (payload == null || payload.length == 0) {
+    boolean jobOutcome =
+        checkJobOutcome(); // Throws a NotModifiedException if the job hasn't completed.
+    // If the job failed, do not enqueue the next step.
+    if (!jobOutcome || payload == null || payload.length == 0) {
       return;
     }
     // If there is a payload, it's a chained task, so enqueue it.

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

@@ -42,8 +42,8 @@ import javax.xml.stream.events.XMLEvent;
 /**
  * Sanitizes sensitive data in incoming/outgoing EPP XML messages.
  *
- * <p>Current implementation masks user credentials (text following <pw> and <newPW> tags) as
- * follows:
+ * <p>Current implementation masks user credentials (text following &lt;pw&gt; and &lt;newPW&gt;
+ * tags) as follows:
  *
  * <ul>
  *   <li>A control character (in ranges [0 - 1F] and [7F - 9F]) is replaced with 'C'.

core/src/main/java/google/registry/flows/custom/DomainCheckFlowCustomLogic.java

@@ -75,7 +75,7 @@ public class DomainCheckFlowCustomLogic extends BaseFlowCustomLogic {
 
     /**
      * The time to perform the domain check as of. This defaults to the current time, but can be
-     * overridden in v>=0.12 of the fee extension.
+     * overridden in v>=0.12 of the fee extension.
      */
     public abstract DateTime asOfDate();
 
@@ -105,7 +105,7 @@ public class DomainCheckFlowCustomLogic extends BaseFlowCustomLogic {
 
     /**
      * The time to perform the domain check as of. This defaults to the current time, but can be
-     * overridden in v>=0.12 of the fee extension.
+     * overridden in v>=0.12 of the fee extension.
      */
     public abstract DateTime asOfDate();
 

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

@@ -534,7 +534,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setPeriodYears(years)
         .setCost(feesAndCredits.getCreateCost())
         .setEventTime(now)
-        .setAllocationToken(allocationToken.map(Key::create).orElse(null))
+        .setAllocationToken(allocationToken.map(AllocationToken::createVKey).orElse(null))
         .setBillingTime(
             now.plus(
                 isAnchorTenant

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

@@ -34,6 +34,7 @@ import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Parent;
+import google.registry.persistence.VKey;
 import java.lang.reflect.Array;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
@@ -139,10 +140,13 @@ public class ModelUtils {
 
       // If the field's type is the same as the field's class object, then it's a non-parameterized
       // type, and thus we just add it directly. We also don't bother looking at the parameterized
-      // types of Key objects, since they are just references to other objects and don't actually
-      // embed themselves in the persisted object anyway.
+      // types of Key and VKey objects, since they are just references to other objects and don't
+      // actually embed themselves in the persisted object anyway.
       Class<?> fieldClazz = field.getType();
       Type fieldType = field.getGenericType();
+      if (VKey.class.equals(fieldClazz)) {
+        continue;
+      }
       builder.add(fieldClazz);
       if (fieldType.equals(fieldClazz) || Key.class.equals(clazz)) {
         continue;

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

@@ -50,7 +50,8 @@ import java.util.stream.Collectors;
 public class OteStats {
 
   /**
-   * Returns the statistics about the OT&E actions that have been taken by a particular registrar.
+   * Returns the statistics about the OT&E actions that have been taken by a particular
+   * registrar.
    */
   public static OteStats getFromRegistrar(String registrarName) {
     return new OteStats().recordRegistrarHistory(registrarName);

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

@@ -30,7 +30,7 @@ public class UpdateAutoTimestamp extends ImmutableObject {
 
   DateTime timestamp;
 
-  /** Returns the timestamp, or {@link START_OF_TIME} if it's null. */
+  /** Returns the timestamp, or {@code START_OF_TIME} if it's null. */
   public DateTime getTimestamp() {
     return Optional.ofNullable(timestamp).orElse(START_OF_TIME);
   }

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

@@ -30,6 +30,7 @@ import com.google.common.collect.Sets;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.IgnoreSave;
 import com.googlecode.objectify.annotation.Index;
 import com.googlecode.objectify.annotation.Parent;
@@ -43,14 +44,28 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
+import google.registry.persistence.VKey;
+import google.registry.persistence.WithLongVKey;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
+import javax.persistence.AttributeOverride;
+import javax.persistence.AttributeOverrides;
+import javax.persistence.Column;
+import javax.persistence.Embedded;
+import javax.persistence.EnumType;
+import javax.persistence.Enumerated;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.MappedSuperclass;
+import javax.persistence.Transient;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 
 /** A billable event in a domain's lifecycle. */
+@MappedSuperclass
+@WithLongVKey
 public abstract class BillingEvent extends ImmutableObject
     implements Buildable, TransferServerApproveEntity {
 
@@ -93,24 +108,41 @@ public abstract class BillingEvent extends ImmutableObject
 
   /** Entity id. */
   @Id
-  long id;
+  @javax.persistence.Id
+  @GeneratedValue(strategy = GenerationType.IDENTITY)
+  Long id;
 
-  @Parent
-  @DoNotHydrate
-  Key<HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<HistoryEntry> parent;
 
   /** The registrar to bill. */
   @Index
+  @Column(nullable = false)
   String clientId;
 
+  /** Revision id of the entry in DomainHistory table that ths bill belongs to. */
+  // TODO(shicong): Add foreign key constraint when DomainHistory table is generated
+  @Ignore
+  @Column(nullable = false)
+  Long domainHistoryRevisionId;
+
+  /** ID of the EPP resource that the bill is for. */
+  // TODO(shicong): Add foreign key constraint when we expand DatastoreHelp for Postgresql
+  @Ignore
+  @Column(nullable = false)
+  String domainRepoId;
+
   /** When this event was created. For recurring events, this is also the recurrence start time. */
   @Index
+  @Column(nullable = false)
   DateTime eventTime;
 
   /** The reason for the bill. */
+  @Enumerated(EnumType.STRING)
+  @Column(nullable = false)
   Reason reason;
 
   /** The fully qualified domain name of the domain that the bill is for. */
+  @Column(name = "domain_name", nullable = false)
   String targetId;
 
   @Nullable
@@ -120,6 +152,14 @@ public abstract class BillingEvent extends ImmutableObject
     return clientId;
   }
 
+  public long getDomainHistoryRevisionId() {
+    return domainHistoryRevisionId;
+  }
+
+  public String getDomainRepoId() {
+    return domainRepoId;
+  }
+
   public DateTime getEventTime() {
     return eventTime;
   }
@@ -163,7 +203,7 @@ public abstract class BillingEvent extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setId(Long id) {
+    public B setId(long id) {
       getInstance().id = id;
       return thisCastToDerived();
     }
@@ -173,6 +213,16 @@ public abstract class BillingEvent extends ImmutableObject
       return thisCastToDerived();
     }
 
+    public B setDomainHistoryRevisionId(long domainHistoryRevisionId) {
+      getInstance().domainHistoryRevisionId = domainHistoryRevisionId;
+      return thisCastToDerived();
+    }
+
+    public B setDomainRepoId(String domainRepoId) {
+      getInstance().domainRepoId = domainRepoId;
+      return thisCastToDerived();
+    }
+
     public B setEventTime(DateTime eventTime) {
       getInstance().eventTime = eventTime;
       return thisCastToDerived();
@@ -194,6 +244,7 @@ public abstract class BillingEvent extends ImmutableObject
     }
 
     public B setParent(Key<HistoryEntry> parentKey) {
+      // TODO(shicong): Figure out how to set domainHistoryRevisionId and domainRepoId
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }
@@ -213,9 +264,23 @@ public abstract class BillingEvent extends ImmutableObject
   /** A one-time billable event. */
   @ReportedOn
   @Entity
+  @javax.persistence.Entity(name = "BillingEvent")
+  @javax.persistence.Table(
+      indexes = {
+        @javax.persistence.Index(columnList = "clientId"),
+        @javax.persistence.Index(columnList = "eventTime"),
+        @javax.persistence.Index(columnList = "billingTime"),
+        @javax.persistence.Index(columnList = "syntheticCreationTime"),
+        @javax.persistence.Index(columnList = "allocation_token_id")
+      })
+  @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
   public static class OneTime extends BillingEvent {
 
     /** The billable value. */
+    @AttributeOverrides({
+      @AttributeOverride(name = "money.amount", column = @Column(name = "cost_amount")),
+      @AttributeOverride(name = "money.currency", column = @Column(name = "cost_currency"))
+    })
     Money cost;
 
     /** When the cost should be billed. */
@@ -223,8 +288,8 @@ public abstract class BillingEvent extends ImmutableObject
     DateTime billingTime;
 
     /**
-     * The period in years of the action being billed for, if applicable, otherwise null.
-     * Used for financial reporting.
+     * The period in years of the action being billed for, if applicable, otherwise null. Used for
+     * financial reporting.
      */
     @IgnoreSave(IfNull.class)
     Integer periodYears = null;
@@ -240,15 +305,21 @@ public abstract class BillingEvent extends ImmutableObject
 
     /**
      * For {@link Flag#SYNTHETIC} events, a {@link Key} to the {@link BillingEvent} from which this
-     * OneTime was created. This is needed in order to properly match billing events against
-     * {@link Cancellation}s.
+     * OneTime was created. This is needed in order to properly match billing events against {@link
+     * Cancellation}s.
      */
-    Key<? extends BillingEvent> cancellationMatchingBillingEvent;
+    @Column(name = "cancellation_matching_billing_recurrence_id")
+    VKey<? extends BillingEvent> cancellationMatchingBillingEvent;
 
     /**
      * The {@link AllocationToken} used in the creation of this event, or null if one was not used.
+     *
+     * <p>TODO(shicong): Add foreign key constraint when AllocationToken schema is generated
      */
-    @Index @Nullable Key<AllocationToken> allocationToken;
+    @Column(name = "allocation_token_id")
+    @Index
+    @Nullable
+    VKey<AllocationToken> allocationToken;
 
     public Money getCost() {
       return cost;
@@ -266,14 +337,18 @@ public abstract class BillingEvent extends ImmutableObject
       return syntheticCreationTime;
     }
 
-    public Key<? extends BillingEvent> getCancellationMatchingBillingEvent() {
+    public VKey<? extends BillingEvent> getCancellationMatchingBillingEvent() {
       return cancellationMatchingBillingEvent;
     }
 
-    public Optional<Key<AllocationToken>> getAllocationToken() {
+    public Optional<VKey<AllocationToken>> getAllocationToken() {
       return Optional.ofNullable(allocationToken);
     }
 
+    public VKey<OneTime> createVKey() {
+      return VKey.createOfy(getClass(), Key.create(this));
+    }
+
     @Override
     public Builder asBuilder() {
       return new Builder(clone(this));
@@ -311,12 +386,12 @@ public abstract class BillingEvent extends ImmutableObject
       }
 
       public Builder setCancellationMatchingBillingEvent(
-          Key<? extends BillingEvent> cancellationMatchingBillingEvent) {
+          VKey<? extends BillingEvent> cancellationMatchingBillingEvent) {
         getInstance().cancellationMatchingBillingEvent = cancellationMatchingBillingEvent;
         return this;
       }
 
-      public Builder setAllocationToken(@Nullable Key<AllocationToken> allocationToken) {
+      public Builder setAllocationToken(@Nullable VKey<AllocationToken> allocationToken) {
         getInstance().allocationToken = allocationToken;
         return this;
       }
@@ -361,6 +436,15 @@ public abstract class BillingEvent extends ImmutableObject
    */
   @ReportedOn
   @Entity
+  @javax.persistence.Entity(name = "BillingRecurrence")
+  @javax.persistence.Table(
+      indexes = {
+        @javax.persistence.Index(columnList = "clientId"),
+        @javax.persistence.Index(columnList = "eventTime"),
+        @javax.persistence.Index(columnList = "recurrenceEndTime"),
+        @javax.persistence.Index(columnList = "recurrence_time_of_year")
+      })
+  @AttributeOverride(name = "id", column = @Column(name = "billing_recurrence_id"))
   public static class Recurring extends BillingEvent {
 
     /**
@@ -384,6 +468,10 @@ public abstract class BillingEvent extends ImmutableObject
      * model, whereas the billing time is a fixed {@link org.joda.time.Duration} later.
      */
     @Index
+    @Embedded
+    @AttributeOverrides({
+      @AttributeOverride(name = "timeString", column = @Column(name = "recurrence_time_of_year"))
+    })
     TimeOfYear recurrenceTimeOfYear;
 
     public DateTime getRecurrenceEndTime() {
@@ -394,6 +482,10 @@ public abstract class BillingEvent extends ImmutableObject
       return recurrenceTimeOfYear;
     }
 
+    public VKey<Recurring> createVKey() {
+      return VKey.createOfy(getClass(), Key.create(this));
+    }
+
     @Override
     public Builder asBuilder() {
       return new Builder(clone(this));
@@ -434,6 +526,14 @@ public abstract class BillingEvent extends ImmutableObject
    */
   @ReportedOn
   @Entity
+  @javax.persistence.Entity(name = "BillingCancellation")
+  @javax.persistence.Table(
+      indexes = {
+        @javax.persistence.Index(columnList = "clientId"),
+        @javax.persistence.Index(columnList = "eventTime"),
+        @javax.persistence.Index(columnList = "billingTime")
+      })
+  @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
   public static class Cancellation extends BillingEvent {
 
     /** The billing time of the charge that is being cancelled. */
@@ -446,7 +546,8 @@ public abstract class BillingEvent extends ImmutableObject
      * <p>Although the type is {@link Key} the name "ref" is preserved for historical reasons.
      */
     @IgnoreSave(IfNull.class)
-    Key<BillingEvent.OneTime> refOneTime = null;
+    @Column(name = "billing_event_id")
+    VKey<BillingEvent.OneTime> refOneTime = null;
 
     /**
      * The recurring billing event to cancel, or null for non-autorenew cancellations.
@@ -454,13 +555,14 @@ public abstract class BillingEvent extends ImmutableObject
      * <p>Although the type is {@link Key} the name "ref" is preserved for historical reasons.
      */
     @IgnoreSave(IfNull.class)
-    Key<BillingEvent.Recurring> refRecurring = null;
+    @Column(name = "billing_recurrence_id")
+    VKey<BillingEvent.Recurring> refRecurring = null;
 
     public DateTime getBillingTime() {
       return billingTime;
     }
 
-    public Key<? extends BillingEvent> getEventKey() {
+    public VKey<? extends BillingEvent> getEventKey() {
       return firstNonNull(refOneTime, refRecurring);
     }
 
@@ -492,13 +594,19 @@ public abstract class BillingEvent extends ImmutableObject
           .setParent(historyEntry);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
-        builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());
+        builder.setOneTimeEventKey(
+            VKey.createOfy(BillingEvent.OneTime.class, gracePeriod.getOneTimeBillingEvent()));
       } else if (gracePeriod.getRecurringBillingEvent() != null) {
-        builder.setRecurringEventKey(gracePeriod.getRecurringBillingEvent());
+        builder.setRecurringEventKey(
+            VKey.createOfy(BillingEvent.Recurring.class, gracePeriod.getRecurringBillingEvent()));
       }
       return builder.build();
     }
 
+    public VKey<Cancellation> createVKey() {
+      return VKey.createOfy(getClass(), Key.create(this));
+    }
+
     @Override
     public Builder asBuilder() {
       return new Builder(clone(this));
@@ -518,12 +626,12 @@ public abstract class BillingEvent extends ImmutableObject
         return this;
       }
 
-      public Builder setOneTimeEventKey(Key<BillingEvent.OneTime> eventKey) {
+      public Builder setOneTimeEventKey(VKey<BillingEvent.OneTime> eventKey) {
         getInstance().refOneTime = eventKey;
         return this;
       }
 
-      public Builder setRecurringEventKey(Key<BillingEvent.Recurring> eventKey) {
+      public Builder setRecurringEventKey(VKey<BillingEvent.Recurring> eventKey) {
         getInstance().refRecurring = eventKey;
         return this;
       }
@@ -540,9 +648,7 @@ public abstract class BillingEvent extends ImmutableObject
     }
   }
 
-  /**
-   * An event representing a modification of an existing one-time billing event.
-   */
+  /** An event representing a modification of an existing one-time billing event. */
   @ReportedOn
   @Entity
   public static class Modification extends BillingEvent {

core/src/main/java/google/registry/model/common/TimeOfYear.java

@@ -29,20 +29,22 @@ import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Index;
 import google.registry.model.ImmutableObject;
 import java.util.List;
+import javax.persistence.Embeddable;
 import org.joda.time.DateTime;
 
 /**
  * A time of year (month, day, millis of day) that can be stored in a sort-friendly format.
  *
- * <p>This is conceptually similar to {@code MonthDay} in Joda or more generally to Joda's
- * {@code Partial}, but the parts we need are too simple to justify a full implementation of
- * {@code Partial}.
+ * <p>This is conceptually similar to {@code MonthDay} in Joda or more generally to Joda's {@code
+ * Partial}, but the parts we need are too simple to justify a full implementation of {@code
+ * Partial}.
  *
  * <p>For simplicity, the native representation of this class's data is its stored format. This
  * allows it to be embeddable with no translation needed and also delays parsing of the string on
  * load until it's actually needed.
  */
 @Embed
+@Embeddable
 public class TimeOfYear extends ImmutableObject {
 
   /**

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

@@ -37,7 +37,7 @@ import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
 /** A collection of {@link ContactResource} commands. */
 public class ContactCommand {
 
-  /** The fields on "chgType" from {@link "http://tools.ietf.org/html/rfc5733"}. */
+  /** The fields on "chgType" from <a href="http://tools.ietf.org/html/rfc5733">RFC5733</a>. */
   @XmlTransient
   public static class ContactCreateOrChange extends ImmutableObject
       implements ResourceCreateOrChange<ContactResource.Builder> {
@@ -111,8 +111,8 @@ public class ContactCommand {
   }
 
   /**
-   * A create command for a {@link ContactResource}, mapping "createType" from
-   * {@link "http://tools.ietf.org/html/rfc5733"}.
+   * A create command for a {@link ContactResource}, mapping "createType" from <a
+   * href="http://tools.ietf.org/html/rfc5733">RFC5733</a>}.
    */
   @XmlType(propOrder = {"contactId", "postalInfo", "voice", "fax", "email", "authInfo", "disclose"})
   @XmlRootElement

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

@@ -27,7 +27,7 @@ import javax.persistence.Embedded;
 import javax.xml.bind.annotation.XmlAttribute;
 import javax.xml.bind.annotation.XmlType;
 
-/** The "discloseType" from {@link "http://tools.ietf.org/html/rfc5733"}. */
+/** The "discloseType" from <a href="http://tools.ietf.org/html/rfc5733">RFC5733</a>. */
 @Embed
 @Embeddable
 @XmlType(propOrder = {"name", "org", "addr", "voice", "fax", "email"})
@@ -76,7 +76,7 @@ public class Disclose extends ImmutableObject {
     return flag;
   }
 
-  /** The "intLocType" from {@link "http://tools.ietf.org/html/rfc5733"}. */
+  /** The "intLocType" from <a href="http://tools.ietf.org/html/rfc5733">RFC5733</a>. */
   @Embed
   public static class PostalInfoChoice extends ImmutableObject {
     @XmlAttribute

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

@@ -32,8 +32,8 @@ import javax.xml.bind.annotation.adapters.NormalizedStringAdapter;
 import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
 
 /**
- * Implementation of both "postalInfoType" and "chgPostalInfoType" from {@link
- * "http://tools.ietf.org/html/rfc5733"}.
+ * Implementation of both "postalInfoType" and "chgPostalInfoType" from <a href=
+ * "http://tools.ietf.org/html/rfc5733">RFC5733</a>.
  */
 @Embed
 @Embeddable

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

@@ -136,18 +136,11 @@ public class DomainBase extends EppResource
   @Index
   String tld;
 
-  /**
-   * References to hosts that are the nameservers for the domain.
-   *
-   * <p>This is a legacy field: we have to preserve it because it is still persisted and indexed in
-   * the datastore, but all external references go through nsHostVKeys.
-   */
-  @Index @ElementCollection @Transient Set<Key<HostResource>> nsHosts;
-
-  @Ignore
+  /** References to hosts that are the nameservers for the domain. */
+  @Index
   @ElementCollection
   @JoinTable(name = "DomainHost")
-  Set<VKey<HostResource>> nsHostVKeys;
+  Set<VKey<HostResource>> nsHosts;
 
   /**
    * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.
@@ -269,11 +262,6 @@ public class DomainBase extends EppResource
 
   @OnLoad
   void load() {
-    nsHostVKeys =
-        nullToEmptyImmutableCopy(nsHosts).stream()
-            .map(hostKey -> VKey.createOfy(HostResource.class, hostKey))
-            .collect(toImmutableSet());
-
     // Reconstitute all of the contacts so that they have VKeys.
     allContacts =
         allContacts.stream().map(contact -> contact.reconstitute()).collect(toImmutableSet());
@@ -363,9 +351,7 @@ public class DomainBase extends EppResource
   }
 
   public ImmutableSet<VKey<HostResource>> getNameservers() {
-    // Since nsHostVKeys gets initialized both from setNameservers() and the OnLoad method, this
-    // should always be valid.
-    return nullToEmptyImmutableCopy(nsHostVKeys);
+    return nullToEmptyImmutableCopy(nsHosts);
   }
 
   public final String getCurrentSponsorClientId() {
@@ -645,14 +631,6 @@ public class DomainBase extends EppResource
 
     Builder(DomainBase instance) {
       super(instance);
-
-      // Convert nsHosts to nsHostVKeys.
-      if (instance.nsHosts != null) {
-        instance.nsHostVKeys =
-            instance.nsHosts.stream()
-                .map(key -> VKey.createOfy(HostResource.class, key))
-                .collect(toImmutableSet());
-      }
     }
 
     @Override
@@ -710,27 +688,12 @@ public class DomainBase extends EppResource
     }
 
     public Builder setNameservers(VKey<HostResource> nameserver) {
-      Optional<Key<HostResource>> nsKey = nameserver.maybeGetOfyKey();
-      if (nsKey.isPresent()) {
-        getInstance().nsHosts = ImmutableSet.of(nsKey.get());
-      } else {
-        getInstance().nsHosts = null;
-      }
-      getInstance().nsHostVKeys = ImmutableSet.of(nameserver);
+      getInstance().nsHosts = ImmutableSet.of(nameserver);
       return thisCastToDerived();
     }
 
     public Builder setNameservers(ImmutableSet<VKey<HostResource>> nameservers) {
-      // If we have all of the ofy keys, we can set nsHosts.  Otherwise, make it null.
-      if (nameservers != null
-          && nameservers.stream().allMatch(key -> key.maybeGetOfyKey().isPresent())) {
-        getInstance().nsHosts =
-            nameservers.stream().map(key -> key.getOfyKey()).collect(toImmutableSet());
-      } else {
-        getInstance().nsHosts = null;
-      }
-
-      getInstance().nsHostVKeys = forceEmptyToNull(nameservers);
+      getInstance().nsHosts = forceEmptyToNull(nameservers);
       return thisCastToDerived();
     }
 

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

@@ -72,10 +72,10 @@ public class DomainCommand {
     T cloneAndLinkReferences(DateTime now) throws InvalidReferencesException;
   }
 
-  /** The fields on "chgType" from {@link "http://tools.ietf.org/html/rfc5731"}. */
+  /** The fields on "chgType" from <a href="http://tools.ietf.org/html/rfc5731">RFC5731</a>. */
   @XmlTransient
-  public static class DomainCreateOrChange<B extends DomainBase.Builder>
-      extends ImmutableObject implements ResourceCreateOrChange<B> {
+  public static class DomainCreateOrChange<B extends DomainBase.Builder> extends ImmutableObject
+      implements ResourceCreateOrChange<B> {
 
     /** The contactId of the registrant who registered this domain. */
     @XmlElement(name = "registrant")
@@ -103,19 +103,20 @@ public class DomainCommand {
   }
 
   /**
-   * A create command for a {@link DomainBase}, mapping "createType" from
-   * {@link "http://tools.ietf.org/html/rfc5731"}.
+   * A create command for a {@link DomainBase}, mapping "createType" from <a
+   * href="http://tools.ietf.org/html/rfc5731">RFC5731</a>.
    */
   @XmlRootElement
-  @XmlType(propOrder = {
-      "fullyQualifiedDomainName",
-      "period",
-      "nameserverFullyQualifiedHostNames",
-      "registrantContactId",
-      "foreignKeyedDesignatedContacts",
-      "authInfo"})
-  public static class Create
-      extends DomainCreateOrChange<DomainBase.Builder>
+  @XmlType(
+      propOrder = {
+        "fullyQualifiedDomainName",
+        "period",
+        "nameserverFullyQualifiedHostNames",
+        "registrantContactId",
+        "foreignKeyedDesignatedContacts",
+        "authInfo"
+      })
+  public static class Create extends DomainCreateOrChange<DomainBase.Builder>
       implements CreateOrUpdate<Create> {
 
     /** Fully qualified domain name, which serves as a unique identifier for this domain. */

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

@@ -20,7 +20,7 @@ import javax.xml.bind.annotation.XmlAttribute;
 import javax.xml.bind.annotation.XmlEnumValue;
 import javax.xml.bind.annotation.XmlValue;
 
-/** The "periodType" from {@link "http://tools.ietf.org/html/rfc5731"}. */
+/** The "periodType" from <a href="http://tools.ietf.org/html/rfc5731">RFC5731</a>. */
 @Embed
 @javax.persistence.Embeddable
 public class Period extends ImmutableObject {

core/src/main/java/google/registry/model/domain/fee12/FeeCheckCommandExtensionItemV12.java

@@ -29,11 +29,13 @@ import org.joda.time.DateTime;
  * An individual price check item in version 0.12 of the fee extension on domain check commands.
  * Items look like:
  *
+ * <pre>{@code
  * <fee:command name="renew" phase="sunrise" subphase="hello">
  *   <fee:period unit="y">1</fee:period>
  *   <fee:class>premium</fee:class>
  *   <fee:date>2017-05-17T13:22:21.0Z</fee:date>
  * </fee:command>
+ * }</pre>
  *
  * In a change from previous versions of the extension, items do not contain domain names; instead,
  * the names from the non-extension check element are used.

core/src/main/java/google/registry/model/domain/launch/LaunchNotice.java

@@ -95,9 +95,6 @@ public class LaunchNotice extends ImmutableObject {
 
   /**
    * Validate the checksum of the notice against the domain label.
-   *
-   * @throws IllegalArgumentException
-   * @throws InvalidChecksumException
    */
   public void validate(String domainLabel) throws InvalidChecksumException {
     // According to http://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.3, a TCNID

core/src/main/java/google/registry/model/domain/rgp/GracePeriodStatus.java

@@ -49,9 +49,9 @@ public enum GracePeriodStatus implements EppEnum {
   AUTO_RENEW("autoRenewPeriod"),
 
   /**
-   * This status value is used to describe a domain for which a <delete> command has been received,
-   * but the domain has not yet been purged because an opportunity exists to restore the domain and
-   * abort the deletion process.
+   * This status value is used to describe a domain for which a &lt;delete&gt; command has been
+   * received, but the domain has not yet been purged because an opportunity exists to restore the
+   * domain and abort the deletion process.
    */
   REDEMPTION("redemptionPeriod"),
 

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

@@ -45,6 +45,8 @@ import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.common.TimedTransitionProperty.TimeMapper;
 import google.registry.model.common.TimedTransitionProperty.TimedTransition;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import google.registry.persistence.WithStringVKey;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -53,6 +55,7 @@ import org.joda.time.DateTime;
 /** An entity representing an allocation token. */
 @ReportedOn
 @Entity
+@WithStringVKey
 public class AllocationToken extends BackupGroupRoot implements Buildable {
 
   // Promotions should only move forward, and ENDED / CANCELLED are terminal states.
@@ -179,6 +182,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return tokenStatusTransitions;
   }
 
+  public VKey<AllocationToken> createVKey() {
+    return VKey.createOfy(getClass(), Key.create(this));
+  }
+
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));

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

@@ -25,7 +25,7 @@ import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
  * An allocation token extension that may be present on EPP domain commands.
  *
  * @see <a href="https://tools.ietf.org/html/draft-ietf-regext-allocation-token-04">the IETF
- *     draft</a> for full details.
+ *     draft</a>
  */
 @XmlRootElement(name = "allocationToken")
 public class AllocationTokenExtension extends ImmutableObject implements CommandExtension {

core/src/main/java/google/registry/model/eppcommon/Address.java

@@ -44,8 +44,9 @@ import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
 /**
  * Container for generic street address.
  *
- * <p>This is the "addrType" type from {@link "http://tools.ietf.org/html/rfc5733"}. It also matches
- * the "addrType" type from {@link "http://tools.ietf.org/html/draft-lozano-tmch-smd"}.
+ * <p>This is the "addrType" type from <a href="http://tools.ietf.org/html/rfc5733">RFC5733</a>. It
+ * also matches the "addrType" type from <a
+ * href="http://tools.ietf.org/html/draft-lozano-tmch-smd">Mark and Signed Mark Objects Mapping</a>.
  *
  * @see google.registry.model.contact.ContactAddress
  * @see google.registry.model.mark.MarkAddress

core/src/main/java/google/registry/model/eppcommon/PhoneNumber.java

@@ -29,8 +29,9 @@ import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
 /**
  * Container for generic E164 phone number.
  *
- * <p>This is the "e164" type from {@link "http://tools.ietf.org/html/rfc5733"}. It also matches the
- * "e164Type" type from {@link "http://tools.ietf.org/html/draft-lozano-tmch-smd"}.
+ * <p>This is the "e164" type from <a href="http://tools.ietf.org/html/rfc5733">RFC5733</a>. It also
+ * matches the "e164Type" type from <a href="http://tools.ietf.org/html/draft-lozano-tmch-smd">Mark
+ * and Signed Mark Objects Mapping</a>
  *
  * <blockquote>
  *

core/src/main/java/google/registry/model/eppoutput/Greeting.java

@@ -24,7 +24,7 @@ import javax.xml.bind.annotation.XmlElementWrapper;
 import org.joda.time.DateTime;
 
 /**
- * A greeting, defined in {@link "http://tools.ietf.org/html/rfc5730"}.
+ * A greeting, defined in <a href="http://tools.ietf.org/html/rfc5730">RFC5730</a>.
  *
  * <p>It would be nice to make this a singleton, but we need the {@link #svDate} field to stay
  * current.

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

@@ -32,7 +32,7 @@ import javax.xml.bind.annotation.XmlType;
 /** A collection of {@link HostResource} commands. */
 public class HostCommand {
 
-  /** The fields on "chgType" from {@link "http://tools.ietf.org/html/rfc5732"}. */
+  /** The fields on "chgType" from <a href="http://tools.ietf.org/html/rfc5732">RFC5732</a>. */
   @XmlTransient
   abstract static class HostCreateOrChange extends AbstractSingleResourceCommand
       implements ResourceCreateOrChange<HostResource.Builder> {
@@ -42,13 +42,13 @@ public class HostCommand {
   }
 
   /**
-   * A create command for a {@link HostResource}, mapping "createType" from
-   * {@link "http://tools.ietf.org/html/rfc5732"}.
+   * A create command for a {@link HostResource}, mapping "createType" from <a
+   * href="http://tools.ietf.org/html/rfc5732">RFC5732</a>.
    */
-  @XmlType(propOrder = {"targetId", "inetAddresses" })
+  @XmlType(propOrder = {"targetId", "inetAddresses"})
   @XmlRootElement
-  public static class Create
-      extends HostCreateOrChange implements ResourceCreateOrChange<HostResource.Builder> {
+  public static class Create extends HostCreateOrChange
+      implements ResourceCreateOrChange<HostResource.Builder> {
     /** IP Addresses for this host. Can be null if this is an external host. */
     @XmlElement(name = "addr")
     Set<InetAddress> inetAddresses;

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

@@ -34,10 +34,9 @@ import com.googlecode.objectify.annotation.EntitySubclass;
 import com.googlecode.objectify.impl.translate.TranslatorFactory;
 import com.googlecode.objectify.impl.translate.opt.joda.MoneyStringTranslatorFactory;
 import google.registry.config.RegistryEnvironment;
+import google.registry.model.Buildable;
 import google.registry.model.EntityClasses;
 import google.registry.model.ImmutableObject;
-import google.registry.model.contact.ContactResource;
-import google.registry.model.host.HostResource;
 import google.registry.model.translators.BloomFilterOfStringTranslatorFactory;
 import google.registry.model.translators.CidrAddressBlockTranslatorFactory;
 import google.registry.model.translators.CommitLogRevisionsTranslatorFactory;
@@ -131,8 +130,7 @@ public class ObjectifyService {
             new InetAddressTranslatorFactory(),
             new MoneyStringTranslatorFactory(),
             new ReadableInstantUtcTranslatorFactory(),
-            new VKeyTranslatorFactory<ContactResource>(ContactResource.class),
-            new VKeyTranslatorFactory<HostResource>(HostResource.class),
+            new VKeyTranslatorFactory(),
             new UpdateAutoTimestampTranslatorFactory())) {
       factory().getTranslators().add(translatorFactory);
     }
@@ -170,10 +168,16 @@ public class ObjectifyService {
       }
       com.googlecode.objectify.ObjectifyService.register(clazz);
       // Autogenerated ids make the commit log code very difficult since we won't always be able
-      // to create a key for an entity immediately when requesting a save. Disallow that here.
-      checkState(
-          !factory().getMetadata(clazz).getKeyMetadata().isIdGeneratable(),
-          "Can't register %s: Autogenerated ids (@Id on a Long) are not supported.", kind);
+      // to create a key for an entity immediately when requesting a save. So, we require such
+      // entities to implement google.registry.model.Buildable as its build() function allocates the
+      // id to the entity.
+      if (factory().getMetadata(clazz).getKeyMetadata().isIdGeneratable()) {
+        checkState(
+            Buildable.class.isAssignableFrom(clazz),
+            "Can't register %s: Entity with autogenerated ids (@Id on a Long) must implement"
+                + " google.registry.model.Buildable.",
+            kind);
+      }
     }
   }
 

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

@@ -62,7 +62,7 @@ import org.joda.time.DateTime;
  * <p>Poll messages are identified externally by registrars using the format defined in {@link
  * PollMessageExternalKeyConverter}.
  *
- * @see <a href="https://tools.ietf.org/html/rfc5730#section-2.9.2.3">RFC5730 - EPP - &ltpoll&gt
+ * @see <a href="https://tools.ietf.org/html/rfc5730#section-2.9.2.3">RFC5730 - EPP - &lt;poll&gt;
  *     Command</a>
  */
 @Entity

core/src/main/java/google/registry/model/registry/label/ReservedList.java

@@ -158,8 +158,8 @@ public final class ReservedList
   /**
    * Gets a ReservedList by name using the caching layer.
    *
-   * @return An Optional<ReservedList> that has a value if a reserved list exists by the given name,
-   *     or absent if not.
+   * @return An Optional&lt;ReservedList&gt; that has a value if a reserved list exists by the given
+   *     name, or absent if not.
    * @throws UncheckedExecutionException if some other error occurs while trying to load the
    *     ReservedList from the cache or Datastore.
    */

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

@@ -22,8 +22,8 @@ public final class IcannReportingTypes {
    * Represents the set of possible ICANN Monthly Registry Functions Activity Report fields.
    *
    * <p>Refer to the <a
-   * href="https://newgtlds.icann.org/sites/default/files/agreements/agreement-approved-09jan14-en.htm#_DV_M278>ICANN
-   * registry agreement Specification 3 Section 2</a> for details.
+   * href="https://newgtlds.icann.org/sites/default/files/agreements/agreement-approved-09jan14-en.htm#_DV_M278">
+   * ICANN registry agreement Specification 3 Section 2</a> for details.
    */
   public enum ActivityReportField {
     DOMAIN_CHECK("srs-dom-check"),

core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java

@@ -14,13 +14,14 @@
 
 package google.registry.model.translators;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
+import static com.google.common.base.Functions.identity;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.model.EntityClasses.ALL_CLASSES;
 
-import com.googlecode.objectify.Key;
+import com.google.appengine.api.datastore.Key;
+import com.google.common.collect.ImmutableMap;
+import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.persistence.VKey;
-import java.io.UnsupportedEncodingException;
-import java.net.URLDecoder;
-import java.net.URLEncoder;
 
 /**
  * Translator factory for VKey.
@@ -28,57 +29,39 @@ import java.net.URLEncoder;
  * <p>These get translated to a string containing the URL safe encoding of the objectify key
  * followed by a (url-unsafe) ampersand delimiter and the SQL key.
  */
-public class VKeyTranslatorFactory<T> extends AbstractSimpleTranslatorFactory<VKey, String> {
-  private final Class<T> refClass;
+public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey, Key> {
 
-  public VKeyTranslatorFactory(Class<T> refClass) {
+  // Class registry allowing us to restore the original class object from the unqualified class
+  // name, which is all the datastore key gives us.
+  // Note that entities annotated with @EntitySubclass are removed because they share the same
+  // kind of the key with their parent class.
+  private static final ImmutableMap<String, Class> CLASS_REGISTRY =
+      ALL_CLASSES.stream()
+          .filter(clazz -> !clazz.isAnnotationPresent(EntitySubclass.class))
+          .collect(toImmutableMap(com.googlecode.objectify.Key::getKind, identity()));
+  ;
+
+  public VKeyTranslatorFactory() {
     super(VKey.class);
-    this.refClass = refClass;
   }
 
   @Override
-  public SimpleTranslator<VKey, String> createTranslator() {
-    return new SimpleTranslator<VKey, String>() {
+  public SimpleTranslator<VKey, Key> createTranslator() {
+    return new SimpleTranslator<VKey, Key>() {
       @Override
-      public VKey loadValue(String datastoreValue) {
-        int pos = datastoreValue.indexOf('&');
-        Key ofyKey = null;
-        String sqlKey = null;
-        if (pos > 0) {
-          // We have an objectify key.
-          ofyKey = Key.create(datastoreValue.substring(0, pos));
-        }
-
-        if (pos < datastoreValue.length() - 1) {
-          // We have an SQL key.
-          sqlKey = decode(datastoreValue.substring(pos + 1));
-        }
-
-        return VKey.create(refClass, sqlKey, ofyKey);
+      public VKey loadValue(Key datastoreValue) {
+        // TODO(mmuller): we need to call a method on refClass to also reconstitute the SQL key.
+        return datastoreValue == null
+            ? null
+            : VKey.createOfy(
+                CLASS_REGISTRY.get(datastoreValue.getKind()),
+                com.googlecode.objectify.Key.create(datastoreValue));
       }
 
       @Override
-      public String saveValue(VKey key) {
-        return ((key.getOfyKey() == null) ? "" : key.getOfyKey().getString())
-            + "&"
-            + ((key.getSqlKey() == null) ? "" : encode(key.getSqlKey().toString()));
+      public Key saveValue(VKey key) {
+        return key == null ? null : key.getOfyKey().getRaw();
       }
     };
   }
-
-  private static String encode(String val) {
-    try {
-      return URLEncoder.encode(val, UTF_8.toString());
-    } catch (UnsupportedEncodingException e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  private static String decode(String encoded) {
-    try {
-      return URLDecoder.decode(encoded, UTF_8.toString());
-    } catch (UnsupportedEncodingException e) {
-      throw new RuntimeException(e);
-    }
-  }
 }

core/src/main/java/google/registry/persistence/converter/BillingCostTransitionConverter.java

@@ -15,7 +15,6 @@
 package google.registry.persistence.converter;
 
 import avro.shaded.com.google.common.collect.Maps;
-import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.registry.Registry.BillingCostTransition;
 import java.util.Map;
 import javax.persistence.Converter;
@@ -23,8 +22,8 @@ import org.joda.money.Money;
 import org.joda.time.DateTime;
 
 /**
- * JPA converter for storing/retrieving {@link TimedTransitionProperty <Money, BillingCostTransition
- * >} objects.
+ * JPA converter for storing/retrieving {@code TimedTransitionProperty<Money,BillingCostTransition>}
+ * objects.
  */
 @Converter(autoApply = true)
 public class BillingCostTransitionConverter

core/src/main/java/google/registry/persistence/converter/BillingEventFlagSetConverter.java

@@ -0,0 +1,35 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.converter;
+
+import google.registry.model.billing.BillingEvent.Flag;
+import java.util.Set;
+import javax.persistence.AttributeConverter;
+import javax.persistence.Converter;
+
+/** JPA {@link AttributeConverter} for storing/retrieving {@link Set}. */
+@Converter(autoApply = true)
+public class BillingEventFlagSetConverter extends StringSetConverterBase<Flag> {
+
+  @Override
+  String toString(Flag element) {
+    return element.name();
+  }
+
+  @Override
+  Flag fromString(String value) {
+    return Flag.valueOf(value);
+  }
+}

core/src/main/java/google/registry/persistence/converter/CidrAddressBlockListConverter.java

@@ -15,12 +15,11 @@
 package google.registry.persistence.converter;
 
 import google.registry.util.CidrAddressBlock;
-import java.util.List;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
 /**
- * JPA {@link AttributeConverter} for storing/retrieving {@link List<CidrAddressBlock>} objects.
+ * JPA {@link AttributeConverter} for storing/retrieving {@code List<CidrAddressBlock>} objects.
  * TODO(shicong): Investigate if we can have one converter for any List type
  */
 @Converter(autoApply = true)

core/src/main/java/google/registry/persistence/converter/CurrencyToBillingConverter.java

@@ -21,7 +21,7 @@ import java.util.Map;
 import javax.persistence.Converter;
 import org.joda.money.CurrencyUnit;
 
-/** JPA converter for storing/retrieving {@link Map <CurrencyUnit, BillingAccountEntry>} objects. */
+/** JPA converter for storing/retrieving {@code Map<CurrencyUnit, BillingAccountEntry>} objects. */
 @Converter(autoApply = true)
 public class CurrencyToBillingConverter
     extends StringMapConverterBase<CurrencyUnit, BillingAccountEntry> {

core/src/main/java/google/registry/persistence/converter/PostalInfoChoiceListConverter.java

@@ -16,11 +16,10 @@ package google.registry.persistence.converter;
 
 import google.registry.model.contact.Disclose.PostalInfoChoice;
 import google.registry.model.contact.PostalInfo;
-import java.util.List;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
-/** JPA {@link AttributeConverter} for storing/retrieving {@link List < PostalInfoChoice >}. */
+/** JPA {@link AttributeConverter} for storing/retrieving {@code List<PostalInfoChoice>}. */
 @Converter(autoApply = true)
 public class PostalInfoChoiceListConverter extends StringListConverterBase<PostalInfoChoice> {
 

core/src/main/java/google/registry/persistence/converter/RegistrarPocSetConverter.java

@@ -15,11 +15,10 @@
 package google.registry.persistence.converter;
 
 import google.registry.model.registrar.RegistrarContact.Type;
-import java.util.Set;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
-/** JPA {@link AttributeConverter} for storing/retrieving {@link Set<Type>}. */
+/** JPA {@link AttributeConverter} for storing/retrieving {@code Set<Type>}. */
 @Converter(autoApply = true)
 public class RegistrarPocSetConverter extends StringSetConverterBase<Type> {
   @Override

core/src/main/java/google/registry/persistence/converter/StatusValueSetConverter.java

@@ -15,11 +15,10 @@
 package google.registry.persistence.converter;
 
 import google.registry.model.eppcommon.StatusValue;
-import java.util.Set;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
-/** JPA {@link AttributeConverter} for storing/retrieving {@link Set<StatusValue>}. */
+/** JPA {@link AttributeConverter} for storing/retrieving {@code Set<StatusValue>}. */
 @Converter(autoApply = true)
 public class StatusValueSetConverter extends StringSetConverterBase<StatusValue> {
 

core/src/main/java/google/registry/persistence/converter/StringCollectionDescriptor.java

@@ -37,7 +37,7 @@ import org.hibernate.type.descriptor.sql.SqlTypeDescriptor;
 /**
  * The {@link JavaTypeDescriptor} and {@link SqlTypeDescriptor} for {@link StringCollection}.
  *
- * <p>A {@link StringCollection} object is a simple wrapper for a {@link Collection<String>} which
+ * <p>A {@link StringCollection} object is a simple wrapper for a {@code Collection<String>} which
  * can be stored as a string array in the database. The {@link JavaTypeDescriptor} and {@link
  * SqlTypeDescriptor} is used by JPA/Hibernate to map between the collection and {@link Array} which
  * is the actual type that JDBC uses to read from and write to the database.

core/src/main/java/google/registry/persistence/converter/StringListConverter.java

@@ -14,11 +14,10 @@
 
 package google.registry.persistence.converter;
 
-import java.util.List;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
-/** JPA {@link AttributeConverter} for storing/retrieving {@link List<String>}. */
+/** JPA {@link AttributeConverter} for storing/retrieving {@code List<String>}. */
 @Converter(autoApply = true)
 public class StringListConverter extends StringListConverterBase<String> {
 

core/src/main/java/google/registry/persistence/converter/StringMapDescriptor.java

@@ -36,7 +36,7 @@ import org.hibernate.type.descriptor.sql.SqlTypeDescriptor;
 /**
  * The {@link JavaTypeDescriptor} and {@link SqlTypeDescriptor} for {@link StringMap}.
  *
- * <p>A {@link StringMap} object is a simple wrapper for a {@link Map <String, String>} which can be
+ * <p>A {@link StringMap} object is a simple wrapper for a {@code Map<String, String>} which can be
  * stored in a column with data type of hstore in the database. The {@link JavaTypeDescriptor} and
  * {@link SqlTypeDescriptor} is used by JPA/Hibernate to map between the map and hstore which is the
  * actual type that JDBC uses to read from and write to the database.
@@ -156,7 +156,7 @@ public class StringMapDescriptor extends AbstractTypeDescriptor<StringMap>
     };
   }
 
-  /** A simple wrapper class for {@link Map<String, String>}. */
+  /** A simple wrapper class for {@code Map<String, String>}. */
   public static class StringMap {
     private Map<String, String> map;
 
@@ -169,7 +169,7 @@ public class StringMapDescriptor extends AbstractTypeDescriptor<StringMap>
       return new StringMap(ImmutableMap.copyOf(map));
     }
 
-    /** Returns the underlying {@link Map<String, String>} object. */
+    /** Returns the underlying {@code Map<String, String>} object. */
     public Map<String, String> getMap() {
       return map;
     }

core/src/main/java/google/registry/persistence/converter/StringSetConverter.java

@@ -14,11 +14,10 @@
 
 package google.registry.persistence.converter;
 
-import java.util.Set;
 import javax.persistence.AttributeConverter;
 import javax.persistence.Converter;
 
-/** JPA {@link AttributeConverter} for storing/retrieving {@link Set<String>}. */
+/** JPA {@link AttributeConverter} for storing/retrieving {@code Set<String>}. */
 @Converter(autoApply = true)
 public class StringSetConverter extends StringSetConverterBase<String> {
 

core/src/main/java/google/registry/persistence/converter/TldStateTransitionConverter.java

@@ -15,7 +15,6 @@
 package google.registry.persistence.converter;
 
 import com.google.common.collect.Maps;
-import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.Registry.TldStateTransition;
 import java.util.Map;
@@ -23,7 +22,7 @@ import javax.persistence.Converter;
 import org.joda.time.DateTime;
 
 /**
- * JPA converter for storing/retrieving {@link TimedTransitionProperty<TldState,
+ * JPA converter for storing/retrieving {@code TimedTransitionProperty<TldState,
  * TldStateTransition>} objects.
  */
 @Converter(autoApply = true)

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

@@ -17,6 +17,7 @@ package google.registry.persistence.transaction;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import google.registry.persistence.VKey;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;

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

@@ -16,6 +16,7 @@ package google.registry.persistence.transaction;
 
 import com.google.appengine.api.utils.SystemProperty;
 import com.google.appengine.api.utils.SystemProperty.Environment.Value;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.DaggerPersistenceComponent;
@@ -26,7 +27,9 @@ import java.util.function.Supplier;
 // TODO: Rename this to PersistenceFactory and move to persistence package.
 public class TransactionManagerFactory {
 
-  private static final TransactionManager TM = createTransactionManager();
+  private static final DatastoreTransactionManager ofyTm = createTransactionManager();
+
+  @NonFinalForTesting private static TransactionManager tm = ofyTm;
 
   /** Supplier for jpaTm so that it is initialized only once, upon first usage. */
   @NonFinalForTesting
@@ -45,10 +48,7 @@ public class TransactionManagerFactory {
     }
   }
 
-  private static TransactionManager createTransactionManager() {
-    // TODO: Determine how to provision TransactionManager after the dual-write. During the
-    // dual-write transitional phase, we need the TransactionManager for both Datastore and Cloud
-    // SQL, and this method returns the one for Datastore.
+  private static DatastoreTransactionManager createTransactionManager() {
     return new DatastoreTransactionManager(null);
   }
 
@@ -67,7 +67,7 @@ public class TransactionManagerFactory {
 
   /** Returns {@link TransactionManager} instance. */
   public static TransactionManager tm() {
-    return TM;
+    return tm;
   }
 
   /** Returns {@link JpaTransactionManager} instance. */
@@ -75,8 +75,20 @@ public class TransactionManagerFactory {
     return jpaTm.get();
   }
 
+  /** Returns {@link DatastoreTransactionManager} instance. */
+  @VisibleForTesting
+  public static DatastoreTransactionManager ofyTm() {
+    return ofyTm;
+  }
+
   /** Sets the return of {@link #jpaTm()} to the given instance of {@link JpaTransactionManager}. */
   public static void setJpaTm(JpaTransactionManager newJpaTm) {
     jpaTm = Suppliers.ofInstance(newJpaTm);
   }
+
+  /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */
+  @VisibleForTesting
+  public static void setTm(TransactionManager newTm) {
+    tm = newTm;
+  }
 }

core/src/main/java/google/registry/rdap/RdapDomainSearchAction.java

@@ -61,10 +61,9 @@ import javax.inject.Inject;
  *     (RDAP) Query Format</a>
  * @see <a href="http://tools.ietf.org/html/rfc7483">RFC 7483: JSON Responses for the Registration
  *     Data Access Protocol (RDAP)</a>
- *
- * TODO(guyben):This isn't required by the RDAP Technical Implementation Guide, and hence should be
- * deleted, at least until it's actually required.
  */
+// TODO: This isn't required by the RDAP Technical Implementation Guide, and hence should be
+// deleted, at least until it's actually required.
 @Action(
     service = Action.Service.PUBAPI,
     path = "/rdap/domains",

core/src/main/java/google/registry/rdap/RdapEntitySearchAction.java

@@ -72,10 +72,9 @@ import javax.inject.Inject;
  *     (RDAP) Query Format</a>
  * @see <a href="http://tools.ietf.org/html/rfc7483">RFC 7483: JSON Responses for the Registration
  *     Data Access Protocol (RDAP)</a>
- *
- * TODO(guyben):This isn't required by the RDAP Technical Implementation Guide, and hence should be
- * deleted, at least until it's actually required.
  */
+// TODO: This isn't required by the RDAP Technical Implementation Guide, and hence should be
+// deleted, at least until it's actually required.
 @Action(
     service = Action.Service.PUBAPI,
     path = "/rdap/entities",

core/src/main/java/google/registry/rdap/RdapSearchActionBase.java

@@ -221,10 +221,9 @@ public abstract class RdapSearchActionBase extends RdapActionBase {
    *
    * <p>This version handles a list of parameter values, all associated with the same name.
    *
-   * <p>Example: If the original parameters were "a=w&a=x&b=y&c=z", and this method is called with
-   * parameterName = "b" and parameterValues of "p" and "q", the result will be
-   * "a=w&a=x&c=z&b=p&b=q". The new values of parameter "b" replace the old ones.
-   *
+   * <p>Example: If the original parameters were "a=w&amp;a=x&amp;b=y&amp;c=z", and this method is
+   * called with parameterName = "b" and parameterValues of "p" and "q", the result will be
+   * "a=w&amp;a=x&amp;c=z&amp;b=p&amp;b=q". The new values of parameter "b" replace the old ones.
    */
   protected String getRequestUrlWithExtraParameter(
       String parameterName, List<String> parameterValues) {

core/src/main/java/google/registry/rde/Ghostryde.java

@@ -120,9 +120,6 @@ public final class Ghostryde {
 
   /**
    * Creates a ghostryde file from an in-memory byte array.
-   *
-   * @throws PGPException
-   * @throws IOException
    */
   public static byte[] encode(byte[] data, PGPPublicKey key)
       throws IOException, PGPException {
@@ -137,9 +134,6 @@ public final class Ghostryde {
 
   /**
    * Deciphers a ghostryde file from an in-memory byte array.
-   *
-   * @throws PGPException
-   * @throws IOException
    */
   public static byte[] decode(byte[] data, PGPPrivateKey key)
       throws IOException, PGPException {

core/src/main/java/google/registry/rde/RdeUtil.java

@@ -46,9 +46,6 @@ public final class RdeUtil {
   /**
    * Look at some bytes from {@code xmlInput} to ensure it appears to be a FULL XML deposit and
    * then use a regular expression to extract the watermark timestamp which is returned.
-   *
-   * @throws IOException
-   * @throws XmlException
    */
   public static DateTime peekWatermark(BufferedInputStream xmlInput)
       throws IOException, XmlException {

core/src/main/java/google/registry/reporting/billing/BillingEmailUtils.java

@@ -34,7 +34,7 @@ import javax.mail.internet.InternetAddress;
 import org.joda.time.YearMonth;
 
 /** Utility functions for sending emails involving monthly invoices. */
-class BillingEmailUtils {
+public class BillingEmailUtils {
 
   private final SendEmailService emailService;
   private final YearMonth yearMonth;

core/src/main/java/google/registry/reporting/billing/GenerateInvoicesAction.java

@@ -43,8 +43,8 @@ import org.joda.time.YearMonth;
  * PublishInvoicesAction} to publish the subsequent output.
  *
  * <p>This action runs the {@link google.registry.beam.invoicing.InvoicingPipeline} beam template,
- * staged at gs://<projectId>-beam/templates/invoicing. The pipeline then generates invoices for the
- * month and stores them on GCS.
+ * staged at gs://&lt;projectId&gt;-beam/templates/invoicing. The pipeline then generates invoices
+ * for the month and stores them on GCS.
  */
 @Action(
     service = Action.Service.BACKEND,

core/src/main/java/google/registry/reporting/spec11/PublishSpec11ReportAction.java

@@ -53,7 +53,7 @@ import org.json.JSONException;
  * pipeline accordingly.
  *
  * <p>This calls {@link Spec11EmailUtils#emailSpec11Reports(LocalDate, SoyTemplateInfo, String,
- * Set)} on success or {@link Spec11EmailUtils#sendAlertEmail(String, String)} on failure.
+ * ImmutableSet)} on success or {@link Spec11EmailUtils#sendAlertEmail(String, String)} on failure.
  */
 @Action(
     service = Action.Service.BACKEND,

core/src/main/java/google/registry/request/HttpException.java

@@ -18,6 +18,7 @@ import static com.google.common.html.HtmlEscapers.htmlEscaper;
 
 import com.google.common.flogger.FluentLogger;
 import java.io.IOException;
+import java.util.logging.Level;
 import javax.servlet.http.HttpServletResponse;
 
 /** Base for exceptions that cause an HTTP error response. */
@@ -28,11 +29,18 @@ public abstract class HttpException extends RuntimeException {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
+  private final Level logLevel;
+
   private final int responseCode;
 
-  protected HttpException(int responseCode, String message, Throwable cause) {
+  protected HttpException(int responseCode, String message, Throwable cause, Level logLevel) {
     super(message, cause);
     this.responseCode = responseCode;
+    this.logLevel = logLevel;
+  }
+
+  protected HttpException(int responseCode, String message, Throwable cause) {
+    this(responseCode, message, cause, Level.INFO);
   }
 
   public final int getResponseCode() {
@@ -57,7 +65,7 @@ public abstract class HttpException extends RuntimeException {
    */
   public final void send(HttpServletResponse rsp) throws IOException {
     rsp.sendError(getResponseCode(), htmlEscaper().escape(getMessage()));
-    logger.atInfo().withCause(getCause()).log("%s", this);
+    logger.at(logLevel).withCause(getCause()).log("%s", this);
   }
 
   /**
@@ -196,7 +204,7 @@ public abstract class HttpException extends RuntimeException {
   /** Exception that causes a 500 response. */
   public static final class InternalServerErrorException extends HttpException {
     public InternalServerErrorException(String message) {
-      super(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, message, null);
+      super(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, message, null, Level.SEVERE);
     }
 
     public InternalServerErrorException(String message, Throwable cause) {

core/src/main/java/google/registry/request/RequestParameters.java

@@ -45,9 +45,9 @@ public final class RequestParameters {
    *
    * <ul>
    *   <li>/foo?bar=hello → hello
-   *   <li>/foo?bar=hello&bar=there → hello
+   *   <li>/foo?bar=hello&amp;bar=there → hello
    *   <li>/foo?bar= → 400 error (empty)
-   *   <li>/foo?bar=&bar=there → 400 error (empty)
+   *   <li>/foo?bar=&amp;bar=there → 400 error (empty)
    *   <li>/foo → 400 error (absent)
    * </ul>
    *

core/src/main/java/google/registry/schema/cursor/CursorDao.java

@@ -59,7 +59,7 @@ public class CursorDao {
         .transact(() -> jpaTm().getEntityManager().find(Cursor.class, new CursorId(type, scope)));
   }
 
-  /** If no scope is given, use {@link Cursor.GLOBAL} as the scope. */
+  /** If no scope is given, use {@link Cursor#GLOBAL} as the scope. */
   public static Cursor load(CursorType type) {
     checkNotNull(type, "The type of the cursor to load must be specified");
     return load(type, Cursor.GLOBAL);

core/src/main/java/google/registry/schema/domain/RegistryLock.java

@@ -214,13 +214,13 @@ public final class RegistryLock extends ImmutableObject implements Buildable, Sq
     return lockCompletionTimestamp != null && unlockCompletionTimestamp == null;
   }
 
-  /** Returns true iff the lock was requested >= 1 hour ago and has not been verified. */
+  /** Returns true iff the lock was requested >= 1 hour ago and has not been verified. */
   public boolean isLockRequestExpired(DateTime now) {
     return !getLockCompletionTimestamp().isPresent()
         && isBeforeOrAt(getLockRequestTimestamp(), now.minusHours(1));
   }
 
-  /** Returns true iff the unlock was requested >= 1 hour ago and has not been verified. */
+  /** Returns true iff the unlock was requested >= 1 hour ago and has not been verified. */
   public boolean isUnlockRequestExpired(DateTime now) {
     Optional<DateTime> unlockRequestTimestamp = getUnlockRequestTimestamp();
     return unlockRequestTimestamp.isPresent()

core/src/main/java/google/registry/storage/drive/DriveConnection.java

@@ -42,7 +42,7 @@ public class DriveConnection {
   /**
    * Creates a folder with the given parent.
    *
-   * @returns the folder id.
+   * @return the folder id.
    */
   public String createFolder(String title, String parentFolderId) throws IOException {
     return drive.files()
@@ -58,7 +58,7 @@ public class DriveConnection {
    * existing file is the desired behavior, use {@link #createOrUpdateFile(String, MediaType,
    * String, byte[])} instead.
    *
-   * @returns the file id.
+   * @return the file id.
    */
   public String createFile(String title, MediaType mimeType, String parentFolderId, byte[] bytes)
       throws IOException {
@@ -76,13 +76,10 @@ public class DriveConnection {
    *
    * @throws IllegalStateException if multiple files with that name exist in the given folder.
    * @throws IOException if communication with Google Drive fails for any reason.
-   * @returns the file id.
+   * @return the file id.
    */
   public String createOrUpdateFile(
-      String title,
-      MediaType mimeType,
-      String parentFolderId,
-      byte[] bytes) throws IOException {
+      String title, MediaType mimeType, String parentFolderId, byte[] bytes) throws IOException {
     List<String> existingFiles = listFiles(parentFolderId, String.format("title = '%s'", title));
     if (existingFiles.size() > 1) {
       throw new IllegalStateException(String.format(
@@ -97,10 +94,10 @@ public class DriveConnection {
   }
 
   /**
-   * Updates the file with the given id in place, setting the title, content, and mime type to
-   * the newly specified values.
+   * Updates the file with the given id in place, setting the title, content, and mime type to the
+   * newly specified values.
    *
-   * @returns the file id.
+   * @return the file id.
    */
   public String updateFile(String fileId, String title, MediaType mimeType, byte[] bytes)
       throws IOException {

core/src/main/java/google/registry/tmch/TmchXmlSignature.java

@@ -77,10 +77,6 @@ public class TmchXmlSignature {
    *
    * @throws GeneralSecurityException for unsupported protocols, certs not signed by the TMCH,
    *     incorrect keys, and for invalid, old, not-yet-valid or revoked certificates.
-   * @throws IOException
-   * @throws MarshalException
-   * @throws ParserConfigurationException
-   * @throws SAXException
    */
   public void verify(byte[] smdXml)
       throws GeneralSecurityException, IOException, MarshalException, ParserConfigurationException,

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

@@ -39,18 +39,14 @@ class CompareDbBackups {
       return;
     }
 
-    ImmutableSet<ComparableEntity> entities1 =
-        new RecordAccumulator()
-            .readDirectory(new File(args[0]), DATA_FILE_MATCHER)
-            .getComparableEntitySet();
-    ImmutableSet<ComparableEntity> entities2 =
-        new RecordAccumulator()
-            .readDirectory(new File(args[1]), DATA_FILE_MATCHER)
-            .getComparableEntitySet();
+    ImmutableSet<EntityWrapper> entities1 =
+        RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER).getEntityWrapperSet();
+    ImmutableSet<EntityWrapper> entities2 =
+        RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER).getEntityWrapperSet();
 
     // Calculate the entities added and removed.
-    SetView<ComparableEntity> added = Sets.difference(entities2, entities1);
-    SetView<ComparableEntity> removed = Sets.difference(entities1, entities2);
+    SetView<EntityWrapper> added = Sets.difference(entities2, entities1);
+    SetView<EntityWrapper> removed = Sets.difference(entities1, entities2);
 
     printHeader(
         String.format("First backup: %d records", entities1.size()),
@@ -58,14 +54,14 @@ class CompareDbBackups {
 
     if (!removed.isEmpty()) {
       printHeader(removed.size() + " records were removed:");
-      for (ComparableEntity entity : removed) {
+      for (EntityWrapper entity : removed) {
         System.out.println(entity);
       }
     }
 
     if (!added.isEmpty()) {
       printHeader(added.size() + " records were added:");
-      for (ComparableEntity entity : added) {
+      for (EntityWrapper entity : added) {
         System.out.println(entity);
       }
     }

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

@@ -15,20 +15,32 @@
 package google.registry.tools;
 
 import com.google.appengine.api.datastore.Entity;
+import com.google.auto.value.AutoValue;
 import com.google.common.base.Objects;
 
-/** Wraps {@link Entity} to do hashCode/equals based on both the entity's key and its properties. */
-final class ComparableEntity {
+/**
+ * Wraps {@link Entity} for ease of processing in collections.
+ *
+ * <p>Note that the {@link #hashCode}/{@link #equals} methods are based on both the entity's key and
+ * its properties.
+ */
+final class EntityWrapper {
+  private static final String TEST_ENTITY_KIND = "TestEntity";
+
   private final Entity entity;
 
-  ComparableEntity(Entity entity) {
+  EntityWrapper(Entity entity) {
     this.entity = entity;
   }
 
+  public Entity getEntity() {
+    return entity;
+  }
+
   @Override
   public boolean equals(Object that) {
-    if (that instanceof ComparableEntity) {
-      ComparableEntity thatEntity = (ComparableEntity) that;
+    if (that instanceof EntityWrapper) {
+      EntityWrapper thatEntity = (EntityWrapper) that;
       return entity.equals(thatEntity.entity)
           && entity.getProperties().equals(thatEntity.entity.getProperties());
     }
@@ -43,6 +55,26 @@ final class ComparableEntity {
 
   @Override
   public String toString() {
-    return "ComparableEntity(" + entity + ")";
+    return "EntityWrapper(" + entity + ")";
+  }
+
+  public static EntityWrapper from(int id, Property... properties) {
+    Entity entity = new Entity(TEST_ENTITY_KIND, id);
+    for (Property prop : properties) {
+      entity.setProperty(prop.name(), prop.value());
+    }
+    return new EntityWrapper(entity);
+  }
+
+  @AutoValue
+  abstract static class Property {
+
+    static Property create(String name, Object value) {
+      return new AutoValue_EntityWrapper_Property(name, value);
+    }
+
+    abstract String name();
+
+    abstract Object value();
   }
 }

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

@@ -14,37 +14,101 @@
 
 package google.registry.tools;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.Optional;
 
 /**
- * Reads records from a set of LevelDB files and builds a gigantic ImmutableList from them.
+ * Iterator that incrementally parses binary data in LevelDb format into records.
  *
- * <p>See <a
- * href="https://github.com/google/leveldb/blob/master/doc/log_format.md">log_format.md</a> for the
- * leveldb log format specification.</a>
+ * <p>The input source is automatically closed when all data have been read.
  *
  * <p>There are several other implementations of this, none of which appeared suitable for our use
  * case: <a href="https://github.com/google/leveldb">The original C++ implementation</a>. <a
- * href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/files/RecordWriteChannel">
- * com.google.appengine.api.files.RecordWriteChannel</a> - Exactly what we need but deprecated. The
+ * href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/files/RecordReadChannel">
+ * com.google.appengine.api.files.RecordReadChannel</a> - Exactly what we need but deprecated. The
  * referenced replacement: <a
  * href="https://github.com/GoogleCloudPlatform/appengine-gcs-client.git">The App Engine GCS
- * Client</a> - Does not appear to have any support for working with LevelDB.
+ * Client</a> - Does not appear to have any support for working with LevelDB. *
+ *
+ * <p>See <a
+ * href="https://github.com/google/leveldb/blob/master/doc/log_format.md">log_format.md</a>
  */
-public final class LevelDbLogReader {
+public final class LevelDbLogReader implements Iterator<byte[]> {
 
   @VisibleForTesting static final int BLOCK_SIZE = 32 * 1024;
   @VisibleForTesting static final int HEADER_SIZE = 7;
 
   private final ByteArrayOutputStream recordContents = new ByteArrayOutputStream();
-  private final ImmutableList.Builder<byte[]> recordListBuilder = new ImmutableList.Builder<>();
+  private final LinkedList<byte[]> recordList = Lists.newLinkedList();
+
+  private final ByteBuffer byteBuffer = ByteBuffer.allocate(BLOCK_SIZE);
+  private final ReadableByteChannel channel;
+
+  LevelDbLogReader(ReadableByteChannel channel) {
+    this.channel = channel;
+  }
+
+  @Override
+  public boolean hasNext() {
+    while (recordList.isEmpty()) {
+      try {
+        Optional<byte[]> block = readFromChannel();
+        if (!block.isPresent()) {
+          return false;
+        }
+        if (block.get().length != BLOCK_SIZE) {
+          throw new IllegalStateException("Data size is not multiple of " + BLOCK_SIZE);
+        }
+        processBlock(block.get());
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    return true;
+  }
+
+  @Override
+  public byte[] next() {
+    checkState(hasNext(), "The next() method called on empty iterator.");
+    return recordList.removeFirst();
+  }
+
+  /**
+   * Returns the next {@link #BLOCK_SIZE} bytes from the input channel, or {@link
+   * Optional#empty()} if there is no more data.
+   */
+  // TODO(weiminyu): use ByteBuffer directly.
+  private Optional<byte[]> readFromChannel() throws IOException {
+    while (channel.isOpen()) {
+      int bytesRead = channel.read(byteBuffer);
+      if (!byteBuffer.hasRemaining() || bytesRead < 0) {
+        byteBuffer.flip();
+        if (!byteBuffer.hasRemaining()) {
+          channel.close();
+          return Optional.empty();
+        }
+        byte[] result = new byte[byteBuffer.remaining()];
+        byteBuffer.get(result);
+        byteBuffer.clear();
+        return Optional.of(result);
+      }
+    }
+    return Optional.empty();
+  }
 
   /** Read a complete block, which must be exactly 32 KB. */
   private void processBlock(byte[] block) {
@@ -63,7 +127,7 @@ public final class LevelDbLogReader {
 
       // If this is the last (or only) chunk in the record, store the full contents into the List.
       if (recordHeader.type == ChunkType.FULL || recordHeader.type == ChunkType.LAST) {
-        recordListBuilder.add(recordContents.toByteArray());
+        recordList.add(recordContents.toByteArray());
         recordContents.reset();
       }
 
@@ -96,40 +160,24 @@ public final class LevelDbLogReader {
     return new RecordHeader(checksum, size, ChunkType.fromCode(type));
   }
 
-  /** Reads all records in the Reader into the record set. */
-  public void readFrom(InputStream source) throws IOException {
-    byte[] block = new byte[BLOCK_SIZE];
-
-    // read until we have no more.
-    while (true) {
-      int amountRead = source.read(block, 0, BLOCK_SIZE);
-      if (amountRead <= 0) {
-        break;
-      }
-      assert amountRead == BLOCK_SIZE;
-
-      processBlock(block);
-    }
+  /** Returns a {@link LevelDbLogReader} over a {@link ReadableByteChannel}. */
+  public static LevelDbLogReader from(ReadableByteChannel channel) {
+    return new LevelDbLogReader(channel);
   }
 
-  /** Reads all records from the file specified by "path" into the record set. */
-  public void readFrom(Path path) throws IOException {
-    readFrom(Files.newInputStream(path));
+  /** Returns a {@link LevelDbLogReader} over an {@link InputStream}. */
+  public static LevelDbLogReader from(InputStream source) {
+    return new LevelDbLogReader(Channels.newChannel(source));
   }
 
-  /** Reads all records from the specified file into the record set. */
-  public void readFrom(String filename) throws IOException {
-    readFrom(FileSystems.getDefault().getPath(filename));
+  /** Returns a {@link LevelDbLogReader} over a file specified by {@link Path}. */
+  public static LevelDbLogReader from(Path path) throws IOException {
+    return from(Files.newInputStream(path));
   }
 
-  /**
-   * Gets the list of records constructed so far.
-   *
-   * <p>Note that this does not invalidate the internal state of the object: we return a copy and
-   * this can be called multiple times.
-   */
-  ImmutableList<byte[]> getRecords() {
-    return recordListBuilder.build();
+  /** Returns a {@link LevelDbLogReader} over a file specified by {@code filename}. */
+  public static LevelDbLogReader from(String filename) throws IOException {
+    return from(FileSystems.getDefault().getPath(filename));
   }
 
   /** Aggregates the fields in a record header. */

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

@@ -15,42 +15,47 @@
 package google.registry.tools;
 
 import com.google.appengine.api.datastore.EntityTranslator;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.function.Predicate;
 
 /** Accumulates Entity records from level db files under a directory hierarchy. */
 class RecordAccumulator {
-  private final LevelDbLogReader reader = new LevelDbLogReader();
+  private final ImmutableList<byte[]> records;
+
+  RecordAccumulator(ImmutableList<byte[]> records) {
+    this.records = records;
+  }
 
   /** Recursively reads all records in the directory. */
-  public final RecordAccumulator readDirectory(File dir, Predicate<File> fileMatcher) {
+  public static RecordAccumulator readDirectory(File dir, Predicate<File> fileMatcher) {
+    ImmutableList.Builder<byte[]> builder = new ImmutableList.Builder<>();
     for (File child : dir.listFiles()) {
       if (child.isDirectory()) {
-        readDirectory(child, fileMatcher);
+        builder.addAll(readDirectory(child, fileMatcher).records);
       } else if (fileMatcher.test(child)) {
         try {
-          reader.readFrom(new FileInputStream(child));
+          builder.addAll(LevelDbLogReader.from(child.getPath()));
         } catch (IOException e) {
           throw new RuntimeException("IOException reading from file: " + child, e);
         }
       }
     }
 
-    return this;
+    return new RecordAccumulator(builder.build());
   }
 
-  /** Creates an entity set from the current set of raw records. */
-  ImmutableSet<ComparableEntity> getComparableEntitySet() {
-    ImmutableSet.Builder<ComparableEntity> builder = new ImmutableSet.Builder<>();
-    for (byte[] rawRecord : reader.getRecords()) {
+  /** Creates an {@link EntityWrapper} set from the current set of raw records. */
+  ImmutableSet<EntityWrapper> getEntityWrapperSet() {
+    ImmutableSet.Builder<EntityWrapper> builder = new ImmutableSet.Builder<>();
+    for (byte[] rawRecord : records) {
       // Parse the entity proto and create an Entity object from it.
       EntityProto proto = new EntityProto();
       proto.parseFrom(rawRecord);
-      ComparableEntity entity = new ComparableEntity(EntityTranslator.createFromPb(proto));
+      EntityWrapper entity = new EntityWrapper(EntityTranslator.createFromPb(proto));
 
       builder.add(entity);
     }

core/src/main/java/google/registry/ui/forms/FormField.java

@@ -44,43 +44,45 @@ import javax.annotation.concurrent.Immutable;
  * Declarative functional fluent form field converter / validator.
  *
  * <p>This class is responsible for converting arbitrary data, sent to us by the web browser, into
- * validated data structures that the server-side code can use. For example:<pre>
+ * validated data structures that the server-side code can use. For example:
  *
- *   private enum Gender { MALE, FEMALE }
+ * <pre>{@code
+ * private enum Gender { MALE, FEMALE }
  *
- *   private static final FormField<String, String> NAME_FIELD = FormField.named("name")
- *       .matches("[a-z]+")
- *       .range(atMost(16))
- *       .required()
- *       .build();
+ * private static final FormField<String, String> NAME_FIELD = FormField.named("name")
+ *     .matches("[a-z]+")
+ *     .range(atMost(16))
+ *     .required()
+ *     .build();
  *
- *   private static final FormField<String, Gender> GENDER_FIELD = FormField.named("gender")
- *       .asEnum(Gender.class)
- *       .required()
- *       .build();
+ * private static final FormField<String, Gender> GENDER_FIELD = FormField.named("gender")
+ *     .asEnum(Gender.class)
+ *     .required()
+ *     .build();
  *
- *   public Person makePerson(Map<String, String> params) {
- *     Person.Builder person = new Person.Builder();
- *     for (String name : NAME_FIELD.extract(params).asSet()) {
- *       person.setName(name);
- *     }
- *     for (Gender name : GENDER_FIELD.extract(params).asSet()) {
- *       person.setGender(name);
- *     }
- *     return person.build();
- *   }</pre>
+ * public Person makePerson(Map<String, String> params) {
+ *   Person.Builder person = new Person.Builder();
+ *   for (String name : NAME_FIELD.extract(params).asSet()) {
+ *     person.setName(name);
+ *   }
+ *   for (Gender name : GENDER_FIELD.extract(params).asSet()) {
+ *     person.setGender(name);
+ *   }
+ *   return person.build();
+ * }
+ * }</pre>
  *
  * <p>This class provides <b>full type-safety</b> <i>if and only if</i> you statically initialize
  * your FormField objects and write a unit test that causes the class to be loaded.
  *
  * <h3>Exception Handling</h3>
  *
- * <p>When values passed to {@link #convert} or {@link #extract} don't meet the contract,
- * {@link FormFieldException} will be thrown, which provides the field name and a short error
- * message that's safe to pass along to the client.
+ * <p>When values passed to {@link #convert} or {@link #extract} don't meet the contract, {@link
+ * FormFieldException} will be thrown, which provides the field name and a short error message
+ * that's safe to pass along to the client.
  *
- * <p>You can safely throw {@code FormFieldException} from within your validator functions, and
- * the field name will automatically be propagated into the exception object for you.
+ * <p>You can safely throw {@code FormFieldException} from within your validator functions, and the
+ * field name will automatically be propagated into the exception object for you.
  *
  * <p>In situations when you're validating lists or maps, you'll end up with a hierarchical field
  * naming structure. For example, if you were validating a list of maps, an error generated by the
@@ -92,23 +94,25 @@ import javax.annotation.concurrent.Immutable;
  * <p>You should never assign a partially constructed {@code FormField.Builder} to a variable or
  * constant. Instead, you should use {@link #asBuilder()} or {@link #asBuilderNamed(String)}.
  *
- * <p>Here is an example of how you might go about defining library definitions:<pre>
+ * <p>Here is an example of how you might go about defining library definitions:
  *
- *   final class FormFields {
- *     private static final FormField<String, String> COUNTRY_CODE =
- *         FormField.named("countryCode")
- *             .range(Range.singleton(2))
- *             .uppercased()
- *             .in(ImmutableSet.copyOf(Locale.getISOCountries()))
- *             .build();
- *   }
+ * <pre>{@code
+ * final class FormFields {
+ *   private static final FormField<String, String> COUNTRY_CODE =
+ *       FormField.named("countryCode")
+ *           .range(Range.singleton(2))
+ *           .uppercased()
+ *           .in(ImmutableSet.copyOf(Locale.getISOCountries()))
+ *           .build();
+ * }
  *
- *   final class Form {
- *     private static final FormField<String, String> COUNTRY_CODE_FIELD =
- *         FormFields.COUNTRY_CODE.asBuilder()
- *             .required()
- *             .build();
- *   }</pre>
+ * final class Form {
+ *   private static final FormField<String, String> COUNTRY_CODE_FIELD =
+ *       FormFields.COUNTRY_CODE.asBuilder()
+ *           .required()
+ *           .build();
+ * }
+ * }</pre>
  *
  * @param <I> input value type
  * @param <O> output value type

core/src/main/java/google/registry/ui/server/SendEmailUtils.java

@@ -105,9 +105,10 @@ public class SendEmailUtils {
     }
   }
 
-  /** Sends an email from Nomulus to the registrarChangesNotificationEmailAddresses.
+  /**
+   * Sends an email from Nomulus to the registrarChangesNotificationEmailAddresses.
    *
-   * <p>See {@link #sendEmail(String, String, ImmutableList<String>)}.
+   * <p>See {@link #sendEmail(String, String, ImmutableList)}.
    */
   public boolean sendEmail(final String subject, String body) {
     return sendEmail(subject, body, ImmutableList.of());

core/src/main/java/google/registry/whois/WhoisHttpAction.java

@@ -58,40 +58,38 @@ import org.joda.time.Duration;
  * <p>You can send AJAX requests to our WHOIS API from your <em>very own</em> website using the
  * following embed code:
  *
- * <pre>
- *   <p>
- *     <input id="query-input" placeholder="Domain, Nameserver, IP, etc." autofocus>
- *     <button id="search-button">Lookup</button>
- *   <p>
- *     <pre id="whois-results"></pre>
- *   <script>
- *    (function() {
- *      var WHOIS_API_URL = 'https://domain-registry-alpha.appspot.com/whois/';
- *      function OnKeyPressQueryInput(ev) {
- *        if (typeof ev == 'undefined' && window.event) {
- *          ev = window.event;
- *        }
- *        if (ev.keyCode == 13) {
- *          document.getElementById('search-button').click();
- *        }
+ * <pre>{@code
+ * <input id="query-input" placeholder="Domain, Nameserver, IP, etc." autofocus>
+ * <button id="search-button">Lookup</button>
+ * <pre id="whois-results"></pre>
+ * <script>
+ *  (function() {
+ *    var WHOIS_API_URL = 'https://domain-registry-alpha.appspot.com/whois/';
+ *    function OnKeyPressQueryInput(ev) {
+ *      if (typeof ev == 'undefined' && window.event) {
+ *        ev = window.event;
  *      }
- *      function OnClickSearchButton() {
- *        var query = document.getElementById('query-input').value;
- *        var req = new XMLHttpRequest();
- *        req.onreadystatechange = function() {
- *          if (req.readyState == 4) {
- *            var results = document.getElementById('whois-results');
- *            results.textContent = req.responseText;
- *          }
- *        };
- *        req.open('GET', WHOIS_API_URL + escape(query), true);
- *        req.send();
+ *      if (ev.keyCode == 13) {
+ *        document.getElementById('search-button').click();
  *      }
- *      document.getElementById('search-button').onclick = OnClickSearchButton;
- *      document.getElementById('query-input').onkeypress = OnKeyPressQueryInput;
- *    })();
- *   </script>
- * </pre>
+ *    }
+ *    function OnClickSearchButton() {
+ *      var query = document.getElementById('query-input').value;
+ *      var req = new XMLHttpRequest();
+ *      req.onreadystatechange = function() {
+ *        if (req.readyState == 4) {
+ *          var results = document.getElementById('whois-results');
+ *          results.textContent = req.responseText;
+ *        }
+ *      };
+ *      req.open('GET', WHOIS_API_URL + escape(query), true);
+ *      req.send();
+ *    }
+ *    document.getElementById('search-button').onclick = OnClickSearchButton;
+ *    document.getElementById('query-input').onkeypress = OnKeyPressQueryInput;
+ *  })();
+ * </script>
+ * }</pre>
  *
  * @see WhoisAction
  */

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

@@ -19,6 +19,9 @@
       *  Move tests to another (sub)project. This is not a big problem, but feels unnatural.
       *  Use Hibernate's ServiceRegistry for bootstrapping (not JPA-compliant)
     -->
+    <class>google.registry.model.billing.BillingEvent$Cancellation</class>
+    <class>google.registry.model.billing.BillingEvent$OneTime</class>
+    <class>google.registry.model.billing.BillingEvent$Recurring</class>
     <class>google.registry.model.contact.ContactResource</class>
     <class>google.registry.model.domain.DomainBase</class>
     <class>google.registry.model.host.HostResource</class>
@@ -36,6 +39,7 @@
 
     <!-- Customized type converters -->
     <class>google.registry.persistence.converter.BillingCostTransitionConverter</class>
+    <class>google.registry.persistence.converter.BillingEventFlagSetConverter</class>
     <class>google.registry.persistence.converter.BloomFilterConverter</class>
     <class>google.registry.persistence.converter.CidrAddressBlockListConverter</class>
     <class>google.registry.persistence.converter.CreateAutoTimestampConverter</class>
@@ -53,6 +57,8 @@
     <class>google.registry.persistence.converter.ZonedDateTimeConverter</class>
 
     <!-- Generated converters for VKey -->
+    <class>google.registry.model.billing.VKeyConverter_BillingEvent</class>
+    <class>google.registry.model.domain.token.VKeyConverter_AllocationToken</class>
     <class>google.registry.model.host.VKeyConverter_HostResource</class>
     <class>google.registry.model.contact.VKeyConverter_ContactResource</class>
 

core/src/main/resources/google/registry/ui/soy/registrar/OteSetupConsole.soy

@@ -55,7 +55,7 @@
 {template .resultSuccess}
   {@param baseClientId: string} /** The base clientId used for the OT&E setup. */
   {@param contactEmail: string} /** The contact's email added to the registrars. */
-  {@param clientIdToTld: map<string, string>} /** The created registrars->TLD mapping. */
+  {@param clientIdToTld: map<string, string>} /** The created registrars-&gt;TLD mapping. */
   {@param password: string} /** The password given for the created registrars. */
   {@param username: string} /** Arbitrary username to display. */
   {@param logoutUrl: string} /** Generated URL for logging out of Google. */

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

@@ -155,7 +155,7 @@ public class ExpandRecurringBillingEventsActionTest
         .setPeriodYears(1)
         .setReason(Reason.RENEW)
         .setSyntheticCreationTime(beginningOfTest)
-        .setCancellationMatchingBillingEvent(Key.create(recurring))
+        .setCancellationMatchingBillingEvent(recurring.createVKey())
         .setTargetId(domain.getFullyQualifiedDomainName());
   }
 
@@ -274,10 +274,12 @@ public class ExpandRecurringBillingEventsActionTest
         .setParent(persistedEntries.get(0))
         .build();
     // Persist an otherwise identical billing event that differs only in recurring event key.
-    BillingEvent.OneTime persisted = expected.asBuilder()
-        .setParent(persistedEntries.get(1))
-        .setCancellationMatchingBillingEvent(Key.create(recurring2))
-        .build();
+    BillingEvent.OneTime persisted =
+        expected
+            .asBuilder()
+            .setParent(persistedEntries.get(1))
+            .setCancellationMatchingBillingEvent(recurring2.createVKey())
+            .build();
     assertCursorAt(beginningOfTest);
     assertBillingEventsForResource(domain, persisted, expected, recurring, recurring2);
   }
@@ -604,19 +606,21 @@ public class ExpandRecurringBillingEventsActionTest
     assertHistoryEntryMatches(
         domain, persistedEntries.get(0), "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"),
         true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setParent(persistedEntries.get(0))
-        .setCancellationMatchingBillingEvent(Key.create(recurring))
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setParent(persistedEntries.get(0))
+            .setCancellationMatchingBillingEvent(recurring.createVKey())
+            .build();
     assertHistoryEntryMatches(
         domain, persistedEntries.get(1), "TheRegistrar", DateTime.parse("2000-05-20T00:00:00Z"),
         true);
-    BillingEvent.OneTime expected2 = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2000-05-20T00:00:00Z"))
-        .setEventTime(DateTime.parse("2000-04-05T00:00:00Z"))
-        .setParent(persistedEntries.get(1))
-        .setCancellationMatchingBillingEvent(Key.create(recurring2))
-        .build();
+    BillingEvent.OneTime expected2 =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2000-05-20T00:00:00Z"))
+            .setEventTime(DateTime.parse("2000-04-05T00:00:00Z"))
+            .setParent(persistedEntries.get(1))
+            .setCancellationMatchingBillingEvent(recurring2.createVKey())
+            .build();
     assertBillingEventsForResource(domain, expected, expected2, recurring, recurring2);
     assertCursorAt(beginningOfTest);
   }

core/src/test/java/google/registry/export/BigqueryPollJobActionTest.java

@@ -17,6 +17,7 @@ package google.registry.export;
 import static com.google.appengine.api.taskqueue.QueueFactory.getQueue;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.testing.TaskQueueHelper.assertNoTasksEnqueued;
 import static google.registry.testing.TaskQueueHelper.assertTasksEnqueued;
 import static google.registry.testing.TestLogHandlerUtils.assertLogMessage;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -174,6 +175,7 @@ public class BigqueryPollJobActionTest {
     action.run();
     assertLogMessage(
         logHandler, SEVERE, String.format("Bigquery job failed - %s:%s", PROJECT_ID, JOB_ID));
+    assertNoTasksEnqueued(CHAINED_QUEUE_NAME);
   }
 
   @Test

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

@@ -41,6 +41,7 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.HistoryEntry.Type;
 import google.registry.monitoring.whitebox.EppMetric;
+import google.registry.persistence.VKey;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeHttpSession;
 import google.registry.testing.FakeResponse;
@@ -339,7 +340,8 @@ public class EppTestCase extends ShardableTestCase {
         .setTargetId(domain.getFullyQualifiedDomainName())
         .setClientId(domain.getCurrentSponsorClientId())
         .setEventTime(deleteTime)
-        .setOneTimeEventKey(findKeyToActualOneTimeBillingEvent(billingEventToCancel))
+        .setOneTimeEventKey(
+            VKey.createOfy(OneTime.class, findKeyToActualOneTimeBillingEvent(billingEventToCancel)))
         .setBillingTime(createTime.plus(Registry.get(domain.getTld()).getAddGracePeriodLength()))
         .setReason(Reason.CREATE)
         .setParent(getOnlyHistoryEntryOfType(domain, Type.DOMAIN_DELETE))
@@ -353,7 +355,8 @@ public class EppTestCase extends ShardableTestCase {
         .setTargetId(domain.getFullyQualifiedDomainName())
         .setClientId(domain.getCurrentSponsorClientId())
         .setEventTime(deleteTime)
-        .setOneTimeEventKey(findKeyToActualOneTimeBillingEvent(billingEventToCancel))
+        .setOneTimeEventKey(
+            VKey.createOfy(OneTime.class, findKeyToActualOneTimeBillingEvent(billingEventToCancel)))
         .setBillingTime(renewTime.plus(Registry.get(domain.getTld()).getRenewGracePeriodLength()))
         .setReason(Reason.RENEW)
         .setParent(getOnlyHistoryEntryOfType(domain, Type.DOMAIN_DELETE))

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

@@ -276,7 +276,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
             .setBillingTime(billingTime)
             .setFlags(expectedBillingFlags)
             .setParent(historyEntry)
-            .setAllocationToken(allocationToken == null ? null : Key.create(allocationToken))
+            .setAllocationToken(allocationToken == null ? null : allocationToken.createVKey())
             .build();
 
     BillingEvent.Recurring renewBillingEvent =

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

@@ -214,7 +214,7 @@ public class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow,
             .setClientId("TheRegistrar")
             .setEventTime(eventTime)
             .setBillingTime(TIME_BEFORE_FLOW.plusDays(1))
-            .setOneTimeEventKey(Key.create(graceBillingEvent))
+            .setOneTimeEventKey(graceBillingEvent.createVKey())
             .setParent(historyEntryDomainDelete)
             .build());
   }

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

@@ -71,6 +71,7 @@ import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferResponse.DomainTransferResponse;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
 import java.util.Arrays;
 import java.util.stream.Stream;
 import org.joda.money.Money;
@@ -402,7 +403,8 @@ public class DomainTransferApproveFlowTest
             .setEventTime(clock.nowUtc()) // The cancellation happens at the moment of transfer.
             .setBillingTime(
                 oldExpirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
-            .setRecurringEventKey(domain.getAutorenewBillingEvent()));
+            .setRecurringEventKey(
+                VKey.createOfy(BillingEvent.Recurring.class, domain.getAutorenewBillingEvent())));
   }
 
   @Test

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

@@ -100,6 +100,7 @@ import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferResponse;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import java.util.Map;
 import java.util.Optional;
@@ -1136,7 +1137,8 @@ public class DomainTransferRequestFlowTest
             .setEventTime(clock.nowUtc().plus(Registry.get("tld").getAutomaticTransferLength()))
             .setBillingTime(autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(existingAutorenewEvent));
+            .setRecurringEventKey(
+                VKey.createOfy(BillingEvent.Recurring.class, existingAutorenewEvent)));
   }
 
   @Test
@@ -1164,7 +1166,8 @@ public class DomainTransferRequestFlowTest
             .setBillingTime(
                 expirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(existingAutorenewEvent));
+            .setRecurringEventKey(
+                VKey.createOfy(BillingEvent.Recurring.class, existingAutorenewEvent)));
   }
 
   @Test

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

@@ -28,6 +28,7 @@ import com.googlecode.objectify.annotation.Parent;
 import com.googlecode.objectify.annotation.Serialize;
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.ofy.Ofy;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
@@ -164,7 +165,8 @@ public abstract class EntityTestCase {
                     : (Class<?>) inner;
           }
           // Descend into persisted ImmutableObject classes, but not anything else.
-          if (ImmutableObject.class.isAssignableFrom(fieldClass)) {
+          if (ImmutableObject.class.isAssignableFrom(fieldClass)
+              && !VKey.class.isAssignableFrom(fieldClass)) {
             getAllPotentiallyIndexedFieldPaths(fieldClass).stream()
                 .map(subfield -> field.getName() + "." + subfield)
                 .distinct()

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

@@ -17,9 +17,11 @@ package google.registry.model.billing;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistActiveDomain;
 import static google.registry.testing.DatastoreHelper.persistResource;
+import static google.registry.testing.SqlHelper.saveRegistrar;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static org.joda.money.CurrencyUnit.USD;
 import static org.joda.time.DateTimeZone.UTC;
@@ -37,19 +39,25 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.util.DateTimeUtils;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Test;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link BillingEvent}. */
 public class BillingEventTest extends EntityTestCase {
   private final DateTime now = DateTime.now(UTC);
 
+  public BillingEventTest() {
+    super(true);
+  }
+
   HistoryEntry historyEntry;
   HistoryEntry historyEntry2;
   DomainBase domain;
+  BillingEvent.OneTime sqlOneTime;
   BillingEvent.OneTime oneTime;
   BillingEvent.OneTime oneTimeSynthetic;
   BillingEvent.Recurring recurring;
@@ -57,7 +65,7 @@ public class BillingEventTest extends EntityTestCase {
   BillingEvent.Cancellation cancellationRecurring;
   BillingEvent.Modification modification;
 
-  @Before
+  @BeforeEach
   public void setUp() {
     createTld("tld");
     domain = persistActiveDomain("foo.tld");
@@ -97,7 +105,18 @@ public class BillingEventTest extends EntityTestCase {
                     .setCost(Money.of(USD, 1))
                     .setEventTime(now)
                     .setBillingTime(now.plusDays(5))
-                    .setAllocationToken(Key.create(allocationToken))));
+                    .setAllocationToken(allocationToken.createVKey())));
+
+    sqlOneTime =
+        oneTime
+            .asBuilder()
+            .setDomainRepoId(domain.getRepoId())
+            .setDomainHistoryRevisionId(1L)
+            .setAllocationToken(
+                VKey.create(
+                    AllocationToken.class, allocationToken.getToken(), Key.create(allocationToken)))
+            .build();
+
     recurring =
         persistResource(
             commonInit(
@@ -107,31 +126,41 @@ public class BillingEventTest extends EntityTestCase {
                     .setReason(Reason.RENEW)
                     .setEventTime(now.plusYears(1))
                     .setRecurrenceEndTime(END_OF_TIME)));
-    oneTimeSynthetic = persistResource(commonInit(
-        new BillingEvent.OneTime.Builder()
-            .setParent(historyEntry)
-            .setReason(Reason.CREATE)
-            .setFlags(ImmutableSet.of(BillingEvent.Flag.ANCHOR_TENANT, BillingEvent.Flag.SYNTHETIC))
-            .setSyntheticCreationTime(now.plusDays(10))
-            .setCancellationMatchingBillingEvent(Key.create(recurring))
-            .setPeriodYears(2)
-            .setCost(Money.of(USD, 1))
-            .setEventTime(now)
-            .setBillingTime(now.plusDays(5))));
-    cancellationOneTime = persistResource(commonInit(
-        new BillingEvent.Cancellation.Builder()
-            .setParent(historyEntry2)
-            .setReason(Reason.CREATE)
-            .setEventTime(now.plusDays(1))
-            .setBillingTime(now.plusDays(5))
-            .setOneTimeEventKey(Key.create(oneTime))));
-    cancellationRecurring = persistResource(commonInit(
-        new BillingEvent.Cancellation.Builder()
-            .setParent(historyEntry2)
-            .setReason(Reason.RENEW)
-            .setEventTime(now.plusDays(1))
-            .setBillingTime(now.plusYears(1).plusDays(45))
-            .setRecurringEventKey(Key.create(recurring))));
+    oneTimeSynthetic =
+        persistResource(
+            commonInit(
+                new BillingEvent.OneTime.Builder()
+                    .setParent(historyEntry)
+                    .setReason(Reason.CREATE)
+                    .setFlags(
+                        ImmutableSet.of(
+                            BillingEvent.Flag.ANCHOR_TENANT, BillingEvent.Flag.SYNTHETIC))
+                    .setSyntheticCreationTime(now.plusDays(10))
+                    .setCancellationMatchingBillingEvent(recurring.createVKey())
+                    .setPeriodYears(2)
+                    .setCost(Money.of(USD, 1))
+                    .setEventTime(now)
+                    .setBillingTime(now.plusDays(5))));
+
+    cancellationOneTime =
+        persistResource(
+            commonInit(
+                new BillingEvent.Cancellation.Builder()
+                    .setParent(historyEntry2)
+                    .setReason(Reason.CREATE)
+                    .setEventTime(now.plusDays(1))
+                    .setBillingTime(now.plusDays(5))
+                    .setOneTimeEventKey(oneTime.createVKey())));
+
+    cancellationRecurring =
+        persistResource(
+            commonInit(
+                new BillingEvent.Cancellation.Builder()
+                    .setParent(historyEntry2)
+                    .setReason(Reason.RENEW)
+                    .setEventTime(now.plusDays(1))
+                    .setBillingTime(now.plusYears(1).plusDays(45))
+                    .setRecurringEventKey(recurring.createVKey())));
     modification = persistResource(commonInit(
         new BillingEvent.Modification.Builder()
             .setParent(historyEntry2)
@@ -149,6 +178,83 @@ public class BillingEventTest extends EntityTestCase {
         .build();
   }
 
+  private void saveNewBillingEvent(BillingEvent billingEvent) {
+    billingEvent.id = null;
+    jpaTm().transact(() -> jpaTm().saveNew(billingEvent));
+  }
+
+  @Test
+  public void testCloudSqlPersistence_OneTime() {
+    saveRegistrar("a registrar");
+    saveNewBillingEvent(sqlOneTime);
+
+    BillingEvent.OneTime persisted =
+        jpaTm()
+            .transact(
+                () -> jpaTm().load(VKey.createSql(BillingEvent.OneTime.class, sqlOneTime.id)));
+    // TODO(shicong): Remove these fixes after the entities are fully compatible
+    BillingEvent.OneTime fixed =
+        persisted
+            .asBuilder()
+            .setParent(sqlOneTime.getParentKey())
+            .setAllocationToken(sqlOneTime.getAllocationToken().get())
+            .build();
+    assertThat(fixed).isEqualTo(sqlOneTime);
+  }
+
+  @Test
+  public void testCloudSqlPersistence_Cancellation() {
+    saveRegistrar("a registrar");
+    saveNewBillingEvent(sqlOneTime);
+    VKey<BillingEvent.OneTime> sqlVKey = VKey.createSql(BillingEvent.OneTime.class, sqlOneTime.id);
+    BillingEvent sqlCancellationOneTime =
+        cancellationOneTime
+            .asBuilder()
+            .setOneTimeEventKey(sqlVKey)
+            .setDomainRepoId(domain.getRepoId())
+            .setDomainHistoryRevisionId(1L)
+            .build();
+    saveNewBillingEvent(sqlCancellationOneTime);
+
+    BillingEvent.Cancellation persisted =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .load(
+                            VKey.createSql(
+                                BillingEvent.Cancellation.class, sqlCancellationOneTime.id)));
+    // TODO(shicong): Remove these fixes after the entities are fully compatible
+    BillingEvent.Cancellation fixed =
+        persisted
+            .asBuilder()
+            .setParent(sqlCancellationOneTime.getParentKey())
+            .setOneTimeEventKey(sqlVKey)
+            .build();
+    assertThat(fixed).isEqualTo(sqlCancellationOneTime);
+  }
+
+  @Test
+  public void testCloudSqlPersistence_Recurring() {
+    saveRegistrar("a registrar");
+    BillingEvent.Recurring sqlRecurring =
+        recurring
+            .asBuilder()
+            .setDomainRepoId(domain.getRepoId())
+            .setDomainHistoryRevisionId(1L)
+            .build();
+    saveNewBillingEvent(sqlRecurring);
+
+    BillingEvent.Recurring persisted =
+        jpaTm()
+            .transact(
+                () -> jpaTm().load(VKey.createSql(BillingEvent.Recurring.class, sqlRecurring.id)));
+    // TODO(shicong): Remove these fixes after the entities are fully compatible
+    BillingEvent.Recurring fixed =
+        persisted.asBuilder().setParent(sqlRecurring.getParentKey()).build();
+    assertThat(fixed).isEqualTo(sqlRecurring);
+  }
+
   @Test
   public void testPersistence() {
     assertThat(ofy().load().entity(oneTime).now()).isEqualTo(oneTime);
@@ -183,8 +289,13 @@ public class BillingEventTest extends EntityTestCase {
 
   @Test
   public void testCancellationMatching() {
-    Key<?> recurringKey = ofy().load().entity(oneTimeSynthetic).now()
-        .getCancellationMatchingBillingEvent();
+    Key<?> recurringKey =
+        ofy()
+            .load()
+            .entity(oneTimeSynthetic)
+            .now()
+            .getCancellationMatchingBillingEvent()
+            .getOfyKey();
     assertThat(ofy().load().key(recurringKey).now()).isEqualTo(recurring);
   }
 
@@ -219,7 +330,7 @@ public class BillingEventTest extends EntityTestCase {
                 oneTime
                     .asBuilder()
                     .setFlags(ImmutableSet.of(BillingEvent.Flag.SYNTHETIC))
-                    .setCancellationMatchingBillingEvent(Key.create(recurring))
+                    .setCancellationMatchingBillingEvent(recurring.createVKey())
                     .build());
     assertThat(thrown)
         .hasMessageThat()
@@ -263,7 +374,7 @@ public class BillingEventTest extends EntityTestCase {
             () ->
                 oneTime
                     .asBuilder()
-                    .setCancellationMatchingBillingEvent(Key.create(recurring))
+                    .setCancellationMatchingBillingEvent(recurring.createVKey())
                     .build());
     assertThat(thrown)
         .hasMessageThat()
@@ -334,8 +445,8 @@ public class BillingEventTest extends EntityTestCase {
             () ->
                 cancellationOneTime
                     .asBuilder()
-                    .setOneTimeEventKey(Key.create(oneTime))
-                    .setRecurringEventKey(Key.create(recurring))
+                    .setOneTimeEventKey(oneTime.createVKey())
+                    .setRecurringEventKey(recurring.createVKey())
                     .build());
     assertThat(thrown).hasMessageThat().contains("exactly one billing event");
   }

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

@@ -119,7 +119,8 @@ public class DomainBaseSqlTest {
         .transact(
             () -> {
               // Persist the contacts.  Note that these need to be persisted before the domain
-              // otherwise we get a foreign key constraint error.
+              // otherwise we get a foreign key constraint error.  If we ever decide to defer the
+              // relevant foreign key checks to commit time, then the order would not matter.
               jpaTm().saveNew(contact);
               jpaTm().saveNew(contact2);
 
@@ -127,7 +128,8 @@ public class DomainBaseSqlTest {
               jpaTm().saveNew(domain);
 
               // Persist the host.  This does _not_ need to be persisted before the domain,
-              // presumably because its relationship is stored in a join table.
+              // because only the row in the join table (DomainHost) is subject to foreign key
+              // constraints, and Hibernate knows to insert it after domain and host.
               jpaTm().saveNew(host);
             });
 

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

@@ -771,60 +771,4 @@ public class DomainBaseTest extends EntityTestCase {
     assertThat(getOnlyElement(clone.getGracePeriods()).getType())
         .isEqualTo(GracePeriodStatus.TRANSFER);
   }
-
-  private static ImmutableSet<Key<HostResource>> getOfyNameservers(DomainBase domain) {
-    return domain.getNameservers().stream().map(key -> key.getOfyKey()).collect(toImmutableSet());
-  }
-
-  @Test
-  public void testNameservers_nsHostsOfyKeys() {
-    assertThat(domain.nsHosts).isEqualTo(getOfyNameservers(domain));
-
-    // Test the setNameserver that functions on a function.
-    VKey<HostResource> host1Key =
-        persistResource(
-                new HostResource.Builder()
-                    .setFullyQualifiedHostName("ns2.example.com")
-                    .setSuperordinateDomain(domainKey)
-                    .setRepoId("2-COM")
-                    .build())
-            .createKey();
-
-    DomainBase dom = new DomainBase.Builder(domain).setNameservers(host1Key).build();
-    assertThat(dom.getNameservers()).isEqualTo(ImmutableSet.of(host1Key));
-    assertThat(getOfyNameservers(dom)).isEqualTo(ImmutableSet.of(host1Key.getOfyKey()));
-
-    // Test that setting to a single host of null throws an NPE.
-    assertThrows(
-        NullPointerException.class,
-        () -> new DomainBase.Builder(domain).setNameservers((VKey<HostResource>) null));
-
-    // Test that setting to a set of values works.
-    VKey<HostResource> host2Key =
-        persistResource(
-                new HostResource.Builder()
-                    .setFullyQualifiedHostName("ns3.example.com")
-                    .setSuperordinateDomain(domainKey)
-                    .setRepoId("3-COM")
-                    .build())
-            .createKey();
-    dom =
-        new DomainBase.Builder(domain).setNameservers(ImmutableSet.of(host1Key, host2Key)).build();
-    assertThat(dom.getNameservers()).isEqualTo(ImmutableSet.of(host1Key, host2Key));
-    assertThat(getOfyNameservers(dom))
-        .isEqualTo(ImmutableSet.of(host1Key.getOfyKey(), host2Key.getOfyKey()));
-
-    // Set of values, passing null.
-    dom =
-        new DomainBase.Builder(domain)
-            .setNameservers((ImmutableSet<VKey<HostResource>>) null)
-            .build();
-    assertThat(dom.nsHostVKeys).isNull();
-    assertThat(dom.nsHosts).isNull();
-
-    // Empty set of values gets translated to null.
-    dom = new DomainBase.Builder(domain).setNameservers(ImmutableSet.of()).build();
-    assertThat(dom.nsHostVKeys).isNull();
-    assertThat(dom.nsHosts).isNull();
-  }
 }

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

@@ -45,7 +45,7 @@ public class JpaEntityCoverage extends ExternalResource {
   private static final ImmutableSet<Class> ALL_JPA_ENTITIES =
       PersistenceXmlUtility.getManagedClasses().stream()
           .filter(e -> !IGNORE_ENTITIES.contains(e.getSimpleName()))
-          .filter(e -> e.getAnnotation(Entity.class) != null)
+          .filter(e -> e.isAnnotationPresent(Entity.class))
           .collect(ImmutableSet.toImmutableSet());
   private static final Set<Class> allCoveredJpaEntities = Sets.newHashSet();
   // Map of test class name to boolean flag indicating if it tests any JPA entities.

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

@@ -28,6 +28,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import org.hibernate.cfg.Environment;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
@@ -180,6 +181,17 @@ public class JpaTestRules {
       return this;
     }
 
+    /**
+     * Enables logging of SQL statements.
+     *
+     * <p>SQL logging is very noisy and disabled by default. This method maybe useful when
+     * troubleshooting a specific test.
+     */
+    public Builder withSqlLogging() {
+      withProperty(Environment.SHOW_SQL, "true");
+      return this;
+    }
+
     /** Builds a {@link JpaIntegrationTestRule} instance. */
     public JpaIntegrationTestRule buildIntegrationTestRule() {
       return new JpaIntegrationTestRule(

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

@@ -0,0 +1,57 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.transaction;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.io.UnsupportedEncodingException;
+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;
+
+/** Unit test for {@link JpaTestRules.Builder#withSqlLogging()}. */
+public class JpaTestRulesSqlLoggingTest {
+
+  // Entity under test: configured to log SQL statements to Stdout.
+  @RegisterExtension
+  JpaUnitTestRule jpaRule = new JpaTestRules.Builder().withSqlLogging().buildUnitTestRule();
+
+  private PrintStream orgStdout;
+  private ByteArrayOutputStream stdoutBuffer;
+
+  @BeforeEach
+  public void beforeEach() {
+    orgStdout = System.out;
+    System.setOut(new PrintStream(stdoutBuffer = new ByteArrayOutputStream()));
+  }
+
+  @AfterEach
+  public void afterEach() {
+    System.setOut(orgStdout);
+  }
+
+  @Test
+  void sqlLog_displayed() throws UnsupportedEncodingException {
+    jpaTm()
+        .transact(() -> jpaTm().getEntityManager().createNativeQuery("select 1").getSingleResult());
+    assertThat(stdoutBuffer.toString(UTF_8.name())).contains("select 1");
+  }
+}

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

@@ -37,7 +37,13 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Unit tests for {@link JpaTransactionManagerImpl}. */
+/**
+ * Unit tests for SQL only APIs defined in {@link JpaTransactionManagerImpl}. Note that the tests
+ * for common APIs in {@link TransactionManager} are added in {@link TransactionManagerTest}.
+ *
+ * <p>TODO(shicong): Remove duplicate tests that covered by TransactionManagerTest by refactoring
+ * the test schema.
+ */
 @RunWith(JUnit4.class)
 public class JpaTransactionManagerImplTest {
 
@@ -62,29 +68,6 @@ public class JpaTransactionManagerImplTest {
           .withEntityClass(TestEntity.class, TestCompoundIdEntity.class)
           .buildUnitTestRule();
 
-  @Test
-  public void inTransaction_returnsCorrespondingResult() {
-    assertThat(jpaTm().inTransaction()).isFalse();
-    jpaTm().transact(() -> assertThat(jpaTm().inTransaction()).isTrue());
-    assertThat(jpaTm().inTransaction()).isFalse();
-  }
-
-  @Test
-  public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
-    assertThrows(IllegalStateException.class, () -> jpaTm().assertInTransaction());
-    jpaTm().transact(() -> jpaTm().assertInTransaction());
-    assertThrows(IllegalStateException.class, () -> jpaTm().assertInTransaction());
-  }
-
-  @Test
-  public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
-    FakeClock txnClock = fakeClock;
-    txnClock.advanceOneMilli();
-    assertThrows(IllegalStateException.class, () -> jpaTm().getTransactionTime());
-    jpaTm().transact(() -> assertThat(jpaTm().getTransactionTime()).isEqualTo(txnClock.nowUtc()));
-    assertThrows(IllegalStateException.class, () -> jpaTm().getTransactionTime());
-  }
-
   @Test
   public void transact_succeeds() {
     assertPersonEmpty();

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

@@ -41,6 +41,8 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.stream.Collectors;
@@ -146,13 +148,15 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
     ImmutableMap properties = PersistenceModule.providesDefaultDatabaseConfigs();
     if (!userProperties.isEmpty()) {
       // If there are user properties, create a new properties object with these added.
-      ImmutableMap.Builder builder = properties.builder();
-      builder.putAll(userProperties);
-      // Forbid Hibernate push to stay consistent with flyway-based schema management.
-      builder.put(Environment.HBM2DDL_AUTO, "none");
-      builder.put(Environment.SHOW_SQL, "true");
-      properties = builder.build();
+      Map<String, String> mergedProperties = Maps.newHashMap();
+      mergedProperties.putAll(properties);
+      mergedProperties.putAll(userProperties);
+      properties = ImmutableMap.copyOf(mergedProperties);
     }
+    // Forbid Hibernate push to stay consistent with flyway-based schema management.
+    checkState(
+        Objects.equals(properties.get(Environment.HBM2DDL_AUTO), "none"),
+        "The HBM2DDL_AUTO property must be 'none'.");
     assertReasonableNumDbConnections();
     emf =
         createEntityManagerFactory(

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

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.model.ofy;
+package google.registry.persistence.transaction;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -23,16 +23,24 @@ import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.model.ImmutableObject;
+import google.registry.model.ofy.DatastoreTransactionManager;
+import google.registry.model.ofy.Ofy;
 import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
 import java.util.NoSuchElementException;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
-public class DatastoreTransactionManagerTest {
+/**
+ * Unit tests for common APIs in {@link DatastoreTransactionManager} and {@link
+ * JpaTransactionManagerImpl}.
+ */
+@DualDatabaseTest
+public class TransactionManagerTest {
 
   private final FakeClock fakeClock = new FakeClock();
 
@@ -51,38 +59,31 @@ public class DatastoreTransactionManagerTest {
           .withClock(fakeClock)
           .withDatastoreAndCloudSql()
           .withOfyTestEntities(TestEntity.class)
+          .withJpaUnitTestEntities(TestEntity.class)
           .build();
 
-  public DatastoreTransactionManagerTest() {}
+  public TransactionManagerTest() {}
 
   @BeforeEach
   public void setUp() {
     inject.setStaticField(Ofy.class, "clock", fakeClock);
   }
 
-  // TODO(mmuller): The tests below are just copy-pasted from JpaTransactionManagerImplTest
-  // (excluding the CompoundId tests and native query tests, which are not relevant to datastore,
-  // and the test methods using "count" which doesn't work for datastore, as well as tests for
-  // functionality that doesn't exist in datastore, like failures based on whether a newly saved or
-  // updated object exists or not). We need to merge these into a single test suite, but first we
-  // should move the JpaUnitTestRule functionality into AppEngineTest and migrate the whole thing
-  // to junit5.
-
-  @Test
+  @TestTemplate
   public void inTransaction_returnsCorrespondingResult() {
     assertThat(tm().inTransaction()).isFalse();
     tm().transact(() -> assertThat(tm().inTransaction()).isTrue());
     assertThat(tm().inTransaction()).isFalse();
   }
 
-  @Test
+  @TestTemplate
   public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
     assertThrows(IllegalStateException.class, () -> tm().assertInTransaction());
     tm().transact(() -> tm().assertInTransaction());
     assertThrows(IllegalStateException.class, () -> tm().assertInTransaction());
   }
 
-  @Test
+  @TestTemplate
   public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
     FakeClock txnClock = fakeClock;
     txnClock.advanceOneMilli();
@@ -91,7 +92,7 @@ public class DatastoreTransactionManagerTest {
     assertThrows(IllegalStateException.class, () -> tm().getTransactionTime());
   }
 
-  @Test
+  @TestTemplate
   public void transact_hasNoEffectWithPartialSuccess() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     assertThrows(
@@ -106,7 +107,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
   }
 
-  @Test
+  @TestTemplate
   public void transact_reusesExistingTransaction() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -115,7 +116,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isTrue();
   }
 
-  @Test
+  @TestTemplate
   public void saveNew_succeeds() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -126,7 +127,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(tm().transact(() -> tm().load(theEntity.key()))).isEqualTo(theEntity);
   }
 
-  @Test
+  @TestTemplate
   public void saveAllNew_succeeds() {
     moreEntities.forEach(
         entity -> assertThat(tm().transact(() -> tm().checkExists(entity))).isFalse());
@@ -137,7 +138,7 @@ public class DatastoreTransactionManagerTest {
         entity -> assertThat(tm().transact(() -> tm().checkExists(entity))).isTrue());
   }
 
-  @Test
+  @TestTemplate
   public void saveNewOrUpdate_persistsNewEntity() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -148,7 +149,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(tm().transact(() -> tm().load(theEntity.key()))).isEqualTo(theEntity);
   }
 
-  @Test
+  @TestTemplate
   public void saveNewOrUpdate_updatesExistingEntity() {
     fakeClock.advanceOneMilli();
     tm().transact(() -> tm().saveNew(theEntity));
@@ -163,7 +164,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(persisted.data).isEqualTo("bar");
   }
 
-  @Test
+  @TestTemplate
   public void saveNewOrUpdateAll_succeeds() {
     moreEntities.forEach(
         entity -> assertThat(tm().transact(() -> tm().checkExists(entity))).isFalse());
@@ -174,13 +175,16 @@ public class DatastoreTransactionManagerTest {
         entity -> assertThat(tm().transact(() -> tm().checkExists(entity))).isTrue());
   }
 
-  @Test
+  @TestTemplate
   public void update_succeeds() {
     fakeClock.advanceOneMilli();
     tm().transact(() -> tm().saveNew(theEntity));
     fakeClock.advanceOneMilli();
     TestEntity persisted =
-        tm().transact(() -> tm().load(VKey.createOfy(TestEntity.class, Key.create(theEntity))));
+        tm().transact(
+                () ->
+                    tm().load(
+                            VKey.create(TestEntity.class, theEntity.name, Key.create(theEntity))));
     fakeClock.advanceOneMilli();
     assertThat(persisted.data).isEqualTo("foo");
     theEntity.data = "bar";
@@ -190,7 +194,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(persisted.data).isEqualTo("bar");
   }
 
-  @Test
+  @TestTemplate
   public void load_succeeds() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -201,7 +205,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(persisted.data).isEqualTo("foo");
   }
 
-  @Test
+  @TestTemplate
   public void load_throwsOnMissingElement() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -209,7 +213,7 @@ public class DatastoreTransactionManagerTest {
         NoSuchElementException.class, () -> tm().transact(() -> tm().load(theEntity.key())));
   }
 
-  @Test
+  @TestTemplate
   public void maybeLoad_succeeds() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -220,14 +224,14 @@ public class DatastoreTransactionManagerTest {
     assertThat(persisted.data).isEqualTo("foo");
   }
 
-  @Test
+  @TestTemplate
   public void maybeLoad_nonExistentObject() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
     assertThat(tm().transact(() -> tm().maybeLoad(theEntity.key())).isPresent()).isFalse();
   }
 
-  @Test
+  @TestTemplate
   public void delete_succeeds() {
     fakeClock.advanceOneMilli();
     tm().transact(() -> tm().saveNew(theEntity));
@@ -239,7 +243,7 @@ public class DatastoreTransactionManagerTest {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
   }
 
-  @Test
+  @TestTemplate
   public void delete_returnsZeroWhenNoEntity() {
     assertThat(tm().transact(() -> tm().checkExists(theEntity))).isFalse();
     fakeClock.advanceOneMilli();
@@ -249,8 +253,9 @@ public class DatastoreTransactionManagerTest {
   }
 
   @Entity(name = "TestEntity")
+  @javax.persistence.Entity(name = "TestEntity")
   private static class TestEntity extends ImmutableObject {
-    @Id private String name;
+    @Id @javax.persistence.Id private String name;
 
     private String data;
 
@@ -262,7 +267,7 @@ public class DatastoreTransactionManagerTest {
     }
 
     public VKey<TestEntity> key() {
-      return VKey.createOfy(TestEntity.class, Key.create(this));
+      return VKey.create(TestEntity.class, name, Key.create(this));
     }
   }
 }

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

@@ -16,6 +16,7 @@ package google.registry.schema.integration;
 
 import static com.google.common.truth.Truth.assert_;
 
+import google.registry.model.billing.BillingEventTest;
 import google.registry.model.contact.ContactResourceTest;
 import google.registry.model.domain.DomainBaseSqlTest;
 import google.registry.model.registry.RegistryLockDaoTest;
@@ -68,6 +69,7 @@ import org.junit.runner.RunWith;
 @SelectClasses({
   // BeforeSuiteTest must be the first entry. See class javadoc for details.
   BeforeSuiteTest.class,
+  BillingEventTest.class,
   ClaimsListDaoTest.class,
   ContactResourceTest.class,
   CursorDaoTest.class,

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

@@ -17,6 +17,8 @@ package google.registry.testing;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.testing.DatastoreHelper.persistSimpleResources;
+import static google.registry.testing.DualDatabaseTestInvocationContextProvider.injectTmForDualDatabaseTest;
+import static google.registry.testing.DualDatabaseTestInvocationContextProvider.restoreTmAfterDualDatabaseTest;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.json.XML.toJSONObject;
@@ -45,6 +47,7 @@ import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestRule;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageRule;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
 import google.registry.util.Clock;
 import java.io.ByteArrayInputStream;
 import java.io.File;
@@ -118,8 +121,11 @@ public final class AppEngineRule extends ExternalResource
    */
   JpaIntegrationWithCoverageExtension jpaIntegrationWithCoverageExtension = null;
 
+  JpaUnitTestRule jpaUnitTestRule;
+
   private boolean withDatastoreAndCloudSql;
   private boolean enableJpaEntityCoverageCheck;
+  private boolean withJpaUnitTest;
   private boolean withLocalModules;
   private boolean withTaskQueue;
   private boolean withUserService;
@@ -131,12 +137,14 @@ public final class AppEngineRule extends ExternalResource
 
   // Test Objectify entity classes to be used with this AppEngineRule instance.
   private ImmutableList<Class<?>> ofyTestEntities;
+  private ImmutableList<Class<?>> jpaTestEntities;
 
   /** Builder for {@link AppEngineRule}. */
   public static class Builder {
 
     private AppEngineRule rule = new AppEngineRule();
-    private ImmutableList.Builder<Class<?>> ofyTestEntities = new ImmutableList.Builder();
+    private ImmutableList.Builder<Class<?>> ofyTestEntities = new ImmutableList.Builder<>();
+    private ImmutableList.Builder<Class<?>> jpaTestEntities = new ImmutableList.Builder<>();
 
     /** Turn on the Datastore service and the Cloud SQL service. */
     public Builder withDatastoreAndCloudSql() {
@@ -205,11 +213,24 @@ public final class AppEngineRule extends ExternalResource
       return this;
     }
 
+    public Builder withJpaUnitTestEntities(Class<?>... entities) {
+      jpaTestEntities.add(entities);
+      rule.withJpaUnitTest = true;
+      return this;
+    }
+
     public AppEngineRule build() {
       checkState(
           !rule.enableJpaEntityCoverageCheck || rule.withDatastoreAndCloudSql,
           "withJpaEntityCoverageCheck enabled without Cloud SQL");
+      checkState(
+          !rule.withJpaUnitTest || rule.withDatastoreAndCloudSql,
+          "withJpaUnitTestEntities enabled without Cloud SQL");
+      checkState(
+          !rule.withJpaUnitTest || !rule.enableJpaEntityCoverageCheck,
+          "withJpaUnitTestEntities cannot be set when enableJpaEntityCoverageCheck");
       rule.ofyTestEntities = this.ofyTestEntities.build();
+      rule.jpaTestEntities = this.jpaTestEntities.build();
       return rule;
     }
   }
@@ -328,11 +349,18 @@ public final class AppEngineRule extends ExternalResource
       if (enableJpaEntityCoverageCheck) {
         jpaIntegrationWithCoverageExtension = builder.buildIntegrationWithCoverageExtension();
         jpaIntegrationWithCoverageExtension.beforeEach(context);
+      } else if (withJpaUnitTest) {
+        jpaUnitTestRule =
+            builder
+                .withEntityClass(jpaTestEntities.toArray(new Class[jpaTestEntities.size()]))
+                .buildUnitTestRule();
+        jpaUnitTestRule.before();
       } else {
         jpaIntegrationTestRule = builder.buildIntegrationTestRule();
         jpaIntegrationTestRule.before();
       }
     }
+    injectTmForDualDatabaseTest(context);
   }
 
   /** Called after each test method. JUnit 5 only. */
@@ -341,11 +369,14 @@ public final class AppEngineRule extends ExternalResource
     if (withDatastoreAndCloudSql) {
       if (enableJpaEntityCoverageCheck) {
         jpaIntegrationWithCoverageExtension.afterEach(context);
+      } else if (withJpaUnitTest) {
+        jpaUnitTestRule.after();
       } else {
         jpaIntegrationTestRule.after();
       }
     }
     after();
+    restoreTmAfterDualDatabaseTest(context);
   }
 
   /**
@@ -560,4 +591,8 @@ public final class AppEngineRule extends ExternalResource
             makeRegistrarContact2(),
             makeRegistrarContact3()));
   }
+
+  boolean isWithDatastoreAndCloudSql() {
+    return withDatastoreAndCloudSql;
+  }
 }

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

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

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

@@ -0,0 +1,125 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.testing;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.google.common.collect.ImmutableList;
+import google.registry.persistence.transaction.TransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
+import java.lang.reflect.Field;
+import java.util.List;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.Extension;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
+import org.junit.jupiter.api.extension.TestInstancePostProcessor;
+import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
+import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider;
+
+/**
+ * Implementation of {@link TestTemplateInvocationContextProvider} to execute tests against
+ * different database. The test annotated with {@link TestTemplate} will be executed twice against
+ * Datastore and PostgresQL respectively.
+ */
+class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocationContextProvider {
+  private static final Namespace NAMESPACE =
+      Namespace.create(DualDatabaseTestInvocationContextProvider.class);
+  private static final String INJECTED_TM_SUPPLIER_KEY = "injected_tm_supplier_key";
+  private static final String ORIGINAL_TM_KEY = "original_tm_key";
+
+  @Override
+  public boolean supportsTestTemplate(ExtensionContext context) {
+    return true;
+  }
+
+  @Override
+  public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(
+      ExtensionContext context) {
+    return Stream.of(
+        createInvocationContext("Test Datastore", TransactionManagerFactory::ofyTm),
+        createInvocationContext("Test PostgreSQL", TransactionManagerFactory::jpaTm));
+  }
+
+  private TestTemplateInvocationContext createInvocationContext(
+      String name, Supplier<? extends TransactionManager> tmSupplier) {
+    return new TestTemplateInvocationContext() {
+      @Override
+      public String getDisplayName(int invocationIndex) {
+        return name;
+      }
+
+      @Override
+      public List<Extension> getAdditionalExtensions() {
+        return ImmutableList.of(new DatabaseSwitchInvocationContext(tmSupplier));
+      }
+    };
+  }
+
+  private static class DatabaseSwitchInvocationContext implements TestInstancePostProcessor {
+
+    private Supplier<? extends TransactionManager> tmSupplier;
+
+    private DatabaseSwitchInvocationContext(Supplier<? extends TransactionManager> tmSupplier) {
+      this.tmSupplier = tmSupplier;
+    }
+
+    @Override
+    public void postProcessTestInstance(Object testInstance, ExtensionContext context)
+        throws Exception {
+      List<Field> appEngineRuleFields =
+          Stream.of(testInstance.getClass().getFields())
+              .filter(field -> field.getType().isAssignableFrom(AppEngineRule.class))
+              .collect(toImmutableList());
+      if (appEngineRuleFields.size() != 1) {
+        throw new IllegalStateException(
+            "@DualDatabaseTest test must have only 1 AppEngineRule field");
+      }
+      appEngineRuleFields.get(0).setAccessible(true);
+      AppEngineRule appEngineRule = (AppEngineRule) appEngineRuleFields.get(0).get(testInstance);
+      if (!appEngineRule.isWithDatastoreAndCloudSql()) {
+        throw new IllegalStateException(
+            "AppEngineRule in @DualDatabaseTest test must set withDatastoreAndCloudSql()");
+      }
+      context.getStore(NAMESPACE).put(INJECTED_TM_SUPPLIER_KEY, tmSupplier);
+    }
+  }
+
+  static void injectTmForDualDatabaseTest(ExtensionContext context) {
+    if (isDualDatabaseTest(context)) {
+      context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm());
+      Supplier<? extends TransactionManager> tmSupplier =
+          (Supplier<? extends TransactionManager>)
+              context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY);
+      TransactionManagerFactory.setTm(tmSupplier.get());
+    }
+  }
+
+  static void restoreTmAfterDualDatabaseTest(ExtensionContext context) {
+    if (isDualDatabaseTest(context)) {
+      TransactionManager original =
+          (TransactionManager) context.getStore(NAMESPACE).get(ORIGINAL_TM_KEY);
+      TransactionManagerFactory.setTm(original);
+    }
+  }
+
+  private static boolean isDualDatabaseTest(ExtensionContext context) {
+    Object testInstance = context.getTestInstance().orElseThrow(RuntimeException::new);
+    return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class);
+  }
+}

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

@@ -19,7 +19,7 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.io.Resources;
 import google.registry.testing.DatastoreEntityExtension;
-import google.registry.tools.LevelDbFileBuilder.Property;
+import google.registry.tools.EntityWrapper.Property;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
@@ -71,30 +71,38 @@ public class CompareDbBackupsTest {
     // Create two directories corresponding to data dumps.
     File dump1 = tempFs.newFolder("dump1");
     LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(dump1, "output-data1"));
-    builder.addEntityProto(
-        BASE_ID,
-        Property.create("eeny", 100L),
-        Property.create("meeny", 200L),
-        Property.create("miney", 300L));
-    builder.addEntityProto(
-        BASE_ID + 1,
-        Property.create("moxey", 100L),
-        Property.create("minney", 200L),
-        Property.create("motz", 300L));
+    builder.addEntity(
+        EntityWrapper.from(
+                BASE_ID,
+                Property.create("eeny", 100L),
+                Property.create("meeny", 200L),
+                Property.create("miney", 300L))
+            .getEntity());
+    builder.addEntity(
+        EntityWrapper.from(
+                BASE_ID + 1,
+                Property.create("moxey", 100L),
+                Property.create("minney", 200L),
+                Property.create("motz", 300L))
+            .getEntity());
     builder.build();
 
     File dump2 = tempFs.newFolder("dump2");
     builder = new LevelDbFileBuilder(new File(dump2, "output-data2"));
-    builder.addEntityProto(
-        BASE_ID + 1,
-        Property.create("moxey", 100L),
-        Property.create("minney", 200L),
-        Property.create("motz", 300L));
-    builder.addEntityProto(
-        BASE_ID + 2,
-        Property.create("blutzy", 100L),
-        Property.create("fishey", 200L),
-        Property.create("strutz", 300L));
+    builder.addEntity(
+        EntityWrapper.from(
+                BASE_ID + 1,
+                Property.create("moxey", 100L),
+                Property.create("minney", 200L),
+                Property.create("motz", 300L))
+            .getEntity());
+    builder.addEntity(
+        EntityWrapper.from(
+                BASE_ID + 2,
+                Property.create("blutzy", 100L),
+                Property.create("fishey", 200L),
+                Property.create("strutz", 300L))
+            .getEntity());
     builder.build();
 
     CompareDbBackups.main(new String[] {dump1.getCanonicalPath(), dump2.getCanonicalPath()});

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

@@ -28,7 +28,7 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
-public final class ComparableEntityTest {
+public final class EntityWrapperTest {
 
   private static final String TEST_ENTITY_KIND = "TestEntity";
   private static final int ARBITRARY_KEY_ID = 1001;
@@ -63,13 +63,13 @@ public final class ComparableEntityTest {
     Entity e2 = EntityTranslator.createFromPb(proto2);
 
     // Ensure that we have a normalized representation.
-    ComparableEntity ce1 = new ComparableEntity(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    EntityWrapper ce1 = new EntityWrapper(e1);
+    EntityWrapper ce2 = new EntityWrapper(e2);
     assertThat(ce1).isEqualTo(ce2);
     assertThat(ce1.hashCode()).isEqualTo(ce2.hashCode());
 
     // Ensure that the original entity is equal.
-    assertThat(new ComparableEntity(entity)).isEqualTo(ce1);
+    assertThat(new EntityWrapper(entity)).isEqualTo(ce1);
   }
 
   @Test
@@ -90,8 +90,8 @@ public final class ComparableEntityTest {
     Entity e1 = EntityTranslator.createFromPb(proto1);
     Entity e2 = EntityTranslator.createFromPb(proto2);
 
-    ComparableEntity ce1 = new ComparableEntity(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    EntityWrapper ce1 = new EntityWrapper(e1);
+    EntityWrapper ce2 = new EntityWrapper(e2);
     assertThat(e1).isEqualTo(e2); // The keys should still be the same.
     assertThat(ce1).isNotEqualTo(ce2);
     assertThat(ce1.hashCode()).isNotEqualTo(ce2.hashCode());
@@ -108,15 +108,15 @@ public final class ComparableEntityTest {
     Entity e1 = EntityTranslator.createFromPb(proto1);
     Entity e2 = EntityTranslator.createFromPb(proto2);
 
-    ComparableEntity ce1 = new ComparableEntity(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    EntityWrapper ce1 = new EntityWrapper(e1);
+    EntityWrapper ce2 = new EntityWrapper(e2);
     assertThat(ce1).isNotEqualTo(ce2);
     assertThat(ce1.hashCode()).isNotEqualTo(ce2.hashCode());
   }
 
   @Test
   public void testComparisonAgainstNonComparableEntities() {
-    ComparableEntity ce = new ComparableEntity(new Entity(TEST_ENTITY_KIND, ARBITRARY_KEY_ID));
+    EntityWrapper ce = new EntityWrapper(new Entity(TEST_ENTITY_KIND, ARBITRARY_KEY_ID));
     // Note: this has to be "isNotEqualTo()" and not isNotNull() because we want to test the
     // equals() method and isNotNull() just checks for "ce != null".
     assertThat(ce).isNotEqualTo(null);

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

@@ -19,7 +19,6 @@ import static google.registry.tools.LevelDbLogReader.HEADER_SIZE;
 
 import com.google.appengine.api.datastore.Entity;
 import com.google.appengine.api.datastore.EntityTranslator;
-import com.google.auto.value.AutoValue;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import google.registry.tools.LevelDbLogReader.ChunkType;
 import java.io.File;
@@ -28,30 +27,19 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 
 /** Utility class for building a leveldb logfile. */
-final class LevelDbFileBuilder {
-  private static final String TEST_ENTITY_KIND = "TestEntity";
-
+public final class LevelDbFileBuilder {
   private final FileOutputStream out;
   private byte[] currentBlock = new byte[BLOCK_SIZE];
 
   // Write position in the current block.
   private int currentPos = 0;
 
-  LevelDbFileBuilder(File file) throws FileNotFoundException {
+  public LevelDbFileBuilder(File file) throws FileNotFoundException {
     out = new FileOutputStream(file);
   }
 
-  /**
-   * Adds a record containing a new entity protobuf to the file.
-   *
-   * <p>Returns the ComparableEntity object rather than "this" so that we can check for the presence
-   * of the entity in the result set.
-   */
-  ComparableEntity addEntityProto(int id, Property... properties) throws IOException {
-    Entity entity = new Entity(TEST_ENTITY_KIND, id);
-    for (Property prop : properties) {
-      entity.setProperty(prop.name(), prop.value());
-    }
+  /** Adds an {@link Entity Datastore Entity object} to the leveldb log file. */
+  LevelDbFileBuilder addEntity(Entity entity) throws IOException {
     EntityProto proto = EntityTranslator.convertToPb(entity);
     byte[] protoBytes = proto.toByteArray();
     if (protoBytes.length > BLOCK_SIZE - (currentPos + HEADER_SIZE)) {
@@ -61,7 +49,7 @@ final class LevelDbFileBuilder {
     }
 
     currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes);
-    return new ComparableEntity(entity);
+    return this;
   }
 
   /** Writes all remaining data and closes the block. */
@@ -69,15 +57,4 @@ final class LevelDbFileBuilder {
     out.write(currentBlock);
     out.close();
   }
-
-  @AutoValue
-  abstract static class Property {
-    static Property create(String name, Object value) {
-      return new AutoValue_LevelDbFileBuilder_Property(name, value);
-    }
-
-    abstract String name();
-
-    abstract Object value();
-  }
 }

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

@@ -15,15 +15,18 @@
 package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.api.datastore.Entity;
 import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.common.collect.ImmutableList;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
+import google.registry.model.contact.ContactResource;
 import google.registry.testing.AppEngineRule;
-import google.registry.tools.LevelDbFileBuilder.Property;
+import google.registry.testing.DatastoreHelper;
+import google.registry.tools.EntityWrapper.Property;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import org.junit.Rule;
 import org.junit.Test;
@@ -46,22 +49,18 @@ public class LevelDbFileBuilderTest {
     File subdir = tempFs.newFolder("folder");
     File logFile = new File(subdir, "testfile");
     LevelDbFileBuilder builder = new LevelDbFileBuilder(logFile);
-    ComparableEntity entity =
-        builder.addEntityProto(
+    EntityWrapper entity =
+        EntityWrapper.from(
             BASE_ID, Property.create("first", 100L), Property.create("second", 200L));
+    builder.addEntity(entity.getEntity());
     builder.build();
 
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new FileInputStream(logFile));
-
-    ImmutableList<byte[]> records = reader.getRecords();
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
     assertThat(records).hasSize(1);
 
     // Reconstitute an entity, make sure that what we've got is the same as what we started with.
-    EntityProto proto = new EntityProto();
-    proto.parseFrom(records.get(0));
-    Entity materializedEntity = EntityTranslator.createFromPb(proto);
-    assertThat(new ComparableEntity(materializedEntity)).isEqualTo(entity);
+    Entity materializedEntity = rawRecordToEntity(records.get(0));
+    assertThat(new EntityWrapper(materializedEntity)).isEqualTo(entity);
   }
 
   @Test
@@ -72,28 +71,50 @@ public class LevelDbFileBuilderTest {
 
     // Generate enough records to cross a block boundary.  These records end up being around 80
     // bytes, so 1000 works.
-    ImmutableList.Builder<ComparableEntity> originalEntitiesBuilder = new ImmutableList.Builder<>();
+    ImmutableList.Builder<EntityWrapper> originalEntitiesBuilder = new ImmutableList.Builder<>();
     for (int i = 0; i < 1000; ++i) {
-      ComparableEntity entity =
-          builder.addEntityProto(
+      EntityWrapper entity =
+          EntityWrapper.from(
               BASE_ID + i, Property.create("first", 100L), Property.create("second", 200L));
+      builder.addEntity(entity.getEntity());
       originalEntitiesBuilder.add(entity);
     }
     builder.build();
-    ImmutableList<ComparableEntity> originalEntities = originalEntitiesBuilder.build();
+    ImmutableList<EntityWrapper> originalEntities = originalEntitiesBuilder.build();
 
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new FileInputStream(logFile));
-
-    ImmutableList<byte[]> records = reader.getRecords();
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
     assertThat(records).hasSize(1000);
     int index = 0;
     for (byte[] record : records) {
-      EntityProto proto = new EntityProto();
-      proto.parseFrom(record);
-      Entity materializedEntity = EntityTranslator.createFromPb(proto);
-      assertThat(new ComparableEntity(materializedEntity)).isEqualTo(originalEntities.get(index));
+      Entity materializedEntity = rawRecordToEntity(record);
+      assertThat(new EntityWrapper(materializedEntity)).isEqualTo(originalEntities.get(index));
       ++index;
     }
   }
+
+  @Test
+  public void testOfyEntityWrite() throws Exception {
+    File subdir = tempFs.newFolder("folder");
+    File logFile = new File(subdir, "testfile");
+    LevelDbFileBuilder builder = new LevelDbFileBuilder(logFile);
+
+    ContactResource contact = DatastoreHelper.newContactResource("contact");
+    builder.addEntity(tm().transact(() -> ofy().save().toEntity(contact)));
+    builder.build();
+
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
+    assertThat(records).hasSize(1);
+    ContactResource ofyEntity = rawRecordToOfyEntity(records.get(0), ContactResource.class);
+    assertThat(ofyEntity.getContactId()).isEqualTo(contact.getContactId());
+  }
+
+  private static Entity rawRecordToEntity(byte[] record) {
+    EntityProto proto = new EntityProto();
+    proto.parseFrom(record);
+    return EntityTranslator.createFromPb(proto);
+  }
+
+  private static <T> T rawRecordToOfyEntity(byte[] record, Class<T> expectedType) {
+    return expectedType.cast(ofy().load().fromEntity(rawRecordToEntity(record)));
+  }
 }

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

@@ -17,6 +17,10 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.tools.LevelDbUtil.MAX_RECORD;
 import static google.registry.tools.LevelDbUtil.addRecord;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+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 com.google.common.primitives.Bytes;
@@ -24,12 +28,9 @@ import google.registry.tools.LevelDbLogReader.ChunkType;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.List;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.Test;
 
-/** LevelDbLogReader tests. */
-@RunWith(JUnit4.class)
+/** Unit tests of {@link LevelDbLogReader}. */
 public final class LevelDbLogReaderTest {
 
   // Size of the test record.  Any value < 256 will do.
@@ -54,28 +55,23 @@ public final class LevelDbLogReaderTest {
   @Test
   public void testSimpleBlock() throws IOException {
     TestBlock block = makeBlockOfRepeatingBytes(0);
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new ByteArrayInputStream(block.data));
-    ImmutableList<byte[]> records = reader.getRecords();
-    assertThat(records).hasSize(block.recordCount);
+    assertThat(readIncrementally(block.data)).hasSize(block.recordCount);
   }
 
   @Test
   public void testLargeRecord() throws IOException {
-    byte[] block = new byte[LevelDbLogReader.BLOCK_SIZE];
-    addRecord(block, 0, ChunkType.FIRST, MAX_RECORD, (byte) 1);
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new ByteArrayInputStream(block));
-    assertThat(reader.getRecords()).isEmpty();
+    byte[] block0 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block0, 0, ChunkType.FIRST, MAX_RECORD, (byte) 1);
+    assertThat(readIncrementally(block0)).isEmpty();
 
-    addRecord(block, 0, ChunkType.MIDDLE, MAX_RECORD, (byte) 2);
-    reader.readFrom(new ByteArrayInputStream(block));
-    assertThat(reader.getRecords()).isEmpty();
+    byte[] block1 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block1, 0, ChunkType.MIDDLE, MAX_RECORD, (byte) 2);
+    assertThat(readIncrementally(block0, block1)).isEmpty();
 
-    addRecord(block, 0, ChunkType.LAST, MAX_RECORD, (byte) 3);
-    reader.readFrom(new ByteArrayInputStream(block));
+    byte[] block2 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block2, 0, ChunkType.LAST, MAX_RECORD, (byte) 3);
 
-    List<byte[]> records = reader.getRecords();
+    List<byte[]> records = readIncrementally(block0, block1, block2);
     assertThat(records).hasSize(1);
     byte[] record = records.get(0);
 
@@ -95,11 +91,24 @@ public final class LevelDbLogReaderTest {
   public void readFromMultiBlockStream() throws IOException {
     TestBlock block0 = makeBlockOfRepeatingBytes(0);
     TestBlock block1 = makeBlockOfRepeatingBytes(138);
-    ByteArrayInputStream source = new ByteArrayInputStream(Bytes.concat(block0.data, block1.data));
+    assertThat(readIncrementally(block0.data, block1.data))
+        .hasSize(block0.recordCount + block1.recordCount);
+  }
 
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(source);
-    assertThat(reader.getRecords()).hasSize(block0.recordCount + block1.recordCount);
+  @Test
+  void read_noData() throws IOException {
+    assertThat(readIncrementally(new byte[0])).isEmpty();
+  }
+
+  @Test
+  void read_failBadFirstBlock() {
+    assertThrows(IllegalStateException.class, () -> readIncrementally(new byte[1]));
+  }
+
+  @Test
+  void read_failBadTrailingBlock() {
+    TestBlock block = makeBlockOfRepeatingBytes(0);
+    assertThrows(IllegalStateException.class, () -> readIncrementally(block.data, new byte[2]));
   }
 
   @Test
@@ -112,6 +121,14 @@ public final class LevelDbLogReaderTest {
     assertThat(ChunkType.fromCode(ChunkType.LAST.getCode())).isEqualTo(ChunkType.LAST);
   }
 
+  @SafeVarargs
+  private static ImmutableList<byte[]> readIncrementally(byte[]... blocks) throws IOException {
+    ByteArrayInputStream source = spy(new ByteArrayInputStream(Bytes.concat(blocks)));
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(source));
+    verify(source, times(1)).close();
+    return records;
+  }
+
   /** Aggregates the bytes of a test block with the record count. */
   private static final class TestBlock {
     final byte[] data;

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

@@ -18,7 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableSet;
 import google.registry.testing.AppEngineRule;
-import google.registry.tools.LevelDbFileBuilder.Property;
+import google.registry.tools.EntityWrapper.Property;
 import java.io.File;
 import java.io.IOException;
 import org.junit.Rule;
@@ -44,39 +44,44 @@ public class RecordAccumulatorTest {
 
     // Note that we need to specify property values as "Long" for property comparisons to work
     // correctly because that's how they are deserialized from protos.
-    ComparableEntity e1 =
-        builder.addEntityProto(
+    EntityWrapper e1 =
+        EntityWrapper.from(
             BASE_ID,
             Property.create("eeny", 100L),
             Property.create("meeny", 200L),
             Property.create("miney", 300L));
-    ComparableEntity e2 =
-        builder.addEntityProto(
+    builder.addEntity(e1.getEntity());
+    EntityWrapper e2 =
+        EntityWrapper.from(
             BASE_ID + 1,
             Property.create("eeny", 100L),
             Property.create("meeny", 200L),
             Property.create("miney", 300L));
+    builder.addEntity(e2.getEntity());
     builder.build();
 
     builder = new LevelDbFileBuilder(new File(subdir, "data2"));
 
     // Duplicate of the record in the other file.
-    builder.addEntityProto(
-        BASE_ID,
-        Property.create("eeny", 100L),
-        Property.create("meeny", 200L),
-        Property.create("miney", 300L));
+    builder.addEntity(
+        EntityWrapper.from(
+                BASE_ID,
+                Property.create("eeny", 100L),
+                Property.create("meeny", 200L),
+                Property.create("miney", 300L))
+            .getEntity());
 
-    ComparableEntity e3 =
-        builder.addEntityProto(
+    EntityWrapper e3 =
+        EntityWrapper.from(
             BASE_ID + 2,
             Property.create("moxy", 100L),
             Property.create("fruvis", 200L),
             Property.create("cortex", 300L));
+    builder.addEntity(e3.getEntity());
     builder.build();
 
-    ImmutableSet<ComparableEntity> entities =
-        new RecordAccumulator().readDirectory(subdir, any -> true).getComparableEntitySet();
+    ImmutableSet<EntityWrapper> entities =
+        RecordAccumulator.readDirectory(subdir, any -> true).getEntityWrapperSet();
     assertThat(entities).containsExactly(e1, e2, e3);
   }
 }

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

@@ -5,11 +5,11 @@ class google.registry.model.UpdateAutoTimestamp {
   org.joda.time.DateTime timestamp;
 }
 class google.registry.model.billing.BillingEvent$Cancellation {
-  @Id long id;
+  @Id java.lang.Long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$OneTime> refOneTime;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$Recurring> refRecurring;
   google.registry.model.billing.BillingEvent$Reason reason;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$OneTime> refOneTime;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> refRecurring;
   java.lang.String clientId;
   java.lang.String targetId;
   java.util.Set<google.registry.model.billing.BillingEvent$Flag> flags;
@@ -26,7 +26,7 @@ enum google.registry.model.billing.BillingEvent$Flag {
   SYNTHETIC;
 }
 class google.registry.model.billing.BillingEvent$Modification {
-  @Id long id;
+  @Id java.lang.Long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
   com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$OneTime> eventRef;
   google.registry.model.billing.BillingEvent$Reason reason;
@@ -38,11 +38,11 @@ class google.registry.model.billing.BillingEvent$Modification {
   org.joda.time.DateTime eventTime;
 }
 class google.registry.model.billing.BillingEvent$OneTime {
-  @Id long id;
+  @Id java.lang.Long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
-  com.googlecode.objectify.Key<? extends google.registry.model.billing.BillingEvent> cancellationMatchingBillingEvent;
-  com.googlecode.objectify.Key<google.registry.model.domain.token.AllocationToken> allocationToken;
   google.registry.model.billing.BillingEvent$Reason reason;
+  google.registry.persistence.VKey<? extends google.registry.model.billing.BillingEvent> cancellationMatchingBillingEvent;
+  google.registry.persistence.VKey<google.registry.model.domain.token.AllocationToken> allocationToken;
   java.lang.Integer periodYears;
   java.lang.String clientId;
   java.lang.String targetId;
@@ -62,7 +62,7 @@ enum google.registry.model.billing.BillingEvent$Reason {
   TRANSFER;
 }
 class google.registry.model.billing.BillingEvent$Recurring {
-  @Id long id;
+  @Id java.lang.Long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
   google.registry.model.billing.BillingEvent$Reason reason;
   google.registry.model.common.TimeOfYear recurrenceTimeOfYear;
@@ -179,11 +179,11 @@ class google.registry.model.domain.DomainBase {
   java.lang.String lastEppUpdateClientId;
   java.lang.String smdId;
   java.lang.String tld;
-  java.util.Set<com.googlecode.objectify.Key<google.registry.model.host.HostResource>> nsHosts;
   java.util.Set<google.registry.model.domain.DesignatedContact> allContacts;
   java.util.Set<google.registry.model.domain.GracePeriod> gracePeriods;
   java.util.Set<google.registry.model.domain.secdns.DelegationSignerData> dsData;
   java.util.Set<google.registry.model.eppcommon.StatusValue> status;
+  java.util.Set<google.registry.persistence.VKey<google.registry.model.host.HostResource>> nsHosts;
   java.util.Set<java.lang.String> subordinateHosts;
   org.joda.time.DateTime deletionTime;
   org.joda.time.DateTime lastEppUpdateTime;