Persist DomainBase.nsHosts VKeys to SQL (#541)

Persist nsHosts in Cloud SQL

Persist the VKey based nameserver hosts field of DomainBase in Cloud SQL with
foreign key constraints.

core/build.gradle

@@ -862,6 +862,8 @@ task standardTest(type: FilteringTest) {
   includeAllTests()
   exclude fragileTestPatterns
   exclude outcastTestPatterns
+  // See SqlIntegrationTestSuite.java
+  exclude '**/*BeforeSuiteTest.*', '**/*AfterSuiteTest.*'
 
   if (rootProject.findProperty("skipDockerIncompatibleTests") == "true") {
     exclude dockerIncompatibleTestPatterns
@@ -894,11 +896,14 @@ createUberJar('nomulus', 'nomulus', 'google.registry.tools.RegistryTool')
 
 // A jar with classes and resources from main sourceSet, excluding internal
 // data. See comments on configurations.nomulus_test above for details.
+// TODO(weiminyu): release process should build this using the public repo to eliminate the need
+// for excludes.
 task nomulusFossJar (type: Jar) {
   archiveBaseName = 'nomulus'
   archiveClassifier = 'public'
   from (project.sourceSets.main.output) {
     exclude 'google/registry/config/files/**'
+    exclude 'google/registry/proxy/config/**'
   }
   from (project.sourceSets.main.output) {
     include 'google/registry/config/files/default-config.yaml'

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

@@ -76,8 +76,10 @@ import javax.annotation.Nullable;
 import javax.persistence.AttributeOverride;
 import javax.persistence.AttributeOverrides;
 import javax.persistence.Column;
+import javax.persistence.Convert;
 import javax.persistence.ElementCollection;
 import javax.persistence.Embedded;
+import javax.persistence.JoinTable;
 import javax.persistence.Transient;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
@@ -141,7 +143,11 @@ public class DomainBase extends EppResource
    */
   @Index @ElementCollection @Transient Set<Key<HostResource>> nsHosts;
 
-  @Ignore @Transient Set<VKey<HostResource>> nsHostVKeys;
+  @Ignore
+  @ElementCollection
+  @JoinTable(name = "DomainHost")
+  @Convert(converter = HostResource.VKeyHostResourceConverter.class)
+  Set<VKey<HostResource>> nsHostVKeys;
 
   /**
    * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.

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

@@ -34,22 +34,25 @@ import google.registry.model.annotations.ReportedOn;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.transfer.TransferData;
 import google.registry.persistence.VKey;
+import google.registry.persistence.converter.VKeyConverter;
 import java.net.InetAddress;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
+import javax.persistence.ElementCollection;
 import org.joda.time.DateTime;
 
 /**
  * A persistable Host resource including mutable and non-mutable fields.
  *
- * <p>A host's {@link TransferData} is stored on the superordinate domain.  Non-subordinate hosts
+ * <p>A host's {@link TransferData} is stored on the superordinate domain. Non-subordinate hosts
  * don't carry a full set of TransferData; all they have is lastTransferTime.
  *
  * @see <a href="https://tools.ietf.org/html/rfc5732">RFC 5732</a>
  */
 @ReportedOn
 @Entity
+@javax.persistence.Entity
 @ExternalMessagingName("host")
 public class HostResource extends EppResource implements ForeignKeyedEppResource {
 
@@ -64,8 +67,7 @@ public class HostResource extends EppResource implements ForeignKeyedEppResource
   String fullyQualifiedHostName;
 
   /** IP Addresses for this host. Can be null if this is an external host. */
-  @Index
-  Set<InetAddress> inetAddresses;
+  @Index @ElementCollection Set<InetAddress> inetAddresses;
 
   /** The superordinate domain of this host, or null if this is an external host. */
   @Index
@@ -210,4 +212,11 @@ public class HostResource extends EppResource implements ForeignKeyedEppResource
       return this;
     }
   }
+
+  public static class VKeyHostResourceConverter extends VKeyConverter<HostResource> {
+    @Override
+    protected Class<HostResource> getAttributeClass() {
+      return HostResource.class;
+    }
+  }
 }

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

@@ -19,6 +19,7 @@ import static com.google.common.base.Strings.emptyToNull;
 import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
+import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.model.Buildable.GenericBuilder;
 import google.registry.model.ImmutableObject;
@@ -83,6 +84,13 @@ public abstract class DomainLabelEntry<T extends Comparable<?>, D extends Domain
           "Label '%s' must be in puny-coded, lower-case form",
           getInstance().label);
       checkArgumentNotNull(getInstance().getValue(), "Value must be specified");
+      // Verify that the label creates a valid SLD if we add a TLD to the end of it.
+      // We require that the label is not already a full domain name including a dot.
+      // Domain name validation is tricky, so let InternetDomainName handle it for us.
+      checkArgument(
+          InternetDomainName.from(getInstance().label + ".tld").parts().size() == 2,
+          "Label %s must not be a multi-level domain name",
+          getInstance().label);
       return super.build();
     }
   }

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

@@ -0,0 +1,196 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence;
+
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.stream.Stream;
+import javax.persistence.Embeddable;
+import javax.persistence.Embedded;
+import javax.persistence.MappedSuperclass;
+import javax.persistence.PostLoad;
+import javax.persistence.PostPersist;
+import javax.persistence.PostRemove;
+import javax.persistence.PostUpdate;
+import javax.persistence.PrePersist;
+import javax.persistence.PreRemove;
+import javax.persistence.PreUpdate;
+
+/**
+ * A listener class to invoke entity callbacks in cases where Hibernate doesn't invoke the callback
+ * as expected.
+ *
+ * <p>JPA defines a few annotations, e.g. {@link PostLoad}, that we can use for the application to
+ * react to certain events that occur inside the persistence mechanism. However, Hibernate only
+ * supports a few basic use cases, e.g. defining a {@link PostLoad} method directly in an {@link
+ * javax.persistence.Entity} class or in an {@link Embeddable} class. If the annotated method is
+ * defined in an {@link Embeddable} class that is a property of another {@link Embeddable} class, or
+ * it is defined in a parent class of the {@link Embeddable} class, Hibernate doesn't invoke it.
+ *
+ * <p>This listener is added in core/src/main/resources/META-INF/orm.xml as a default entity
+ * listener whose annotated methods will be invoked by Hibernate when corresponding events happen.
+ * For example, {@link EntityCallbacksListener#prePersist} will be invoked before the entity is
+ * persisted to the database, then it will recursively invoke any other {@link PrePersist} method
+ * that should be invoked but not handled by Hibernate due to the bug.
+ *
+ * @see <a
+ *     href="https://docs.jboss.org/hibernate/orm/current/userguide/html_single/Hibernate_User_Guide.html#events-jpa-callbacks">JPA
+ *     Callbacks</a>
+ * @see <a href="https://hibernate.atlassian.net/browse/HHH-13316">HHH-13316</a>
+ */
+public class EntityCallbacksListener {
+
+  @PrePersist
+  void prePersist(Object entity) {
+    EntityCallbackExecutor.create(PrePersist.class).execute(entity, entity.getClass());
+  }
+
+  @PreRemove
+  void preRemove(Object entity) {
+    EntityCallbackExecutor.create(PreRemove.class).execute(entity, entity.getClass());
+  }
+
+  @PostPersist
+  void postPersist(Object entity) {
+    EntityCallbackExecutor.create(PostPersist.class).execute(entity, entity.getClass());
+  }
+
+  @PostRemove
+  void postRemove(Object entity) {
+    EntityCallbackExecutor.create(PostRemove.class).execute(entity, entity.getClass());
+  }
+
+  @PreUpdate
+  void preUpdate(Object entity) {
+    EntityCallbackExecutor.create(PreUpdate.class).execute(entity, entity.getClass());
+  }
+
+  @PostUpdate
+  void postUpdate(Object entity) {
+    EntityCallbackExecutor.create(PostUpdate.class).execute(entity, entity.getClass());
+  }
+
+  @PostLoad
+  void postLoad(Object entity) {
+    EntityCallbackExecutor.create(PostLoad.class).execute(entity, entity.getClass());
+  }
+
+  private static class EntityCallbackExecutor {
+    Class<? extends Annotation> callbackType;
+
+    private EntityCallbackExecutor(Class<? extends Annotation> callbackType) {
+      this.callbackType = callbackType;
+    }
+
+    private static EntityCallbackExecutor create(Class<? extends Annotation> callbackType) {
+      return new EntityCallbackExecutor(callbackType);
+    }
+
+    /**
+     * Executes eligible callbacks in {@link Embedded} properties recursively.
+     *
+     * @param entity the Java object of the entity class
+     * @param entityType either the type of the entity or an ancestor type
+     */
+    private void execute(Object entity, Class<?> entityType) {
+      Class<?> parentType = entityType.getSuperclass();
+      if (parentType != null && parentType.isAnnotationPresent(MappedSuperclass.class)) {
+        execute(entity, parentType);
+      }
+
+      findEmbeddedProperties(entity, entityType)
+          .forEach(
+              normalEmbedded -> {
+                // For each normal embedded property, we don't execute its callback method because
+                // it is handled by Hibernate. However, for the embedded property defined in the
+                // entity's parent class, we need to treat it as a nested embedded property and
+                // invoke its callback function.
+                if (entity.getClass().equals(entityType)) {
+                  executeCallbackForNormalEmbeddedProperty(
+                      normalEmbedded, normalEmbedded.getClass());
+                } else {
+                  executeCallbackForNestedEmbeddedProperty(
+                      normalEmbedded, normalEmbedded.getClass());
+                }
+              });
+    }
+
+    private void executeCallbackForNestedEmbeddedProperty(
+        Object nestedEmbeddedObject, Class<?> nestedEmbeddedType) {
+      Class<?> parentType = nestedEmbeddedType.getSuperclass();
+      if (parentType != null && parentType.isAnnotationPresent(MappedSuperclass.class)) {
+        executeCallbackForNestedEmbeddedProperty(nestedEmbeddedObject, parentType);
+      }
+
+      findEmbeddedProperties(nestedEmbeddedObject, nestedEmbeddedType)
+          .forEach(
+              embeddedProperty ->
+                  executeCallbackForNestedEmbeddedProperty(
+                      embeddedProperty, embeddedProperty.getClass()));
+
+      for (Method method : nestedEmbeddedType.getDeclaredMethods()) {
+        if (method.isAnnotationPresent(callbackType)) {
+          invokeMethod(method, nestedEmbeddedObject);
+        }
+      }
+    }
+
+    private void executeCallbackForNormalEmbeddedProperty(
+        Object normalEmbeddedObject, Class<?> normalEmbeddedType) {
+      Class<?> parentType = normalEmbeddedType.getSuperclass();
+      if (parentType != null && parentType.isAnnotationPresent(MappedSuperclass.class)) {
+        executeCallbackForNormalEmbeddedProperty(normalEmbeddedObject, parentType);
+      }
+
+      findEmbeddedProperties(normalEmbeddedObject, normalEmbeddedType)
+          .forEach(
+              embeddedProperty ->
+                  executeCallbackForNestedEmbeddedProperty(
+                      embeddedProperty, embeddedProperty.getClass()));
+    }
+
+    private Stream<Object> findEmbeddedProperties(Object object, Class<?> clazz) {
+      return Arrays.stream(clazz.getDeclaredFields())
+          .filter(
+              field ->
+                  field.isAnnotationPresent(Embedded.class)
+                      || field.getType().isAnnotationPresent(Embeddable.class))
+          .map(field -> getFieldObject(field, object))
+          .filter(Objects::nonNull);
+    }
+
+    private static Object getFieldObject(Field field, Object object) {
+      field.setAccessible(true);
+      try {
+        return field.get(object);
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    private static void invokeMethod(Method method, Object object) {
+      method.setAccessible(true);
+      try {
+        method.invoke(object);
+      } catch (IllegalAccessException | InvocationTargetException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+}

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

@@ -0,0 +1,41 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.converter;
+
+import google.registry.persistence.VKey;
+import javax.annotation.Nullable;
+import javax.persistence.AttributeConverter;
+
+/**
+ * Converts VKey to a string column.
+ */
+public abstract class VKeyConverter<T> implements AttributeConverter<VKey<T>, String> {
+  @Override
+  @Nullable
+  public String convertToDatabaseColumn(@Nullable VKey<T> attribute) {
+    return attribute == null ? null : (String) attribute.getSqlKey();
+  }
+
+  @Override
+  @Nullable
+  public VKey<T> convertToEntityAttribute(@Nullable String dbData) {
+    return dbData == null ? null : VKey.createSql(getAttributeClass(), dbData);
+  }
+
+  /**
+   * Returns the class of the attribute.
+   */
+  protected abstract Class<T> getAttributeClass();
+}

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

@@ -14,6 +14,8 @@
 
 package google.registry.tools;
 
+import static com.google.common.base.Strings.isNullOrEmpty;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.collect.Multimap;
@@ -39,6 +41,11 @@ final class CheckDomainCommand extends NonMutatingEppToolCommand {
       required = true)
   private List<String> mainParameters;
 
+  @Parameter(
+      names = {"-t", "--token"},
+      description = "Allocation token to use in the command, if desired")
+  private String allocationToken;
+
   @Inject
   @Config("registryAdminClientId")
   String registryAdminClientId;
@@ -53,7 +60,11 @@ final class CheckDomainCommand extends NonMutatingEppToolCommand {
     Multimap<String, String> domainNameMap = validateAndGroupDomainNamesByTld(mainParameters);
     for (Collection<String> values : domainNameMap.asMap().values()) {
       setSoyTemplate(DomainCheckSoyInfo.getInstance(), DomainCheckSoyInfo.DOMAINCHECK);
-      addSoyRecord(clientId, new SoyMapData("domainNames", values));
+      SoyMapData soyMapData = new SoyMapData("domainNames", values);
+      if (!isNullOrEmpty(allocationToken)) {
+        soyMapData.put("allocationToken", allocationToken);
+      }
+      addSoyRecord(clientId, soyMapData);
     }
   }
 }

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

@@ -10,4 +10,11 @@
       <basic name="amount" access="FIELD"/>
     </attributes>
   </embeddable>
+  <persistence-unit-metadata>
+    <persistence-unit-defaults>
+      <entity-listeners>
+        <entity-listener class="google.registry.persistence.EntityCallbacksListener" />
+      </entity-listeners>
+    </persistence-unit-defaults>
+  </persistence-unit-metadata>
 </entity-mappings>

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

@@ -20,6 +20,7 @@
       *  Use Hibernate's ServiceRegistry for bootstrapping (not JPA-compliant)
     -->
     <class>google.registry.model.domain.DomainBase</class>
+    <class>google.registry.model.host.HostResource</class>
     <class>google.registry.model.registrar.Registrar</class>
     <class>google.registry.model.registrar.RegistrarContact</class>
     <class>google.registry.schema.domain.RegistryLock</class>
@@ -45,6 +46,7 @@
     <class>google.registry.persistence.converter.StringListConverter</class>
     <class>google.registry.persistence.converter.StringSetConverter</class>
     <class>google.registry.persistence.converter.UpdateAutoTimestampConverter</class>
+    <class>google.registry.persistence.converter.VKeyConverter</class>
     <class>google.registry.persistence.converter.ZonedDateTimeConverter</class>
 
     <!-- TODO(weiminyu): check out application-layer validation. -->

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

@@ -19,6 +19,7 @@
  */
 {template .domaincheck stricthtml="false"}
 {@param domainNames: list<string>}
+{@param? allocationToken: string|null}
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
 <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
   <command>
@@ -39,6 +40,12 @@
         </fee:domain>
       {/for}
       </fee:check>
+      {if isNonnull($allocationToken)}
+        <allocationToken:allocationToken
+            xmlns:allocationToken="urn:ietf:params:xml:ns:allocationToken-1.0">
+          {$allocationToken}
+        </allocationToken:allocationToken>
+      {/if}
     </extension>
     <clTRID>RegistryTool</clTRID>
   </command>

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

@@ -50,12 +50,23 @@ public abstract class EntityTestCase {
 
   protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
 
-  @Rule @RegisterExtension
-  public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+  @Rule @RegisterExtension public final AppEngineRule appEngine;
 
   @Rule @RegisterExtension public InjectRule inject = new InjectRule();
 
+  protected EntityTestCase() {
+    this(false);
+  }
+
+  protected EntityTestCase(boolean enableJpaEntityCheck) {
+    appEngine =
+        AppEngineRule.builder()
+            .withDatastoreAndCloudSql()
+            .enableJpaEntityCoverageCheck(enableJpaEntityCheck)
+            .withClock(fakeClock)
+            .build();
+  }
+
   @Before
   @BeforeEach
   public void injectClock() {

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

@@ -17,32 +17,58 @@ package google.registry.model.domain;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.joda.time.DateTimeZone.UTC;
+import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.model.EntityTestCase;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DesignatedContact.Type;
 import google.registry.model.domain.launch.LaunchNotice;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.host.HostResource;
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import java.sql.SQLException;
 import javax.persistence.EntityManager;
+import javax.persistence.RollbackException;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Verify that we can store/retrieve DomainBase objects from a SQL database. */
-public class DomainBaseSqlTest extends EntityTestCase {
+public class DomainBaseSqlTest {
+
+  protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  @RegisterExtension
+  @Order(value = 1)
+  DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
+
+  @RegisterExtension
+  JpaIntegrationWithCoverageExtension jpa =
+      new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationWithCoverageExtension();
 
   DomainBase domain;
   Key<ContactResource> contactKey;
   Key<ContactResource> contact2Key;
+  VKey<HostResource> host1VKey;
+  HostResource host;
 
   @BeforeEach
   public void setUp() {
     contactKey = Key.create(ContactResource.class, "contact_id1");
     contact2Key = Key.create(ContactResource.class, "contact_id2");
 
+    host1VKey = VKey.createSql(HostResource.class, "host1");
+
     domain =
         new DomainBase.Builder()
             .setFullyQualifiedDomainName("example.com")
@@ -51,6 +77,7 @@ public class DomainBaseSqlTest extends EntityTestCase {
             .setLastEppUpdateTime(fakeClock.nowUtc())
             .setLastEppUpdateClientId("AnotherRegistrar")
             .setLastTransferTime(fakeClock.nowUtc())
+            .setNameservers(host1VKey)
             .setStatusValues(
                 ImmutableSet.of(
                     StatusValue.CLIENT_DELETE_PROHIBITED,
@@ -70,6 +97,12 @@ public class DomainBaseSqlTest extends EntityTestCase {
                 LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
             .setSmdId("smdid")
             .build();
+
+    host =
+        new HostResource.Builder()
+            .setRepoId("host1")
+            .setFullyQualifiedHostName("ns1.example.com")
+            .build();
   }
 
   @Test
@@ -80,6 +113,9 @@ public class DomainBaseSqlTest extends EntityTestCase {
               // Persist the domain.
               EntityManager em = jpaTm().getEntityManager();
               em.persist(domain);
+
+              // Persist the host.
+              em.persist(host);
             });
 
     jpaTm()
@@ -108,4 +144,31 @@ public class DomainBaseSqlTest extends EntityTestCase {
               assertThat(result).isEqualTo(org);
             });
   }
+
+  @Test
+  public void testForeignKeyConstraints() {
+    Exception e =
+        assertThrows(
+            RollbackException.class,
+            () -> {
+              jpaTm()
+                  .transact(
+                      () -> {
+                        // Persist the domain without the associated host object.
+                        EntityManager em = jpaTm().getEntityManager();
+                        em.persist(domain);
+                      });
+            });
+    assertThat(e)
+        .hasCauseThat() // ConstraintViolationException
+        .hasCauseThat() // ConstraintViolationException
+        .hasCauseThat()
+        .isInstanceOf(SQLException.class);
+    assertThat(e)
+        .hasCauseThat() // ConstraintViolationException
+        .hasCauseThat() // ConstraintViolationException
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains("\"DomainHost\" violates foreign key constraint \"fk_domainhost_host");
+  }
 }

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

@@ -41,7 +41,11 @@ public final class RegistryLockDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   @Test
   public void testSaveAndLoad_success() {

core/src/test/java/google/registry/persistence/EntityCallbacksListenerTest.java

@@ -0,0 +1,302 @@
+// 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;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.common.collect.ImmutableSet;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import java.lang.reflect.Method;
+import javax.persistence.Embeddable;
+import javax.persistence.Embedded;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.MappedSuperclass;
+import javax.persistence.PostLoad;
+import javax.persistence.PostPersist;
+import javax.persistence.PostRemove;
+import javax.persistence.PostUpdate;
+import javax.persistence.PrePersist;
+import javax.persistence.PreRemove;
+import javax.persistence.PreUpdate;
+import javax.persistence.Transient;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link EntityCallbacksListener}. */
+@RunWith(JUnit4.class)
+public class EntityCallbacksListenerTest {
+
+  @Rule
+  public final JpaUnitTestRule jpaRule =
+      new JpaTestRules.Builder().withEntityClass(TestEntity.class).buildUnitTestRule();
+
+  @Test
+  public void verifyAllCallbacks_executedExpectedTimes() {
+    TestEntity testPersist = new TestEntity();
+    jpaTm().transact(() -> jpaTm().saveNew(testPersist));
+    checkAll(testPersist, 1, 0, 0, 0);
+
+    TestEntity testUpdate = new TestEntity();
+    TestEntity updated =
+        jpaTm()
+            .transact(
+                () -> {
+                  TestEntity merged = jpaTm().getEntityManager().merge(testUpdate);
+                  merged.foo++;
+                  jpaTm().getEntityManager().flush();
+                  return merged;
+                });
+    // Note that when we get the merged entity, its @PostLoad callbacks are also invoked
+    checkAll(updated, 0, 1, 0, 1);
+
+    TestEntity testLoad =
+        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestEntity.class, "id"))).get();
+    checkAll(testLoad, 0, 0, 0, 1);
+
+    TestEntity testRemove =
+        jpaTm()
+            .transact(
+                () -> {
+                  TestEntity removed = jpaTm().load(VKey.createSql(TestEntity.class, "id")).get();
+                  jpaTm().getEntityManager().remove(removed);
+                  return removed;
+                });
+    checkAll(testRemove, 0, 0, 1, 1);
+  }
+
+  @Test
+  public void verifyAllManagedEntities_haveNoMethodWithEmbedded() {
+    ImmutableSet<Class> violations =
+        PersistenceXmlUtility.getManagedClasses().stream()
+            .filter(clazz -> clazz.isAnnotationPresent(Entity.class))
+            .filter(EntityCallbacksListenerTest::hasMethodAnnotatedWithEmbedded)
+            .collect(toImmutableSet());
+    assertWithMessage(
+            "Found entity classes having methods annotated with @Embedded. EntityCallbacksListener"
+                + " only supports annotating fields with @Embedded.")
+        .that(violations)
+        .isEmpty();
+  }
+
+  @Test
+  public void verifyHasMethodAnnotatedWithEmbedded_work() {
+    assertThat(hasMethodAnnotatedWithEmbedded(ViolationEntity.class)).isTrue();
+  }
+
+  private static boolean hasMethodAnnotatedWithEmbedded(Class<?> entityType) {
+    boolean result = false;
+    Class<?> parentType = entityType.getSuperclass();
+    if (parentType != null && parentType.isAnnotationPresent(MappedSuperclass.class)) {
+      result = hasMethodAnnotatedWithEmbedded(parentType);
+    }
+    for (Method method : entityType.getDeclaredMethods()) {
+      if (method.isAnnotationPresent(Embedded.class)) {
+        result = true;
+        break;
+      }
+    }
+    return result;
+  }
+
+  private static void checkAll(
+      TestEntity testEntity,
+      int expectedPersist,
+      int expectedUpdate,
+      int expectedRemove,
+      int expectedLoad) {
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPostPersist)
+        .isEqualTo(expectedPersist);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPrePersist)
+        .isEqualTo(expectedPersist);
+
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPreUpdate)
+        .isEqualTo(expectedUpdate);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPostUpdate)
+        .isEqualTo(expectedUpdate);
+
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPreRemove)
+        .isEqualTo(expectedRemove);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPostRemove)
+        .isEqualTo(expectedRemove);
+
+    assertThat(testEntity.entityPostLoad).isEqualTo(expectedLoad);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedPostLoad).isEqualTo(expectedLoad);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedNested.entityEmbeddedNestedPostLoad)
+        .isEqualTo(expectedLoad);
+    assertThat(testEntity.entityEmbedded.entityEmbeddedParentPostLoad).isEqualTo(expectedLoad);
+
+    assertThat(testEntity.parentPostLoad).isEqualTo(expectedLoad);
+    assertThat(testEntity.parentEmbedded.parentEmbeddedPostLoad).isEqualTo(expectedLoad);
+    assertThat(testEntity.parentEmbedded.parentEmbeddedNested.parentEmbeddedNestedPostLoad)
+        .isEqualTo(expectedLoad);
+    assertThat(testEntity.parentEmbedded.parentEmbeddedParentPostLoad).isEqualTo(expectedLoad);
+  }
+
+  @Entity(name = "TestEntity")
+  private static class TestEntity extends ParentEntity {
+    @Id String name = "id";
+    int foo = 0;
+
+    @Transient int entityPostLoad = 0;
+
+    @Embedded EntityEmbedded entityEmbedded = new EntityEmbedded();
+
+    @PostLoad
+    void entityPostLoad() {
+      entityPostLoad++;
+    }
+  }
+
+  @Embeddable
+  private static class EntityEmbedded extends EntityEmbeddedParent {
+    @Embedded EntityEmbeddedNested entityEmbeddedNested = new EntityEmbeddedNested();
+
+    @Transient int entityEmbeddedPostLoad = 0;
+
+    String entityEmbedded = "placeholder";
+
+    @PostLoad
+    void entityEmbeddedPrePersist() {
+      entityEmbeddedPostLoad++;
+    }
+  }
+
+  @MappedSuperclass
+  private static class EntityEmbeddedParent {
+    @Transient int entityEmbeddedParentPostLoad = 0;
+
+    String entityEmbeddedParent = "placeholder";
+
+    @PostLoad
+    void entityEmbeddedParentPostLoad() {
+      entityEmbeddedParentPostLoad++;
+    }
+  }
+
+  @Embeddable
+  private static class EntityEmbeddedNested {
+    @Transient int entityEmbeddedNestedPrePersist = 0;
+    @Transient int entityEmbeddedNestedPreRemove = 0;
+    @Transient int entityEmbeddedNestedPostPersist = 0;
+    @Transient int entityEmbeddedNestedPostRemove = 0;
+    @Transient int entityEmbeddedNestedPreUpdate = 0;
+    @Transient int entityEmbeddedNestedPostUpdate = 0;
+    @Transient int entityEmbeddedNestedPostLoad = 0;
+
+    String entityEmbeddedNested = "placeholder";
+
+    @PrePersist
+    void entityEmbeddedNestedPrePersist() {
+      entityEmbeddedNestedPrePersist++;
+    }
+
+    @PreRemove
+    void entityEmbeddedNestedPreRemove() {
+      entityEmbeddedNestedPreRemove++;
+    }
+
+    @PostPersist
+    void entityEmbeddedNestedPostPersist() {
+      entityEmbeddedNestedPostPersist++;
+    }
+
+    @PostRemove
+    void entityEmbeddedNestedPostRemove() {
+      entityEmbeddedNestedPostRemove++;
+    }
+
+    @PreUpdate
+    void entityEmbeddedNestedPreUpdate() {
+      entityEmbeddedNestedPreUpdate++;
+    }
+
+    @PostUpdate
+    void entityEmbeddedNestedPostUpdate() {
+      entityEmbeddedNestedPostUpdate++;
+    }
+
+    @PostLoad
+    void entityEmbeddedNestedPostLoad() {
+      entityEmbeddedNestedPostLoad++;
+    }
+  }
+
+  @MappedSuperclass
+  private static class ParentEntity {
+    @Embedded ParentEmbedded parentEmbedded = new ParentEmbedded();
+    @Transient int parentPostLoad = 0;
+
+    String parentEntity = "placeholder";
+
+    @PostLoad
+    void parentPostLoad() {
+      parentPostLoad++;
+    }
+  }
+
+  @Embeddable
+  private static class ParentEmbedded extends ParentEmbeddedParent {
+    @Transient int parentEmbeddedPostLoad = 0;
+
+    String parentEmbedded = "placeholder";
+
+    @Embedded ParentEmbeddedNested parentEmbeddedNested = new ParentEmbeddedNested();
+
+    @PostLoad
+    void parentEmbeddedPostLoad() {
+      parentEmbeddedPostLoad++;
+    }
+  }
+
+  @Embeddable
+  private static class ParentEmbeddedNested {
+    @Transient int parentEmbeddedNestedPostLoad = 0;
+
+    String parentEmbeddedNested = "placeholder";
+
+    @PostLoad
+    void parentEmbeddedNestedPostLoad() {
+      parentEmbeddedNestedPostLoad++;
+    }
+  }
+
+  @MappedSuperclass
+  private static class ParentEmbeddedParent {
+    @Transient int parentEmbeddedParentPostLoad = 0;
+
+    String parentEmbeddedParent = "placeholder";
+
+    @PostLoad
+    void parentEmbeddedParentPostLoad() {
+      parentEmbeddedParentPostLoad++;
+    }
+  }
+
+  @Entity
+  private static class ViolationEntity {
+
+    @Embedded
+    EntityEmbedded getEntityEmbedded() {
+      return new EntityEmbedded();
+    }
+  }
+}

core/src/test/java/google/registry/persistence/converter/VKeyConverterTest.java

@@ -0,0 +1,79 @@
+// 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 static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import javax.persistence.Convert;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test SQL persistence of VKey. */
+@RunWith(JUnit4.class)
+public class VKeyConverterTest {
+
+  @Rule
+  public final JpaUnitTestRule jpaRule =
+      new JpaTestRules.Builder().withEntityClass(TestEntity.class).buildUnitTestRule();
+
+  public VKeyConverterTest() {}
+
+  @Test
+  public void testRoundTrip() {
+    TestEntity original =
+        new TestEntity("TheRealSpartacus", VKey.createSql(TestEntity.class, "ImSpartacus!"));
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
+
+    TestEntity retrieved =
+        jpaTm().transact(
+            () -> jpaTm().getEntityManager().find(TestEntity.class, "TheRealSpartacus"));
+    assertThat(retrieved.other.getSqlKey()).isEqualTo("ImSpartacus!");
+  }
+
+  static class TestEntityVKeyConverter extends VKeyConverter<TestEntity> {
+
+    @Override
+    protected Class<TestEntity> getAttributeClass() {
+      return TestEntity.class;
+    }
+  }
+
+  @Entity(name = "TestEntity")
+  static class TestEntity {
+    @Id
+    String id;
+
+    // Specifying "@Converter(autoApply = true) on TestEntityVKeyConverter this doesn't seem to
+    // work.
+    @Convert(converter = TestEntityVKeyConverter.class)
+    VKey<TestEntity> other;
+
+    TestEntity(String id, VKey<TestEntity> other) {
+      this.id = id;
+      this.other = other;
+    }
+
+    /** Default constructor, needed for hibernate. */
+    public TestEntity() {}
+  }
+}

core/src/test/java/google/registry/schema/cursor/CursorDaoTest.java

@@ -43,7 +43,11 @@ public class CursorDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   @Test
   public void save_worksSuccessfullyOnNewCursor() {

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

@@ -14,19 +14,22 @@
 
 package google.registry.schema.integration;
 
-import com.google.common.truth.Expect;
+import static com.google.common.truth.Truth.assert_;
+
 import google.registry.model.domain.DomainBaseSqlTest;
 import google.registry.model.registry.RegistryLockDaoTest;
 import google.registry.persistence.transaction.JpaEntityCoverage;
 import google.registry.schema.cursor.CursorDaoTest;
+import google.registry.schema.integration.SqlIntegrationTestSuite.AfterSuiteTest;
+import google.registry.schema.integration.SqlIntegrationTestSuite.BeforeSuiteTest;
 import google.registry.schema.registrar.RegistrarDaoTest;
 import google.registry.schema.server.LockDaoTest;
 import google.registry.schema.tld.PremiumListDaoTest;
 import google.registry.schema.tld.ReservedListDaoTest;
 import google.registry.schema.tmch.ClaimsListDaoTest;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
-import org.junit.ClassRule;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
 import org.junit.platform.runner.JUnitPlatform;
 import org.junit.platform.suite.api.SelectClasses;
 import org.junit.runner.RunWith;
@@ -47,9 +50,23 @@ import org.junit.runner.RunWith;
  * <p>Note that with {@code JpaIntegrationWithCoverageRule}, each method starts with an empty
  * database. Therefore this is not the right place for verifying backward data compatibility in
  * end-to-end functional tests.
+ *
+ * <p>As of April 2020, none of the before/after annotations ({@code BeforeClass} and {@code
+ * AfterClass} in JUnit 4, or {@code BeforeAll} and {@code AfterAll} in JUnit5) work in a test suite
+ * run with {@link JUnitPlatform the current JUnit 5 runner}. However, staying with the JUnit 4
+ * runner would prevent any member tests from migrating to JUnit 5.
+ *
+ * <p>This class uses a hack to work with the current JUnit 5 runner. {@link BeforeSuiteTest} is
+ * added to the front of the suite class list and invokes the suite's setup method, and {@link
+ * AfterSuiteTest} is added to the tail of the suite class list and invokes the suite's teardown
+ * method. This works because the member tests are run in the order they are declared (See {@code
+ * org.junit.platform.engine.support.descriptor.AbstractTestDescriptor#addChild}). Should the
+ * ordering changes in the future, we will only get false alarms.
  */
 @RunWith(JUnitPlatform.class)
 @SelectClasses({
+  // BeforeSuiteTest must be the first entry. See class javadoc for details.
+  BeforeSuiteTest.class,
   ClaimsListDaoTest.class,
   CursorDaoTest.class,
   DomainBaseSqlTest.class,
@@ -57,27 +74,56 @@ import org.junit.runner.RunWith;
   PremiumListDaoTest.class,
   RegistrarDaoTest.class,
   RegistryLockDaoTest.class,
-  ReservedListDaoTest.class
+  ReservedListDaoTest.class,
+  // AfterSuiteTest must be the last entry. See class javadoc for details.
+  AfterSuiteTest.class
 })
 public class SqlIntegrationTestSuite {
 
-  @ClassRule public static final Expect expect = Expect.create();
-
-  @BeforeClass
+  @BeforeAll // Not yet supported in JUnit 5. Called through BeforeSuiteTest.
   public static void initJpaEntityCoverage() {
     JpaEntityCoverage.init();
   }
 
-  @AfterClass
+  @AfterAll // Not yet supported in JUnit 5. Called through AfterSuiteTest.
   public static void checkJpaEntityCoverage() {
-    expect
+    // TODO(weiminyu): collect both assertion errors like Truth's Expect does in JUnit 4.
+    assert_()
         .withMessage("Tests are missing for the following JPA entities:")
         .that(JpaEntityCoverage.getUncoveredEntities())
         .isEmpty();
-    expect
+    assert_()
         .withMessage(
             "The following classes do not test JPA entities. Please remove them from this suite")
         .that(JpaEntityCoverage.getIrrelevantTestClasses())
         .isEmpty();
   }
+
+  /**
+   * Hack for calling {@link SqlIntegrationTestSuite#initJpaEntityCoverage()} before all real tests
+   * in suite. See outer class javadoc for details.
+   *
+   * <p>The 'Test' suffix in class name is required.
+   */
+  static class BeforeSuiteTest {
+
+    @Test
+    void beforeAll() {
+      initJpaEntityCoverage();
+    }
+  }
+
+  /**
+   * Hack for invoking {@link SqlIntegrationTestSuite#checkJpaEntityCoverage()} after all real tests
+   * in suite. See outer class javadoc for details.
+   *
+   * <p>The 'Test' suffix in class name is required.
+   */
+  static class AfterSuiteTest {
+
+    @Test
+    void afterSuite() {
+      checkJpaEntityCoverage();
+    }
+  }
 }

core/src/test/java/google/registry/schema/registrar/RegistrarDaoTest.java

@@ -16,18 +16,35 @@ package google.registry.schema.registrar;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableList;
-import google.registry.model.EntityTestCase;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarAddress;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for persisting {@link Registrar} entities. */
-public class RegistrarDaoTest extends EntityTestCase {
+public class RegistrarDaoTest {
+
+  protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  @RegisterExtension
+  @Order(value = 1)
+  DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
+
+  @RegisterExtension
+  JpaIntegrationWithCoverageExtension jpa =
+      new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationWithCoverageExtension();
 
   private final VKey<Registrar> registrarKey = VKey.createSql(Registrar.class, "registrarId");
 

core/src/test/java/google/registry/schema/tld/PremiumListDaoTest.java

@@ -45,7 +45,11 @@ public class PremiumListDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   private static final ImmutableMap<String, BigDecimal> TEST_PRICES =
       ImmutableMap.of(

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

@@ -41,6 +41,7 @@ import google.registry.model.registrar.Registrar.State;
 import google.registry.model.registrar.RegistrarAddress;
 import google.registry.model.registrar.RegistrarContact;
 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.util.Clock;
@@ -104,10 +105,19 @@ public final class AppEngineRule extends ExternalResource
 
   /** A rule-within-a-rule to provide a temporary folder for AppEngineRule's internal temp files. */
   TemporaryFolder temporaryFolder = new TemporaryFolder();
-  // Sets up a SQL database when running on JUnit 5.
+  /**
+   * Sets up a SQL database when running on JUnit 5. This is for test classes that are not member of
+   * the {@code SqlIntegrationTestSuite}.
+   */
+  JpaIntegrationTestRule jpaIntegrationTestRule = null;
+  /**
+   * Sets up a SQL database when running on JUnit 5 and records the JPA entities tested by each test
+   * class. This is for {@code SqlIntegrationTestSuite} members.
+   */
   JpaIntegrationWithCoverageExtension jpaIntegrationWithCoverageExtension = null;
 
   private boolean withDatastoreAndCloudSql;
+  private boolean enableJpaEntityCoverageCheck;
   private boolean withLocalModules;
   private boolean withTaskQueue;
   private boolean withUserService;
@@ -127,6 +137,14 @@ public final class AppEngineRule extends ExternalResource
       rule.withDatastoreAndCloudSql = true;
       return this;
     }
+    /**
+     * Enables JPA entity coverage check if {@code enabled} is true. This should only be enabled for
+     * members of SqlIntegrationTestSuite.
+     */
+    public Builder enableJpaEntityCoverageCheck(boolean enabled) {
+      rule.enableJpaEntityCoverageCheck = enabled;
+      return this;
+    }
 
     /** Turn on the use of local modules. */
     public Builder withLocalModules() {
@@ -164,6 +182,9 @@ public final class AppEngineRule extends ExternalResource
     }
 
     public AppEngineRule build() {
+      checkState(
+          !rule.enableJpaEntityCoverageCheck || rule.withDatastoreAndCloudSql,
+          "withJpaEntityCoverageCheck enabled without Cloud SQL");
       return rule;
     }
   }
@@ -279,8 +300,13 @@ public final class AppEngineRule extends ExternalResource
       if (clock != null) {
         builder.withClock(clock);
       }
-      jpaIntegrationWithCoverageExtension = builder.buildIntegrationWithCoverageExtension();
-      jpaIntegrationWithCoverageExtension.beforeEach(context);
+      if (enableJpaEntityCoverageCheck) {
+        jpaIntegrationWithCoverageExtension = builder.buildIntegrationWithCoverageExtension();
+        jpaIntegrationWithCoverageExtension.beforeEach(context);
+      } else {
+        jpaIntegrationTestRule = builder.buildIntegrationTestRule();
+        jpaIntegrationTestRule.before();
+      }
     }
   }
 
@@ -288,9 +314,11 @@ public final class AppEngineRule extends ExternalResource
   @Override
   public void afterEach(ExtensionContext context) throws Exception {
     if (withDatastoreAndCloudSql) {
-      checkState(
-          jpaIntegrationWithCoverageExtension != null, "Null jpaIntegrationWithCoverageExtension");
-      jpaIntegrationWithCoverageExtension.afterEach(context);
+      if (enableJpaEntityCoverageCheck) {
+        jpaIntegrationWithCoverageExtension.afterEach(context);
+      } else {
+        jpaIntegrationTestRule.after();
+      }
     }
     after();
   }
@@ -310,7 +338,10 @@ public final class AppEngineRule extends ExternalResource
       if (clock != null) {
         builder.withClock(clock);
       }
-      statement = builder.buildIntegrationWithCoverageRule().apply(base, description);
+      checkState(
+          !enableJpaEntityCoverageCheck,
+          "JUnit4 tests must not enable withJpaEntityCoverageCheck.");
+      statement = builder.buildIntegrationTestRule().apply(base, description);
     }
     return super.apply(statement, description);
   }

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

@@ -0,0 +1,104 @@
+// 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 com.google.apphosting.api.ApiProxy;
+import com.google.apphosting.api.ApiProxy.Environment;
+import java.util.Map;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
+
+/**
+ * Allows instantiation of Datastore {@code Entity entities} without the heavyweight {@code
+ * AppEngineRule}.
+ *
+ * <p>When used together with {@code JpaIntegrationWithCoverageExtension}, this extension must be
+ * registered first. For consistency's sake, it is recommended that the field for this extension be
+ * annotated with {@code @org.junit.jupiter.api.Order(value = 1)}. Please refer to {@link
+ * google.registry.model.domain.DomainBaseSqlTest} for example, and to <a
+ * href="https://junit.org/junit5/docs/current/user-guide/#extensions-registration-programmatic">
+ * JUnit 5 User Guide</a> for details of extension ordering.
+ */
+public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCallback {
+
+  private static final Environment PLACEHOLDER_ENV = new PlaceholderEnvironment();
+
+  @Override
+  public void beforeEach(ExtensionContext context) {
+    ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context) {
+    // Clear the cached instance.
+    ApiProxy.setEnvironmentForCurrentThread(null);
+  }
+
+  private static final class PlaceholderEnvironment implements Environment {
+
+    @Override
+    public String getAppId() {
+      return "PlaceholderAppId";
+    }
+
+    @Override
+    public Map<String, Object> getAttributes() {
+      return ImmutableMap.of();
+    }
+
+    @Override
+    public String getModuleId() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public String getVersionId() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public String getEmail() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean isLoggedIn() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean isAdmin() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public String getAuthDomain() {
+      throw new UnsupportedOperationException();
+    }
+
+    @SuppressWarnings("deprecation")
+    @Override
+    public String getRequestNamespace() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public long getRemainingMillis() {
+      throw new UnsupportedOperationException();
+    }
+  }
+}

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

@@ -76,6 +76,12 @@ public class CheckDomainCommandTest extends EppToolCommandTestCase<CheckDomainCo
     eppVerifier.expectDryRun().expectClientId("adminreg").verifySent("domain_check_fee.xml");
   }
 
+  @Test
+  public void testSuccess_allocationToken_reserved() throws Exception {
+    runCommand("--client=NewRegistrar", "--token=abc123", "example.tld");
+    eppVerifier.expectDryRun().verifySent("domain_check_fee_allocationtoken.xml");
+  }
+
   @Test
   public void testFailure_NoMainParameter() {
     assertThrows(ParameterException.class, () -> runCommand("--client=NewRegistrar"));

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

@@ -18,21 +18,19 @@ import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.io.Resources;
-import google.registry.testing.AppEngineRule;
+import google.registry.testing.DatastoreEntityExtension;
 import google.registry.tools.LevelDbFileBuilder.Property;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.IOException;
 import java.io.PrintStream;
 import java.net.URL;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
+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;
 import org.junit.rules.TemporaryFolder;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
-@RunWith(JUnit4.class)
 public class CompareDbBackupsTest {
 
   private static final int BASE_ID = 1001;
@@ -41,20 +39,22 @@ public class CompareDbBackupsTest {
   private final ByteArrayOutputStream stdout = new ByteArrayOutputStream();
   private PrintStream orgStdout;
 
-  @Rule public final TemporaryFolder tempFs = new TemporaryFolder();
+  public final TemporaryFolder tempFs = new TemporaryFolder();
 
-  @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  @RegisterExtension
+  public DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
 
-  @Before
-  public void before() {
+  @BeforeEach
+  public void before() throws IOException {
     orgStdout = System.out;
     System.setOut(new PrintStream(stdout));
+    tempFs.create();
   }
 
-  @After
+  @AfterEach
   public void after() {
     System.setOut(orgStdout);
+    tempFs.delete();
   }
 
   @Test

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

@@ -61,18 +61,43 @@ public abstract class CreateOrUpdateReservedListCommandTestCase<
 
   @Test
   public void testFailure_fileDoesntExist() {
-    assertThrows(
-        ParameterException.class,
-        () ->
-            runCommandForced(
-                "--name=xn--q9jyb4c-blah", "--input=" + reservedTermsPath + "-nonexistent"));
+    assertThat(
+            assertThrows(
+                ParameterException.class,
+                () ->
+                    runCommandForced(
+                        "--name=xn--q9jyb4c_common-reserved",
+                        "--input=" + reservedTermsPath + "-nonexistent")))
+        .hasMessageThat()
+        .contains("-i not found");
   }
 
   @Test
   public void testFailure_fileDoesntParse() {
-    assertThrows(
-        IllegalArgumentException.class,
-        () -> runCommandForced("--name=xn--q9jyb4c-blork", "--input=" + invalidReservedTermsPath));
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    runCommandForced(
+                        "--name=xn--q9jyb4c_common-reserved",
+                        "--input=" + invalidReservedTermsPath)))
+        .hasMessageThat()
+        .contains("No enum constant");
+  }
+
+  @Test
+  public void testFailure_invalidLabel_includesFullDomainName() throws Exception {
+    Files.asCharSink(new File(invalidReservedTermsPath), UTF_8)
+        .write("example.tld,FULLY_BLOCKED\n\n");
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    runCommandForced(
+                        "--name=xn--q9jyb4c_common-reserved",
+                        "--input=" + invalidReservedTermsPath)))
+        .hasMessageThat()
+        .isEqualTo("Label example.tld must not be a multi-level domain name");
   }
 
   google.registry.schema.tld.ReservedList createCloudSqlReservedList(

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

@@ -26,12 +26,18 @@ import com.google.common.collect.ImmutableMap;
 import google.registry.model.registry.label.ReservedList;
 import google.registry.schema.tld.ReservedList.ReservedEntry;
 import google.registry.schema.tld.ReservedListDao;
+import org.junit.Before;
 import org.junit.Test;
 
 /** Unit tests for {@link UpdateReservedListCommand}. */
 public class UpdateReservedListCommandTest extends
     CreateOrUpdateReservedListCommandTestCase<UpdateReservedListCommand> {
 
+  @Before
+  public void setup() {
+    populateInitialReservedListInDatastore(true);
+  }
+
   private void populateInitialReservedListInDatastore(boolean shouldPublish) {
     persistResource(
         new ReservedList.Builder()
@@ -63,7 +69,6 @@ public class UpdateReservedListCommandTest extends
 
   @Test
   public void testSuccess_lastUpdateTime_updatedCorrectly() throws Exception {
-    populateInitialReservedListInDatastore(true);
     ReservedList original = ReservedList.get("xn--q9jyb4c_common-reserved").get();
     runCommandForced("--input=" + reservedTermsPath);
     ReservedList updated = ReservedList.get("xn--q9jyb4c_common-reserved").get();
@@ -90,7 +95,6 @@ public class UpdateReservedListCommandTest extends
   }
 
   private void runSuccessfulUpdateTest(String... args) throws Exception {
-    populateInitialReservedListInDatastore(true);
     runCommandForced(args);
     assertThat(ReservedList.get("xn--q9jyb4c_common-reserved")).isPresent();
     ReservedList reservedList = ReservedList.get("xn--q9jyb4c_common-reserved").get();
@@ -114,7 +118,6 @@ public class UpdateReservedListCommandTest extends
 
   @Test
   public void testSaveToCloudSql_succeeds() throws Exception {
-    populateInitialReservedListInDatastore(true);
     populateInitialReservedListInCloudSql(true);
     runCommandForced("--name=xn--q9jyb4c_common-reserved", "--input=" + reservedTermsPath);
     verifyXnq9jyb4cInDatastore();
@@ -126,7 +129,6 @@ public class UpdateReservedListCommandTest extends
     // Note that, during the dual-write phase, we always save the reserved list to Cloud SQL without
     // checking if there is a list with same name. This is to backfill the existing list in Cloud
     // Datastore when we update it.
-    populateInitialReservedListInDatastore(true);
     runCommandForced("--name=xn--q9jyb4c_common-reserved", "--input=" + reservedTermsPath);
     verifyXnq9jyb4cInDatastore();
     assertThat(ReservedListDao.checkExists("xn--q9jyb4c_common-reserved")).isTrue();

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

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <check>
+      <domain:check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>example.tld</domain:name>
+      </domain:check>
+    </check>
+    <extension>
+      <fee:check xmlns:fee="urn:ietf:params:xml:ns:fee-0.6">
+        <fee:domain>
+          <fee:name>example.tld</fee:name>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+        </fee:domain>
+      </fee:check>
+      <allocationToken:allocationToken
+          xmlns:allocationToken="urn:ietf:params:xml:ns:allocationToken-1.0">
+        abc123
+      </allocationToken:allocationToken>
+    </extension>
+    <clTRID>RegistryTool</clTRID>
+  </command>
+</epp>

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

@@ -0,0 +1,54 @@
+-- Copyright 2020 The Nomulus Authors. All Rights Reserved.
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+create table "DomainHost" (
+   domain_repo_id text not null,
+    ns_host_v_keys text
+);
+
+create table "HostResource" (
+   repo_id text not null,
+    creation_client_id text,
+    creation_time timestamptz,
+    current_sponsor_client_id text,
+    deletion_time timestamptz,
+    last_epp_update_client_id text,
+    last_epp_update_time timestamptz,
+    statuses text[],
+    fully_qualified_host_name text,
+    last_superordinate_change timestamptz,
+    last_transfer_time timestamptz,
+    superordinate_domain bytea,
+    primary key (repo_id)
+);
+
+create table "HostResource_inetAddresses" (
+   host_resource_repo_id text not null,
+    inet_addresses bytea
+);
+
+alter table if exists "DomainHost"
+   add constraint FKfmi7bdink53swivs390m2btxg
+   foreign key (domain_repo_id)
+   references "Domain";
+
+alter table if exists "DomainHost"
+  add constraint FK_DomainHost_host_valid
+  foreign key (ns_host_v_keys)
+  references "HostResource";
+
+alter table if exists "HostResource_inetAddresses"
+   add constraint FK6unwhfkcu3oq6q347fxvpagv
+   foreign key (host_resource_repo_id)
+   references "HostResource";

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

@@ -67,6 +67,11 @@
         primary key (repo_id)
     );
 
+    create table "DomainHost" (
+       domain_repo_id text not null,
+        ns_host_v_keys text
+    );
+
     create table "GracePeriod" (
        id  bigserial not null,
         billing_event_one_time bytea,
@@ -77,6 +82,27 @@
         primary key (id)
     );
 
+    create table "HostResource" (
+       repo_id text not null,
+        creation_client_id text,
+        creation_time timestamptz,
+        current_sponsor_client_id text,
+        deletion_time timestamptz,
+        last_epp_update_client_id text,
+        last_epp_update_time timestamptz,
+        statuses text[],
+        fully_qualified_host_name text,
+        last_superordinate_change timestamptz,
+        last_transfer_time timestamptz,
+        superordinate_domain bytea,
+        primary key (repo_id)
+    );
+
+    create table "HostResource_inetAddresses" (
+       host_resource_repo_id text not null,
+        inet_addresses bytea
+    );
+
     create table "Lock" (
        resource_name text not null,
         tld text not null,
@@ -222,6 +248,16 @@ create index reservedlist_name_idx on "ReservedList" (name);
        foreign key (revision_id) 
        references "ClaimsList";
 
+    alter table if exists "DomainHost" 
+       add constraint FKeq1guccbre1yk3oosgp2io554 
+       foreign key (domain_repo_id) 
+       references "Domain";
+
+    alter table if exists "HostResource_inetAddresses" 
+       add constraint FK6unwhfkcu3oq6q347fxvpagv 
+       foreign key (host_resource_repo_id) 
+       references "HostResource";
+
     alter table if exists "PremiumEntry" 
        add constraint FKo0gw90lpo1tuee56l0nb6y6g5 
        foreign key (revision_id) 

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

@@ -116,6 +116,46 @@ CREATE TABLE public."Domain" (
 );
 
 
+--
+-- Name: DomainHost; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."DomainHost" (
+    domain_repo_id text NOT NULL,
+    ns_host_v_keys text
+);
+
+
+--
+-- Name: HostResource; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."HostResource" (
+    repo_id text NOT NULL,
+    creation_client_id text,
+    creation_time timestamp with time zone,
+    current_sponsor_client_id text,
+    deletion_time timestamp with time zone,
+    last_epp_update_client_id text,
+    last_epp_update_time timestamp with time zone,
+    statuses text[],
+    fully_qualified_host_name text,
+    last_superordinate_change timestamp with time zone,
+    last_transfer_time timestamp with time zone,
+    superordinate_domain bytea
+);
+
+
+--
+-- Name: HostResource_inetAddresses; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."HostResource_inetAddresses" (
+    host_resource_repo_id text NOT NULL,
+    inet_addresses bytea
+);
+
+
 --
 -- Name: Lock; Type: TABLE; Schema: public; Owner: -
 --
@@ -390,6 +430,14 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT "Domain_pkey" PRIMARY KEY (repo_id);
 
 
+--
+-- Name: HostResource HostResource_pkey; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."HostResource"
+    ADD CONSTRAINT "HostResource_pkey" PRIMARY KEY (repo_id);
+
+
 --
 -- Name: Lock Lock_pkey; Type: CONSTRAINT; Schema: public; Owner: -
 --
@@ -562,6 +610,30 @@ ALTER TABLE ONLY public."ClaimsEntry"
     ADD CONSTRAINT fk6sc6at5hedffc0nhdcab6ivuq FOREIGN KEY (revision_id) REFERENCES public."ClaimsList"(revision_id);
 
 
+--
+-- Name: HostResource_inetAddresses fk6unwhfkcu3oq6q347fxvpagv; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."HostResource_inetAddresses"
+    ADD CONSTRAINT fk6unwhfkcu3oq6q347fxvpagv FOREIGN KEY (host_resource_repo_id) REFERENCES public."HostResource"(repo_id);
+
+
+--
+-- Name: DomainHost fk_domainhost_host_valid; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHost"
+    ADD CONSTRAINT fk_domainhost_host_valid FOREIGN KEY (ns_host_v_keys) REFERENCES public."HostResource"(repo_id);
+
+
+--
+-- Name: DomainHost fkfmi7bdink53swivs390m2btxg; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHost"
+    ADD CONSTRAINT fkfmi7bdink53swivs390m2btxg FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
+
+
 --
 -- Name: ReservedEntry fkgq03rk0bt1hb915dnyvd3vnfc; Type: FK CONSTRAINT; Schema: public; Owner: -
 --