Fix show-sql which stopped working (#596 )

* Fix show-sql which stopped working Made show-sql property configurable in JpaUnitTestRules. Added a few comments on foreign key constraint behavior.
Do not enqueue update snapshot task if import fails (#578 )
2026-06-09 16:33:02 +00:00 · 2020-05-21 12:20:56 -04:00 · 2020-05-21 11:40:45 -04:00 · 2020-05-20 15:12:27 -04:00 · 2020-05-20 14:24:45 -04:00 · 2020-05-20 12:54:13 -04:00
41 changed files with 1045 additions and 323 deletions
@@ -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
@@ -73,8 +73,10 @@ public class BigqueryPollJobAction implements Runnable {

  @Override
  public void run() {
-    checkJobOutcome();  // Throws a NotModifiedException if the job hasn't completed.
-    if (payload == null || payload.length == 0) {
+    boolean jobOutcome =
+        checkJobOutcome(); // Throws a NotModifiedException if the job hasn't completed.
+    // If the job failed, do not enqueue the next step.
+    if (!jobOutcome || payload == null || payload.length == 0) {
      return;
    }
    // If there is a payload, it's a chained task, so enqueue it.
@@ -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();
    }

@@ -36,8 +36,6 @@ import com.googlecode.objectify.impl.translate.opt.joda.MoneyStringTranslatorFac
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.EntityClasses;
 import google.registry.model.ImmutableObject;
-import google.registry.model.contact.ContactResource;
-import google.registry.model.host.HostResource;
 import google.registry.model.translators.BloomFilterOfStringTranslatorFactory;
 import google.registry.model.translators.CidrAddressBlockTranslatorFactory;
 import google.registry.model.translators.CommitLogRevisionsTranslatorFactory;
@@ -131,8 +129,7 @@ public class ObjectifyService {
            new InetAddressTranslatorFactory(),
            new MoneyStringTranslatorFactory(),
            new ReadableInstantUtcTranslatorFactory(),
-            new VKeyTranslatorFactory<ContactResource>(ContactResource.class),
-            new VKeyTranslatorFactory<HostResource>(HostResource.class),
+            new VKeyTranslatorFactory(),
            new UpdateAutoTimestampTranslatorFactory())) {
      factory().getTranslators().add(translatorFactory);
    }
@@ -14,13 +14,14 @@

 package google.registry.model.translators;

-import static java.nio.charset.StandardCharsets.UTF_8;
+import static com.google.common.base.Functions.identity;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.model.EntityClasses.ALL_CLASSES;

-import com.googlecode.objectify.Key;
+import com.google.appengine.api.datastore.Key;
+import com.google.common.collect.ImmutableMap;
+import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.persistence.VKey;
-import java.io.UnsupportedEncodingException;
-import java.net.URLDecoder;
-import java.net.URLEncoder;

 /**
 * Translator factory for VKey.
@@ -28,57 +29,39 @@ import java.net.URLEncoder;
 * <p>These get translated to a string containing the URL safe encoding of the objectify key
 * followed by a (url-unsafe) ampersand delimiter and the SQL key.
 */
-public class VKeyTranslatorFactory<T> extends AbstractSimpleTranslatorFactory<VKey, String> {
-  private final Class<T> refClass;
+public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey, Key> {

-  public VKeyTranslatorFactory(Class<T> refClass) {
+  // Class registry allowing us to restore the original class object from the unqualified class
+  // name, which is all the datastore key gives us.
+  // Note that entities annotated with @EntitySubclass are removed because they share the same
+  // kind of the key with their parent class.
+  private static final ImmutableMap<String, Class> CLASS_REGISTRY =
+      ALL_CLASSES.stream()
+          .filter(clazz -> !clazz.isAnnotationPresent(EntitySubclass.class))
+          .collect(toImmutableMap(com.googlecode.objectify.Key::getKind, identity()));
+  ;
+
+  public VKeyTranslatorFactory() {
    super(VKey.class);
-    this.refClass = refClass;
  }

  @Override
-  public SimpleTranslator<VKey, String> createTranslator() {
-    return new SimpleTranslator<VKey, String>() {
+  public SimpleTranslator<VKey, Key> createTranslator() {
+    return new SimpleTranslator<VKey, Key>() {
      @Override
-      public VKey loadValue(String datastoreValue) {
-        int pos = datastoreValue.indexOf('&');
-        Key ofyKey = null;
-        String sqlKey = null;
-        if (pos > 0) {
-          // We have an objectify key.
-          ofyKey = Key.create(datastoreValue.substring(0, pos));
-        }
-
-        if (pos < datastoreValue.length() - 1) {
-          // We have an SQL key.
-          sqlKey = decode(datastoreValue.substring(pos + 1));
-        }
-
-        return VKey.create(refClass, sqlKey, ofyKey);
+      public VKey loadValue(Key datastoreValue) {
+        // TODO(mmuller): we need to call a method on refClass to also reconstitute the SQL key.
+        return datastoreValue == null
+            ? null
+            : VKey.createOfy(
+                CLASS_REGISTRY.get(datastoreValue.getKind()),
+                com.googlecode.objectify.Key.create(datastoreValue));
      }

      @Override
-      public String saveValue(VKey key) {
-        return ((key.getOfyKey() == null) ? "" : key.getOfyKey().getString())
-            + "&"
-            + ((key.getSqlKey() == null) ? "" : encode(key.getSqlKey().toString()));
+      public Key saveValue(VKey key) {
+        return key == null ? null : key.getOfyKey().getRaw();
      }
    };
  }
-
-  private static String encode(String val) {
-    try {
-      return URLEncoder.encode(val, UTF_8.toString());
-    } catch (UnsupportedEncodingException e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  private static String decode(String encoded) {
-    try {
-      return URLDecoder.decode(encoded, UTF_8.toString());
-    } catch (UnsupportedEncodingException e) {
-      throw new RuntimeException(e);
-    }
-  }
 }
@@ -0,0 +1,48 @@
+// 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 avro.shaded.com.google.common.collect.Maps;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.registry.Registry.BillingCostTransition;
+import java.util.Map;
+import javax.persistence.Converter;
+import org.joda.money.Money;
+import org.joda.time.DateTime;
+
+/**
+ * JPA converter for storing/retrieving {@link TimedTransitionProperty <Money, BillingCostTransition
+ * >} objects.
+ */
+@Converter(autoApply = true)
+public class BillingCostTransitionConverter
+    extends TimedTransitionPropertyConverterBase<Money, BillingCostTransition> {
+
+  @Override
+  Map.Entry<String, String> convertToDatabaseMapEntry(
+      Map.Entry<DateTime, BillingCostTransition> entry) {
+    return Maps.immutableEntry(entry.getKey().toString(), entry.getValue().getValue().toString());
+  }
+
+  @Override
+  Map.Entry<DateTime, Money> convertToEntityMapEntry(Map.Entry<String, String> entry) {
+    return Maps.immutableEntry(DateTime.parse(entry.getKey()), Money.parse(entry.getValue()));
+  }
+
+  @Override
+  Class<BillingCostTransition> getTimedTransitionSubclass() {
+    return BillingCostTransition.class;
+  }
+}
@@ -0,0 +1,63 @@
+// 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.collect.ImmutableMap.toImmutableMap;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.common.TimedTransitionProperty.TimedTransition;
+import google.registry.persistence.converter.StringMapDescriptor.StringMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import javax.persistence.AttributeConverter;
+import org.joda.time.DateTime;
+
+/**
+ * Base JPA converter for {@link TimedTransitionProperty} objects that are stored in a column with
+ * data type of hstore in the database.
+ */
+public abstract class TimedTransitionPropertyConverterBase<K, V extends TimedTransition<K>>
+    implements AttributeConverter<TimedTransitionProperty<K, V>, StringMap> {
+
+  abstract Map.Entry<String, String> convertToDatabaseMapEntry(Map.Entry<DateTime, V> entry);
+
+  abstract Map.Entry<DateTime, K> convertToEntityMapEntry(Map.Entry<String, String> entry);
+
+  abstract Class<V> getTimedTransitionSubclass();
+
+  @Override
+  public StringMap convertToDatabaseColumn(@Nullable TimedTransitionProperty<K, V> attribute) {
+    return attribute == null
+        ? null
+        : StringMap.create(
+            attribute.entrySet().stream()
+                .map(this::convertToDatabaseMapEntry)
+                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)));
+  }
+
+  @Override
+  public TimedTransitionProperty<K, V> convertToEntityAttribute(@Nullable StringMap dbData) {
+    if (dbData == null) {
+      return null;
+    }
+    Map<DateTime, K> map =
+        dbData.getMap().entrySet().stream()
+            .map(this::convertToEntityMapEntry)
+            .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+    return TimedTransitionProperty.fromValueMap(
+        ImmutableSortedMap.copyOf(map), getTimedTransitionSubclass());
+  }
+}
@@ -0,0 +1,48 @@
+// 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 com.google.common.collect.Maps;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.registry.Registry.TldState;
+import google.registry.model.registry.Registry.TldStateTransition;
+import java.util.Map;
+import javax.persistence.Converter;
+import org.joda.time.DateTime;
+
+/**
+ * JPA converter for storing/retrieving {@link TimedTransitionProperty<TldState,
+ * TldStateTransition>} objects.
+ */
+@Converter(autoApply = true)
+public class TldStateTransitionConverter
+    extends TimedTransitionPropertyConverterBase<TldState, TldStateTransition> {
+
+  @Override
+  Map.Entry<String, String> convertToDatabaseMapEntry(
+      Map.Entry<DateTime, TldStateTransition> entry) {
+    return Maps.immutableEntry(entry.getKey().toString(), entry.getValue().getValue().name());
+  }
+
+  @Override
+  Map.Entry<DateTime, TldState> convertToEntityMapEntry(Map.Entry<String, String> entry) {
+    return Maps.immutableEntry(DateTime.parse(entry.getKey()), TldState.valueOf(entry.getValue()));
+  }
+
+  @Override
+  Class<TldStateTransition> getTimedTransitionSubclass() {
+    return TldStateTransition.class;
+  }
+}
@@ -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;
+  }
 }
@@ -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) {
@@ -40,12 +40,10 @@ class CompareDbBackups {
    }

    ImmutableSet<ComparableEntity> entities1 =
-        new RecordAccumulator()
-            .readDirectory(new File(args[0]), DATA_FILE_MATCHER)
+        RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER)
            .getComparableEntitySet();
    ImmutableSet<ComparableEntity> entities2 =
-        new RecordAccumulator()
-            .readDirectory(new File(args[1]), DATA_FILE_MATCHER)
+        RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER)
            .getComparableEntitySet();

    // Calculate the entities added and removed.
@@ -14,17 +14,27 @@

 package google.registry.tools;

+import static com.google.common.base.Preconditions.checkState;
+
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.Optional;

 /**
- * Reads records from a set of LevelDB files and builds a gigantic ImmutableList from them.
+ * Iterator that incrementally parses binary data in LevelDb format into records.
+ *
+ * <p>The input source is automatically closed when all data have been read.
 *
 * <p>See <a
 * href="https://github.com/google/leveldb/blob/master/doc/log_format.md">log_format.md</a> for the
@@ -32,19 +42,74 @@ import java.nio.file.Path;
 *
 * <p>There are several other implementations of this, none of which appeared suitable for our use
 * case: <a href="https://github.com/google/leveldb">The original C++ implementation</a>. <a
- * href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/files/RecordWriteChannel">
- * com.google.appengine.api.files.RecordWriteChannel</a> - Exactly what we need but deprecated. The
+ * href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/files/RecordReadChannel">
+ * com.google.appengine.api.files.RecordReadChannel</a> - Exactly what we need but deprecated. The
 * referenced replacement: <a
 * href="https://github.com/GoogleCloudPlatform/appengine-gcs-client.git">The App Engine GCS
 * Client</a> - Does not appear to have any support for working with LevelDB.
 */
-public final class LevelDbLogReader {
+public final class LevelDbLogReader implements Iterator<byte[]> {

  @VisibleForTesting static final int BLOCK_SIZE = 32 * 1024;
  @VisibleForTesting static final int HEADER_SIZE = 7;

  private final ByteArrayOutputStream recordContents = new ByteArrayOutputStream();
-  private final ImmutableList.Builder<byte[]> recordListBuilder = new ImmutableList.Builder<>();
+  private final LinkedList<byte[]> recordList = Lists.newLinkedList();
+
+  private final ByteBuffer byteBuffer = ByteBuffer.allocate(BLOCK_SIZE);
+  private final ReadableByteChannel channel;
+
+  LevelDbLogReader(ReadableByteChannel channel) {
+    this.channel = channel;
+  }
+
+  @Override
+  public boolean hasNext() {
+    while (recordList.isEmpty()) {
+      try {
+        Optional<byte[]> block = readFromChannel();
+        if (!block.isPresent()) {
+          return false;
+        }
+        if (block.get().length != BLOCK_SIZE) {
+          throw new IllegalStateException("Data size is not multiple of " + BLOCK_SIZE);
+        }
+        processBlock(block.get());
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    return true;
+  }
+
+  @Override
+  public byte[] next() {
+    checkState(hasNext(), "The next() method called on empty iterator.");
+    return recordList.removeFirst();
+  }
+
+  /**
+   * Returns the next {@link #BLOCK_SIZE} bytes from the input channel, or {@link
+   * Optional#empty()} if there is no more data.
+   */
+  // TODO(weiminyu): use ByteBuffer directly.
+  private Optional<byte[]> readFromChannel() throws IOException {
+    while (channel.isOpen()) {
+      int bytesRead = channel.read(byteBuffer);
+      if (!byteBuffer.hasRemaining() || bytesRead < 0) {
+        byteBuffer.flip();
+        if (!byteBuffer.hasRemaining()) {
+          channel.close();
+          return Optional.empty();
+        }
+        byte[] result = new byte[byteBuffer.remaining()];
+        byteBuffer.get(result);
+        byteBuffer.clear();
+        return Optional.of(result);
+      }
+    }
+    return Optional.empty();
+  }

  /** Read a complete block, which must be exactly 32 KB. */
  private void processBlock(byte[] block) {
@@ -63,7 +128,7 @@ public final class LevelDbLogReader {

      // If this is the last (or only) chunk in the record, store the full contents into the List.
      if (recordHeader.type == ChunkType.FULL || recordHeader.type == ChunkType.LAST) {
-        recordListBuilder.add(recordContents.toByteArray());
+        recordList.add(recordContents.toByteArray());
        recordContents.reset();
      }

@@ -96,40 +161,24 @@ public final class LevelDbLogReader {
    return new RecordHeader(checksum, size, ChunkType.fromCode(type));
  }

-  /** Reads all records in the Reader into the record set. */
-  public void readFrom(InputStream source) throws IOException {
-    byte[] block = new byte[BLOCK_SIZE];
-
-    // read until we have no more.
-    while (true) {
-      int amountRead = source.read(block, 0, BLOCK_SIZE);
-      if (amountRead <= 0) {
-        break;
-      }
-      assert amountRead == BLOCK_SIZE;
-
-      processBlock(block);
-    }
+  /** Returns a {@link LevelDbLogReader} over a {@link ReadableByteChannel}. */
+  public static LevelDbLogReader from(ReadableByteChannel channel) {
+    return new LevelDbLogReader(channel);
  }

-  /** Reads all records from the file specified by "path" into the record set. */
-  public void readFrom(Path path) throws IOException {
-    readFrom(Files.newInputStream(path));
+  /** Returns a {@link LevelDbLogReader} over an {@link InputStream}. */
+  public static LevelDbLogReader from(InputStream source) {
+    return new LevelDbLogReader(Channels.newChannel(source));
  }

-  /** Reads all records from the specified file into the record set. */
-  public void readFrom(String filename) throws IOException {
-    readFrom(FileSystems.getDefault().getPath(filename));
+  /** Returns a {@link LevelDbLogReader} over a file specified by {@link Path}. */
+  public static LevelDbLogReader from(Path path) throws IOException {
+    return from(Files.newInputStream(path));
  }

-  /**
-   * Gets the list of records constructed so far.
-   *
-   * <p>Note that this does not invalidate the internal state of the object: we return a copy and
-   * this can be called multiple times.
-   */
-  ImmutableList<byte[]> getRecords() {
-    return recordListBuilder.build();
+  /** Returns a {@link LevelDbLogReader} over a file specified by {@code filename}. */
+  public static LevelDbLogReader from(String filename) throws IOException {
+    return from(FileSystems.getDefault().getPath(filename));
  }

  /** Aggregates the fields in a record header. */
@@ -15,38 +15,43 @@
 package google.registry.tools;

 import com.google.appengine.api.datastore.EntityTranslator;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.function.Predicate;

 /** Accumulates Entity records from level db files under a directory hierarchy. */
 class RecordAccumulator {
-  private final LevelDbLogReader reader = new LevelDbLogReader();
+  private final ImmutableList<byte[]> records;
+
+  RecordAccumulator(ImmutableList<byte[]> records) {
+    this.records = records;
+  }

  /** Recursively reads all records in the directory. */
-  public final RecordAccumulator readDirectory(File dir, Predicate<File> fileMatcher) {
+  public static RecordAccumulator readDirectory(File dir, Predicate<File> fileMatcher) {
+    ImmutableList.Builder<byte[]> builder = new ImmutableList.Builder<>();
    for (File child : dir.listFiles()) {
      if (child.isDirectory()) {
-        readDirectory(child, fileMatcher);
+        builder.addAll(readDirectory(child, fileMatcher).records);
      } else if (fileMatcher.test(child)) {
        try {
-          reader.readFrom(new FileInputStream(child));
+          builder.addAll(LevelDbLogReader.from(child.getPath()));
        } catch (IOException e) {
          throw new RuntimeException("IOException reading from file: " + child, e);
        }
      }
    }

-    return this;
+    return new RecordAccumulator(builder.build());
  }

  /** Creates an entity set from the current set of raw records. */
  ImmutableSet<ComparableEntity> getComparableEntitySet() {
    ImmutableSet.Builder<ComparableEntity> builder = new ImmutableSet.Builder<>();
-    for (byte[] rawRecord : reader.getRecords()) {
+    for (byte[] rawRecord : records) {
      // Parse the entity proto and create an Entity object from it.
      EntityProto proto = new EntityProto();
      proto.parseFrom(rawRecord);
@@ -35,6 +35,7 @@
    <class>google.registry.model.domain.GracePeriod</class>

    <!-- Customized type converters -->
+    <class>google.registry.persistence.converter.BillingCostTransitionConverter</class>
    <class>google.registry.persistence.converter.BloomFilterConverter</class>
    <class>google.registry.persistence.converter.CidrAddressBlockListConverter</class>
    <class>google.registry.persistence.converter.CreateAutoTimestampConverter</class>
@@ -47,6 +48,7 @@
    <class>google.registry.persistence.converter.StatusValueSetConverter</class>
    <class>google.registry.persistence.converter.StringListConverter</class>
    <class>google.registry.persistence.converter.StringSetConverter</class>
+    <class>google.registry.persistence.converter.TldStateTransitionConverter</class>
    <class>google.registry.persistence.converter.UpdateAutoTimestampConverter</class>
    <class>google.registry.persistence.converter.ZonedDateTimeConverter</class>

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

  @Test
@@ -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()
@@ -119,7 +119,8 @@ public class DomainBaseSqlTest {
        .transact(
            () -> {
              // Persist the contacts.  Note that these need to be persisted before the domain
-              // otherwise we get a foreign key constraint error.
+              // otherwise we get a foreign key constraint error.  If we ever decide to defer the
+              // relevant foreign key checks to commit time, then the order would not matter.
              jpaTm().saveNew(contact);
              jpaTm().saveNew(contact2);

@@ -127,7 +128,8 @@ public class DomainBaseSqlTest {
              jpaTm().saveNew(domain);

              // Persist the host.  This does _not_ need to be persisted before the domain,
-              // presumably because its relationship is stored in a join table.
+              // because only the row in the join table (DomainHost) is subject to foreign key
+              // constraints, and Hibernate knows to insert it after domain and host.
              jpaTm().saveNew(host);
            });

@@ -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();
-  }
 }
@@ -0,0 +1,77 @@
+// 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 static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.joda.money.CurrencyUnit.USD;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.ImmutableObject;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.registry.Registry.BillingCostTransition;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import org.joda.money.Money;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link BillingCostTransitionConverter}. */
+public class BillingCostTransitionConverterTest {
+
+  @RegisterExtension
+  public final JpaUnitTestRule jpa =
+      new JpaTestRules.Builder()
+          .withInitScript("sql/flyway/V14__load_extension_for_hstore.sql")
+          .withEntityClass(TestEntity.class)
+          .buildUnitTestRule();
+
+  private static final ImmutableSortedMap<DateTime, Money> values =
+      ImmutableSortedMap.of(
+          START_OF_TIME,
+          Money.of(USD, 8),
+          DateTime.parse("2001-01-01T00:00:00.0Z"),
+          Money.of(USD, 0));
+
+  @Test
+  void roundTripConversion_returnsSameTimedTransitionProperty() {
+    TimedTransitionProperty<Money, BillingCostTransition> timedTransitionProperty =
+        TimedTransitionProperty.fromValueMap(values, BillingCostTransition.class);
+    TestEntity testEntity = new TestEntity(timedTransitionProperty);
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(testEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(persisted.timedTransitionProperty).containsExactlyEntriesIn(timedTransitionProperty);
+  }
+
+  @Entity(name = "TestEntity")
+  private static class TestEntity extends ImmutableObject {
+
+    @Id String name = "id";
+
+    TimedTransitionProperty<Money, BillingCostTransition> timedTransitionProperty;
+
+    private TestEntity() {}
+
+    private TestEntity(
+        TimedTransitionProperty<Money, BillingCostTransition> timedTransitionProperty) {
+      this.timedTransitionProperty = timedTransitionProperty;
+    }
+  }
+}
@@ -0,0 +1,181 @@
+// 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 static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Maps;
+import google.registry.model.ImmutableObject;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.common.TimedTransitionProperty.TimedTransition;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import java.util.Map;
+import javax.persistence.Converter;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.NoResultException;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link TimedTransitionPropertyConverterBase}. */
+public class TimedTransitionPropertyConverterBaseTest {
+
+  @RegisterExtension
+  public final JpaUnitTestRule jpa =
+      new JpaTestRules.Builder()
+          .withInitScript("sql/flyway/V14__load_extension_for_hstore.sql")
+          .withEntityClass(TestTimedTransitionPropertyConverter.class, TestEntity.class)
+          .buildUnitTestRule();
+
+  private static final DateTime DATE_1 = DateTime.parse("2001-01-01T00:00:00.000Z");
+  private static final DateTime DATE_2 = DateTime.parse("2002-01-01T00:00:00.000Z");
+
+  private static final ImmutableSortedMap<DateTime, String> VALUES =
+      ImmutableSortedMap.of(
+          START_OF_TIME, "val1",
+          DATE_1, "val2",
+          DATE_2, "val3");
+
+  private static final TimedTransitionProperty<String, TestTransition> TIMED_TRANSITION_PROPERTY =
+      TimedTransitionProperty.fromValueMap(VALUES, TestTransition.class);
+
+  @Test
+  void roundTripConversion_returnsSameTimedTransitionProperty() {
+    TestEntity testEntity = new TestEntity(TIMED_TRANSITION_PROPERTY);
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(testEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(persisted.property).containsExactlyEntriesIn(TIMED_TRANSITION_PROPERTY);
+  }
+
+  @Test
+  void testUpdateColumn_succeeds() {
+    TestEntity testEntity = new TestEntity(TIMED_TRANSITION_PROPERTY);
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(testEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(persisted.property).containsExactlyEntriesIn(TIMED_TRANSITION_PROPERTY);
+    ImmutableSortedMap<DateTime, String> newValues = ImmutableSortedMap.of(START_OF_TIME, "val4");
+    persisted.property = TimedTransitionProperty.fromValueMap(newValues, TestTransition.class);
+    jpaTm().transact(() -> jpaTm().getEntityManager().merge(persisted));
+    TestEntity updated =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(updated.property.toValueMap()).isEqualTo(newValues);
+  }
+
+  @Test
+  void testNullValue_writesAndReadsNullSuccessfully() {
+    TestEntity testEntity = new TestEntity(null);
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(testEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(persisted.property).isNull();
+  }
+
+  @Test
+  void testNativeQuery_succeeds() {
+    executeNativeQuery(
+        "INSERT INTO \"TestEntity\" (name, property) VALUES ('id',"
+            + " 'val1=>1970-01-01T00:00:00.000Z, val2=>2001-01-01T00:00:00.000Z')");
+
+    assertThat(
+            getSingleResultFromNativeQuery(
+                "SELECT property -> 'val1' FROM \"TestEntity\" WHERE name = 'id'"))
+        .isEqualTo(START_OF_TIME.toString());
+    assertThat(
+            getSingleResultFromNativeQuery(
+                "SELECT property -> 'val2' FROM \"TestEntity\" WHERE name = 'id'"))
+        .isEqualTo(DATE_1.toString());
+
+    executeNativeQuery(
+        "UPDATE \"TestEntity\" SET property = 'val3=>2002-01-01T00:00:00.000Z' WHERE name = 'id'");
+
+    assertThat(
+            getSingleResultFromNativeQuery(
+                "SELECT property -> 'val3' FROM \"TestEntity\" WHERE name = 'id'"))
+        .isEqualTo(DATE_2.toString());
+
+    executeNativeQuery("DELETE FROM \"TestEntity\" WHERE name = 'id'");
+
+    assertThrows(
+        NoResultException.class,
+        () ->
+            getSingleResultFromNativeQuery(
+                "SELECT property -> 'val3' FROM \"TestEntity\" WHERE name = 'id'"));
+  }
+
+  private static Object getSingleResultFromNativeQuery(String sql) {
+    return jpaTm()
+        .transact(() -> jpaTm().getEntityManager().createNativeQuery(sql).getSingleResult());
+  }
+
+  private static void executeNativeQuery(String sql) {
+    jpaTm().transact(() -> jpaTm().getEntityManager().createNativeQuery(sql).executeUpdate());
+  }
+
+  public static class TestTransition extends TimedTransition<String> {
+    private String transition;
+
+    @Override
+    public String getValue() {
+      return transition;
+    }
+
+    @Override
+    protected void setValue(String transition) {
+      this.transition = transition;
+    }
+  }
+
+  @Converter(autoApply = true)
+  private static class TestTimedTransitionPropertyConverter
+      extends TimedTransitionPropertyConverterBase<String, TestTransition> {
+
+    @Override
+    Map.Entry<DateTime, String> convertToEntityMapEntry(Map.Entry<String, String> entry) {
+      return Maps.immutableEntry(DateTime.parse(entry.getKey()), entry.getValue());
+    }
+
+    @Override
+    Class<TestTransition> getTimedTransitionSubclass() {
+      return TestTransition.class;
+    }
+
+    @Override
+    Map.Entry<String, String> convertToDatabaseMapEntry(Map.Entry<DateTime, TestTransition> entry) {
+      return Maps.immutableEntry(entry.getKey().toString(), entry.getValue().getValue());
+    }
+  }
+
+  @Entity(name = "TestEntity") // Override entity name to avoid the nested class reference.
+  private static class TestEntity extends ImmutableObject {
+
+    @Id String name = "id";
+
+    TimedTransitionProperty<String, TestTransition> property;
+
+    private TestEntity() {}
+
+    private TestEntity(TimedTransitionProperty<String, TestTransition> timedTransitionProperty) {
+      this.property = timedTransitionProperty;
+    }
+  }
+}
@@ -0,0 +1,80 @@
+// 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 static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.ImmutableObject;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.model.registry.Registry.TldState;
+import google.registry.model.registry.Registry.TldStateTransition;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link TldStateTransitionConverter}. */
+public class TldStateTransitionConverterTest {
+
+  @RegisterExtension
+  public final JpaUnitTestRule jpa =
+      new JpaTestRules.Builder()
+          .withInitScript("sql/flyway/V14__load_extension_for_hstore.sql")
+          .withEntityClass(TestEntity.class)
+          .buildUnitTestRule();
+
+  private static final ImmutableSortedMap<DateTime, TldState> values =
+      ImmutableSortedMap.of(
+          START_OF_TIME,
+          TldState.PREDELEGATION,
+          DateTime.parse("2001-01-01T00:00:00.0Z"),
+          TldState.QUIET_PERIOD,
+          DateTime.parse("2002-01-01T00:00:00.0Z"),
+          TldState.PDT,
+          DateTime.parse("2003-01-01T00:00:00.0Z"),
+          TldState.GENERAL_AVAILABILITY);
+
+  @Test
+  void roundTripConversion_returnsSameTimedTransitionProperty() {
+    TimedTransitionProperty<TldState, TldStateTransition> timedTransitionProperty =
+        TimedTransitionProperty.fromValueMap(values, TldStateTransition.class);
+    TestEntity testEntity = new TestEntity(timedTransitionProperty);
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(testEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    assertThat(persisted.timedTransitionProperty).containsExactlyEntriesIn(timedTransitionProperty);
+  }
+
+  @Entity(name = "TestEntity")
+  private static class TestEntity extends ImmutableObject {
+
+    @Id String name = "id";
+
+    TimedTransitionProperty<TldState, TldStateTransition> timedTransitionProperty;
+
+    private TestEntity() {}
+
+    private TestEntity(
+        TimedTransitionProperty<TldState, TldStateTransition> timedTransitionProperty) {
+      this.timedTransitionProperty = timedTransitionProperty;
+    }
+  }
+}
@@ -28,6 +28,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import org.hibernate.cfg.Environment;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
@@ -60,10 +61,11 @@ public class JpaTestRules {

  /**
   * Junit rule for unit tests with JPA framework, when the underlying database is populated by the
-   * optional init script (which must not be the Nomulus Cloud SQL schema).
+   * optional init script (which must not be the Nomulus Cloud SQL schema). This rule can also be
+   * used as am extension for JUnit5 tests.
   */
-  public static class JpaUnitTestRule extends JpaTransactionManagerRule {
-
+  public static class JpaUnitTestRule extends JpaTransactionManagerRule
+      implements BeforeEachCallback, AfterEachCallback {
    private JpaUnitTestRule(
        Clock clock,
        Optional<String> initScriptPath,
@@ -71,6 +73,16 @@ public class JpaTestRules {
        ImmutableMap<String, String> userProperties) {
      super(clock, initScriptPath, extraEntityClasses, userProperties);
    }
+
+    @Override
+    public void beforeEach(ExtensionContext context) throws Exception {
+      this.before();
+    }
+
+    @Override
+    public void afterEach(ExtensionContext context) throws Exception {
+      this.after();
+    }
  }

  /**
@@ -169,6 +181,17 @@ public class JpaTestRules {
      return this;
    }

+    /**
+     * Enables logging of SQL statements.
+     *
+     * <p>SQL logging is very noisy and disabled by default. This method maybe useful when
+     * troubleshooting a specific test.
+     */
+    public Builder withSqlLogging() {
+      withProperty(Environment.SHOW_SQL, "true");
+      return this;
+    }
+
    /** Builds a {@link JpaIntegrationTestRule} instance. */
    public JpaIntegrationTestRule buildIntegrationTestRule() {
      return new JpaIntegrationTestRule(
@@ -195,7 +218,9 @@ public class JpaTestRules {
      return new JpaIntegrationWithCoverageExtension(buildIntegrationTestRule());
    }

-    /** Builds a {@link JpaUnitTestRule} instance. */
+    /**
+     * Builds a {@link JpaUnitTestRule} instance that can also be used as an extension for JUnit5.
+     */
    public JpaUnitTestRule buildUnitTestRule() {
      checkState(
          !Objects.equals(GOLDEN_SCHEMA_SQL_PATH, initScript),
@@ -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();
@@ -41,6 +41,8 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.stream.Collectors;
@@ -146,13 +148,15 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
    ImmutableMap properties = PersistenceModule.providesDefaultDatabaseConfigs();
    if (!userProperties.isEmpty()) {
      // If there are user properties, create a new properties object with these added.
-      ImmutableMap.Builder builder = properties.builder();
-      builder.putAll(userProperties);
-      // Forbid Hibernate push to stay consistent with flyway-based schema management.
-      builder.put(Environment.HBM2DDL_AUTO, "none");
-      builder.put(Environment.SHOW_SQL, "true");
-      properties = builder.build();
+      Map<String, String> mergedProperties = Maps.newHashMap();
+      mergedProperties.putAll(properties);
+      mergedProperties.putAll(userProperties);
+      properties = ImmutableMap.copyOf(mergedProperties);
    }
+    // Forbid Hibernate push to stay consistent with flyway-based schema management.
+    checkState(
+        Objects.equals(properties.get(Environment.HBM2DDL_AUTO), "none"),
+        "The HBM2DDL_AUTO property must be 'none'.");
    assertReasonableNumDbConnections();
    emf =
        createEntityManagerFactory(
@@ -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));
    }
  }
 }
@@ -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;
+  }
 }
@@ -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 {}
@@ -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);
+  }
+}
@@ -23,7 +23,6 @@ import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import google.registry.testing.AppEngineRule;
 import google.registry.tools.LevelDbFileBuilder.Property;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import org.junit.Rule;
 import org.junit.Test;
@@ -51,10 +50,7 @@ public class LevelDbFileBuilderTest {
            BASE_ID, Property.create("first", 100L), Property.create("second", 200L));
    builder.build();

-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new FileInputStream(logFile));
-
-    ImmutableList<byte[]> records = reader.getRecords();
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
    assertThat(records).hasSize(1);

    // Reconstitute an entity, make sure that what we've got is the same as what we started with.
@@ -82,10 +78,7 @@ public class LevelDbFileBuilderTest {
    builder.build();
    ImmutableList<ComparableEntity> originalEntities = originalEntitiesBuilder.build();

-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new FileInputStream(logFile));
-
-    ImmutableList<byte[]> records = reader.getRecords();
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
    assertThat(records).hasSize(1000);
    int index = 0;
    for (byte[] record : records) {
@@ -17,6 +17,10 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.tools.LevelDbUtil.MAX_RECORD;
 import static google.registry.tools.LevelDbUtil.addRecord;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;

 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Bytes;
@@ -24,12 +28,9 @@ import google.registry.tools.LevelDbLogReader.ChunkType;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.List;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.Test;

-/** LevelDbLogReader tests. */
-@RunWith(JUnit4.class)
+/** Unit tests of {@link LevelDbLogReader}. */
 public final class LevelDbLogReaderTest {

  // Size of the test record.  Any value < 256 will do.
@@ -54,28 +55,23 @@ public final class LevelDbLogReaderTest {
  @Test
  public void testSimpleBlock() throws IOException {
    TestBlock block = makeBlockOfRepeatingBytes(0);
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new ByteArrayInputStream(block.data));
-    ImmutableList<byte[]> records = reader.getRecords();
-    assertThat(records).hasSize(block.recordCount);
+    assertThat(readIncrementally(block.data)).hasSize(block.recordCount);
  }

  @Test
  public void testLargeRecord() throws IOException {
-    byte[] block = new byte[LevelDbLogReader.BLOCK_SIZE];
-    addRecord(block, 0, ChunkType.FIRST, MAX_RECORD, (byte) 1);
-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(new ByteArrayInputStream(block));
-    assertThat(reader.getRecords()).isEmpty();
+    byte[] block0 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block0, 0, ChunkType.FIRST, MAX_RECORD, (byte) 1);
+    assertThat(readIncrementally(block0)).isEmpty();

-    addRecord(block, 0, ChunkType.MIDDLE, MAX_RECORD, (byte) 2);
-    reader.readFrom(new ByteArrayInputStream(block));
-    assertThat(reader.getRecords()).isEmpty();
+    byte[] block1 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block1, 0, ChunkType.MIDDLE, MAX_RECORD, (byte) 2);
+    assertThat(readIncrementally(block0, block1)).isEmpty();

-    addRecord(block, 0, ChunkType.LAST, MAX_RECORD, (byte) 3);
-    reader.readFrom(new ByteArrayInputStream(block));
+    byte[] block2 = new byte[LevelDbLogReader.BLOCK_SIZE];
+    addRecord(block2, 0, ChunkType.LAST, MAX_RECORD, (byte) 3);

-    List<byte[]> records = reader.getRecords();
+    List<byte[]> records = readIncrementally(block0, block1, block2);
    assertThat(records).hasSize(1);
    byte[] record = records.get(0);

@@ -95,11 +91,24 @@ public final class LevelDbLogReaderTest {
  public void readFromMultiBlockStream() throws IOException {
    TestBlock block0 = makeBlockOfRepeatingBytes(0);
    TestBlock block1 = makeBlockOfRepeatingBytes(138);
-    ByteArrayInputStream source = new ByteArrayInputStream(Bytes.concat(block0.data, block1.data));
+    assertThat(readIncrementally(block0.data, block1.data))
+        .hasSize(block0.recordCount + block1.recordCount);
+  }

-    LevelDbLogReader reader = new LevelDbLogReader();
-    reader.readFrom(source);
-    assertThat(reader.getRecords()).hasSize(block0.recordCount + block1.recordCount);
+  @Test
+  void read_noData() throws IOException {
+    assertThat(readIncrementally(new byte[0])).isEmpty();
+  }
+
+  @Test
+  void read_failBadFirstBlock() {
+    assertThrows(IllegalStateException.class, () -> readIncrementally(new byte[1]));
+  }
+
+  @Test
+  void read_failBadTrailingBlock() {
+    TestBlock block = makeBlockOfRepeatingBytes(0);
+    assertThrows(IllegalStateException.class, () -> readIncrementally(block.data, new byte[2]));
  }

  @Test
@@ -112,6 +121,14 @@ public final class LevelDbLogReaderTest {
    assertThat(ChunkType.fromCode(ChunkType.LAST.getCode())).isEqualTo(ChunkType.LAST);
  }

+  @SafeVarargs
+  private static ImmutableList<byte[]> readIncrementally(byte[]... blocks) throws IOException {
+    ByteArrayInputStream source = spy(new ByteArrayInputStream(Bytes.concat(blocks)));
+    ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(source));
+    verify(source, times(1)).close();
+    return records;
+  }
+
  /** Aggregates the bytes of a test block with the record count. */
  private static final class TestBlock {
    final byte[] data;
@@ -76,7 +76,7 @@ public class RecordAccumulatorTest {
    builder.build();

    ImmutableSet<ComparableEntity> entities =
-        new RecordAccumulator().readDirectory(subdir, any -> true).getComparableEntitySet();
+        RecordAccumulator.readDirectory(subdir, any -> true).getComparableEntitySet();
    assertThat(entities).containsExactly(e1, e2, e3);
  }
 }
@@ -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;
@@ -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";
+
@@ -122,7 +122,7 @@

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

    create table "GracePeriod" (
@@ -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);


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

@@ -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) {
@@ -473,9 +473,9 @@
      "dev": true
    },
    "eventemitter3": {
-      "version": "4.0.0",
-      "resolved": "https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.0.tgz",
-      "integrity": "sha512-qerSRB0p+UDEssxTtm6EDKcE7W4OaoisfIMl4CngyEhjpYglocpNg6UEqCvemdGhosAsg4sO2dXJOdyBifPGCg==",
+      "version": "4.0.4",
+      "resolved": "https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.4.tgz",
+      "integrity": "sha512-rlaVLnVxtxvoyLsQQFBx53YmXHDxRIzzTLbdfxqi4yocpSjAxXwkU0cScM5JgSKMqEhrZpnvQ2D9gjylR0AimQ==",
      "dev": true
    },
    "extend": {
@@ -673,9 +673,9 @@
      }
    },
    "http-proxy": {
-      "version": "1.18.0",
-      "resolved": "https://registry.npmjs.org/http-proxy/-/http-proxy-1.18.0.tgz",
-      "integrity": "sha512-84I2iJM/n1d4Hdgc6y2+qY5mDaz2PUVjlg9znE9byl+q0uC3DeByqBGReQu5tpLK0TAqTIXScRUV+dg7+bUPpQ==",
+      "version": "1.18.1",
+      "resolved": "https://registry.npmjs.org/http-proxy/-/http-proxy-1.18.1.tgz",
+      "integrity": "sha512-7mz/721AbnJwIVbnaSv1Cz3Am0ZLT/UBwkC92VlxhXv/k/BBQfM2fXElQNC27BVGr0uwUpplYPQM9LnaBMR5NQ==",
      "dev": true,
      "requires": {
        "eventemitter3": "^4.0.0",
Author	SHA1	Message	Date
Weimin Yu	54f1357d83	Fix show-sql which stopped working (#596 ) * Fix show-sql which stopped working Made show-sql property configurable in JpaUnitTestRules. Added a few comments on foreign key constraint behavior.	2020-05-21 12:20:56 -04:00
Lai Jiang	c73d154084	Do not enqueue update snapshot task if import fails (#578 ) If the import from Datastore to BigQuery fails, there is no point enqueuing a job to update the snapshot view. Also when there's an error updating the snapshot view, log it at severe level. The HTTP exception thrown is logged at info and triggers a retry implicitly. I'm not sure if we want this behavior though. Do we want to retry upon snapshot updating failures? Unless the failurs are transient, retrying doesn't help. In our case the failure (End of time out of range in Standard SQL) is not transient.	2020-05-21 11:40:45 -04:00
gbrodman	259d2e2cdc	Run "npm audit fix" to fix a vulnerability (#592 )	2020-05-20 15:12:27 -04:00
Shicong Huang	0f174d9ce0	Add all existing entities to VKeyTranslatorFactory (#595 ) EntityClasses.ALL_CLASSES has all of our registered entities so we can use it to initialize VKeyTranslatorFactory.classRegistry to avoid adding them one by one. Also, this PR changed to use Key.getKind() to get the kind of the entity to solve the problem that when the entity class is an inner class, its kind should still be the class name instead of OuterClass$InnerClass.	2020-05-20 14:24:45 -04:00
Weimin Yu	ca2edb6a17	Close input channel in LevelDbLogReader (#594 ) * Close input channel in LevelDbLogReader Input channel should be closed when all data has been read.	2020-05-20 12:54:13 -04:00
Weimin Yu	3947ac6ef7	Read LevelDb incrementally (#593 ) * Read LevelDb incrementally Made LevelDbLogReader an iterator over a LevelDb data stream, Reducing memory footprint which is important when used in a Dataflow pipeline.	2020-05-20 10:26:34 -04:00
Michael Muller	579a3d0ac1	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	2020-05-19 14:10:28 -04:00
Lai Jiang	5fe929b027	Log InternalServerErrorException at SEVERE (#585 ) Normal HttpException logs at INFO because they usual do not indicate anything out of the ordinary and is meant to convey to the client that there is some expected error. However InternalServerErrorException is something that we do care about being alerted for so we log it at SEVERE.	2020-05-18 22:55:13 -04:00
Lai Jiang	fb335b7d89	Upgrade to Gradle 6.4.1 (#590 )	2020-05-18 16:47:02 -04:00
Shicong Huang	a0f4013d53	Add JUnit5 extension to run test twice against different databases (#588 ) * Add JUnit5 extension to run test against different databases * Fix typos * Add some explanation	2020-05-18 11:06:21 -04:00
Lai Jiang	5e596bb389	Upgrade to Gradle 6.4 (#589 )	2020-05-14 14:57:24 -04:00
Lai Jiang	f62fd82803	Log information about SSL connection from the client (#586 )	2020-05-14 09:38:33 -04:00
sarahcaseybot	b7353ef338	Add TimedTransitionProperty Converters (#561 ) * Add TimedTldStateTransitionMapConverter * Move timedTransitions to a base class and add BillingCostTransitionConverter * Add test of TimedTransitionPropertyConverterBase * clean up tests * Switch tests to JUnit 5 * Make JpaUnitTestRule an extension	2020-05-12 11:46:19 -04:00