Make VKey persist to datastore as a key (#591)

* Make VKey persist to datastore as a key

Convert nsHosts entirely to VKey as a proof-of-concept.

Tested as follows:
    1) Deployed to crash, verified that nameservers were visible for several
       domains (indicating that we are able to load a set of Keys as VKeys)
    2) Updated the set of nameservers for a domain (removing some initial
       hosts) and verified that the changes went through.
    3) Deployed the old version to crash, verified that I was able to retrieve
       the newly saved VKeys as Keys.
    4) Modified the hosts for the same domain (adding back one of the hosts)
       and verified that the change took effect.
    5) Redeployed this change to crash, again updated the nameservers to add
       another host.
    6) Again restored the old version, verified that the new hosts were
       visible.

* Changes in response to review

* Convert to a single VKeyTranslatorFactory instance

* Moved vkey field rename to V25

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/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/ofy/ObjectifyService.java

@@ -131,8 +131,7 @@ public class ObjectifyService {
             new InetAddressTranslatorFactory(),
             new MoneyStringTranslatorFactory(),
             new ReadableInstantUtcTranslatorFactory(),
-            new VKeyTranslatorFactory<ContactResource>(ContactResource.class),
-            new VKeyTranslatorFactory<HostResource>(HostResource.class),
+            new VKeyTranslatorFactory(HostResource.class, ContactResource.class),
             new UpdateAutoTimestampTranslatorFactory())) {
       factory().getTranslators().add(translatorFactory);
     }

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

@@ -14,13 +14,15 @@
 
 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 com.googlecode.objectify.Key;
+import com.google.appengine.api.datastore.Key;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
 import google.registry.persistence.VKey;
-import java.io.UnsupportedEncodingException;
-import java.net.URLDecoder;
-import java.net.URLEncoder;
+import java.util.List;
+import java.util.stream.Stream;
 
 /**
  * Translator factory for VKey.
@@ -28,57 +30,42 @@ 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.
+  private final ImmutableMap<String, Class> classRegistry;
+
+  public VKeyTranslatorFactory(Class... refClasses) {
     super(VKey.class);
-    this.refClass = refClass;
+
+    // Store a registry of all classes by their unqualified name.
+    classRegistry =
+        Stream.of(refClasses)
+            .collect(
+                toImmutableMap(
+                    clazz -> {
+                      List<String> nameComponent = Splitter.on('.').splitToList(clazz.getName());
+                      return nameComponent.get(nameComponent.size() - 1);
+                    },
+                    identity()));
   }
 
   @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 VKey.createOfy(
+            classRegistry.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.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/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/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/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/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/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/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/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/resources/google/registry/model/schema.txt

@@ -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;

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

@@ -0,0 +1,22 @@
+-- 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.
+
+ALTER TABLE "DomainHost" DROP CONSTRAINT FK_DomainHost_host_valid;
+ALTER TABLE "DomainHost" ADD COLUMN ns_hosts text;
+ALTER TABLE "DomainHost" DROP COLUMN ns_host_v_keys;
+ALTER TABLE "DomainHost"
+  ADD CONSTRAINT FK_DomainHost_host_valid
+  FOREIGN KEY (ns_hosts)
+  REFERENCES "HostResource";
+

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

@@ -122,7 +122,7 @@
 
     create table "DomainHost" (
        domain_repo_id text not null,
-        ns_host_v_keys text
+        ns_hosts text
     );
 
     create table "GracePeriod" (

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

@@ -179,7 +179,7 @@ CREATE TABLE public."Domain" (
 
 CREATE TABLE public."DomainHost" (
     domain_repo_id text NOT NULL,
-    ns_host_v_keys text
+    ns_hosts text
 );
 
 
@@ -788,7 +788,7 @@ ALTER TABLE ONLY public."Domain"
 --
 
 ALTER TABLE ONLY public."DomainHost"
-    ADD CONSTRAINT fk_domainhost_host_valid FOREIGN KEY (ns_host_v_keys) REFERENCES public."HostResource"(repo_id);
+    ADD CONSTRAINT fk_domainhost_host_valid FOREIGN KEY (ns_hosts) REFERENCES public."HostResource"(repo_id);
 
 
 --

gradle/wrapper/gradle-wrapper.properties

@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-6.3-all.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-6.4.1-all.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists

gradlew

@@ -82,6 +82,7 @@ esac
 
 CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar
 
+
 # Determine the Java command to use to start the JVM.
 if [ -n "$JAVA_HOME" ] ; then
     if [ -x "$JAVA_HOME/jre/sh/java" ] ; then
@@ -129,6 +130,7 @@ fi
 if [ "$cygwin" = "true" -o "$msys" = "true" ] ; then
     APP_HOME=`cygpath --path --mixed "$APP_HOME"`
     CLASSPATH=`cygpath --path --mixed "$CLASSPATH"`
+    
     JAVACMD=`cygpath --unix "$JAVACMD"`
 
     # We build the pattern for arguments to be converted via cygpath

gradlew.bat

@@ -84,6 +84,7 @@ set CMD_LINE_ARGS=%*
 
 set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
 
+
 @rem Execute Gradle
 "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
 

networking/src/main/java/google/registry/networking/handler/SslServerInitializer.java

@@ -37,6 +37,7 @@ import java.security.cert.CertificateExpiredException;
 import java.security.cert.CertificateNotYetValidException;
 import java.security.cert.X509Certificate;
 import java.util.function.Supplier;
+import javax.net.ssl.SSLSession;
 
 /**
  * Adds a server side SSL handler to the channel pipeline.
@@ -108,9 +109,21 @@ public class SslServerInitializer<C extends Channel> extends ChannelInitializer<
               .addListener(
                   future -> {
                     if (future.isSuccess()) {
+                      SSLSession sslSession = sslHandler.engine().getSession();
                       X509Certificate clientCertificate =
-                          (X509Certificate)
-                              sslHandler.engine().getSession().getPeerCertificates()[0];
+                          (X509Certificate) sslSession.getPeerCertificates()[0];
+                      logger.atInfo().log(
+                          "--SSL Information--\n"
+                              + "Client Certificate Hash: %s\n"
+                              + "SSL Protocol: %s\n"
+                              + "Cipher Suite: %s\n"
+                              + "Not Before: %s\n"
+                              + "Not After: %s\n",
+                          getCertificateHash(clientCertificate),
+                          sslSession.getProtocol(),
+                          sslSession.getCipherSuite(),
+                          clientCertificate.getNotBefore(),
+                          clientCertificate.getNotAfter());
                       try {
                         clientCertificate.checkValidity();
                       } catch (CertificateNotYetValidException | CertificateExpiredException e) {