Add loadOnlyOf method to tm() (#1162)

* Add loadOnlyOf method to tm()

In addition there's a bit of a refator of SqlReplayCheckpoint to make it
more in line with the other singletons. This method is useful for the
singleton classes where we expect at most one entity to exist, e.g.
ServerSecret.

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

@@ -262,13 +262,17 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
-    Query<T> query = getOfy().load().type(clazz);
-    // If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
-    // query to give us strong transactional consistency.
-    if (clazz.isAnnotationPresent(InCrossTld.class)) {
-      query = query.ancestor(getCrossTldKey());
-    }
-    return ImmutableList.copyOf(query);
+    return ImmutableList.copyOf(getPossibleAncestorQuery(clazz));
+  }
+
+  @Override
+  public <T> Optional<T> loadSingleton(Class<T> clazz) {
+    List<T> elements = getPossibleAncestorQuery(clazz).limit(2).list();
+    checkArgument(
+        elements.size() <= 1,
+        "Expected at most one entity of type %s, found at least two",
+        clazz.getSimpleName());
+    return elements.stream().findFirst();
   }
 
   @Override
@@ -401,6 +405,17 @@ public class DatastoreTransactionManager implements TransactionManager {
     return obj;
   }
 
+  /** A query for returning any/all results of an object, with an ancestor if possible. */
+  private <T> Query<T> getPossibleAncestorQuery(Class<T> clazz) {
+    Query<T> query = getOfy().load().type(clazz);
+    // If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
+    // query to give us strong transactional consistency.
+    if (clazz.isAnnotationPresent(InCrossTld.class)) {
+      query = query.ancestor(getCrossTldKey());
+    }
+    return query;
+  }
+
   private static class DatastoreQueryComposerImpl<T> extends QueryComposer<T> {
 
     DatastoreQueryComposerImpl(Class<T> entityClass) {

core/src/main/java/google/registry/model/server/ServerSecret.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.server;
 
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -22,7 +21,6 @@ import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.primitives.Longs;
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.OnLoad;
@@ -30,7 +28,6 @@ import com.googlecode.objectify.annotation.Unindex;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.CrossTldSingleton;
-import google.registry.persistence.VKey;
 import google.registry.schema.replay.NonReplicatedEntity;
 import java.nio.ByteBuffer;
 import java.util.Optional;
@@ -65,14 +62,9 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
               });
 
   private static ServerSecret retrieveAndSaveSecret() {
-    VKey<ServerSecret> vkey =
-        VKey.create(
-            ServerSecret.class,
-            SINGLETON_ID,
-            Key.create(getCrossTldKey(), ServerSecret.class, SINGLETON_ID));
     if (tm().isOfy()) {
       // Attempt a quick load if we're in ofy first to short-circuit sans transaction
-      Optional<ServerSecret> secretWithoutTransaction = tm().loadByKeyIfPresent(vkey);
+      Optional<ServerSecret> secretWithoutTransaction = tm().loadSingleton(ServerSecret.class);
       if (secretWithoutTransaction.isPresent()) {
         return secretWithoutTransaction.get();
       }
@@ -81,7 +73,7 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
             () -> {
               // Make sure we're in a transaction and attempt to load any existing secret, then
               // create it if it's absent.
-              Optional<ServerSecret> secret = tm().loadByKeyIfPresent(vkey);
+              Optional<ServerSecret> secret = tm().loadSingleton(ServerSecret.class);
               if (!secret.isPresent()) {
                 secret = Optional.of(create(UUID.randomUUID()));
                 tm().insertWithoutBackup(secret.get());

core/src/main/java/google/registry/model/tmch/TmchCrl.java

@@ -15,17 +15,14 @@
 package google.registry.model.tmch;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.CrossTldSingleton;
-import google.registry.persistence.VKey;
 import google.registry.schema.replay.NonReplicatedEntity;
 import java.util.Optional;
 import javax.annotation.concurrent.Immutable;
@@ -50,13 +47,7 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
 
   /** Returns the singleton instance of this entity, without memoization. */
   public static Optional<TmchCrl> get() {
-    return tm().transact(
-            () ->
-                tm().loadByKeyIfPresent(
-                        VKey.create(
-                            TmchCrl.class,
-                            SINGLETON_ID,
-                            Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID))));
+    return tm().transact(() -> tm().loadSingleton(TmchCrl.class));
   }
 
   /**

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

@@ -475,6 +475,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             .getResultList());
   }
 
+  @Override
+  public <T> Optional<T> loadSingleton(Class<T> clazz) {
+    assertInTransaction();
+    List<T> elements =
+        getEntityManager()
+            .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
+            .setMaxResults(2)
+            .getResultList();
+    checkArgument(
+        elements.size() <= 1,
+        "Expected at most one entity of type %s, found at least two",
+        clazz.getSimpleName());
+    return elements.stream().findFirst();
+  }
+
   private int internalDelete(VKey<?> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();

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

@@ -246,6 +246,13 @@ public interface TransactionManager {
    */
   <T> ImmutableList<T> loadAllOf(Class<T> clazz);
 
+  /**
+   * Loads the only instance of this particular class, or empty if none exists.
+   *
+   * <p>Throws an exception if there is more than one element in the table.
+   */
+  <T> Optional<T> loadSingleton(Class<T> clazz);
+
   /** Deletes the entity by its id. */
   void delete(VKey<?> key);
 

core/src/main/java/google/registry/schema/replay/SqlReplayCheckpoint.java

@@ -14,26 +14,24 @@
 
 package google.registry.schema.replay;
 
-import static google.registry.model.common.CrossTldSingleton.SINGLETON_ID;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
+import google.registry.model.common.CrossTldSingleton;
+import javax.persistence.Column;
 import javax.persistence.Entity;
-import javax.persistence.Id;
 import org.joda.time.DateTime;
 
 @Entity
-public class SqlReplayCheckpoint implements SqlOnlyEntity {
-
-  // Hibernate doesn't allow us to have a converted DateTime as our primary key so we need this
-  @Id private long revisionId = SINGLETON_ID;
+public class SqlReplayCheckpoint extends CrossTldSingleton implements SqlOnlyEntity {
 
+  @Column(nullable = false)
   private DateTime lastReplayTime;
 
   public static DateTime get() {
     jpaTm().assertInTransaction();
-    return jpaTm().loadAllOf(SqlReplayCheckpoint.class).stream()
-        .findFirst()
+    return jpaTm()
+        .loadSingleton(SqlReplayCheckpoint.class)
         .map(checkpoint -> checkpoint.lastReplayTime)
         .orElse(START_OF_TIME);
   }

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

@@ -17,7 +17,9 @@ package google.registry.persistence.transaction;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.common.collect.ImmutableList;
@@ -372,6 +374,27 @@ public class TransactionManagerTest {
         () -> tm().transact(() -> tm().loadAllOf(TestEntity.class)));
   }
 
+  @TestOfyAndSql
+  void loadSingleton_returnsValue_orEmpty() {
+    assertEntityNotExist(theEntity);
+    assertThat(transactIfJpaTm(() -> tm().loadSingleton(TestEntity.class))).isEmpty();
+
+    tm().transact(() -> tm().insert(theEntity));
+    assertThat(transactIfJpaTm(() -> tm().loadSingleton(TestEntity.class))).hasValue(theEntity);
+  }
+
+  @TestOfyAndSql
+  void loadSingleton_exceptionOnMultiple() {
+    assertAllEntitiesNotExist(moreEntities);
+    tm().transact(() -> tm().insertAll(moreEntities));
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> transactIfJpaTm(() -> tm().loadSingleton(TestEntity.class))))
+        .hasMessageThat()
+        .isEqualTo("Expected at most one entity of type TestEntity, found at least two");
+  }
+
   private static void assertEntityExists(TestEntity entity) {
     assertThat(tm().transact(() -> tm().exists(entity))).isTrue();
   }