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();
   }

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

@@ -93,3 +93,4 @@ V92__singletons.sql
 V93__defer_all_fkeys.sql
 V94__rename_lock_scope.sql
 V95__add_contacts_indexes_on_domain.sql
+V96__rename_sql_checkpoint_fields.sql

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

@@ -0,0 +1,15 @@
+-- Copyright 2021 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 IF EXISTS "SqlReplayCheckpoint" RENAME revision_id TO id;

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

@@ -695,9 +695,9 @@
     );
 
     create table "SqlReplayCheckpoint" (
-       revision_id int8 not null,
-        last_replay_time timestamptz,
-        primary key (revision_id)
+       id int8 not null,
+        last_replay_time timestamptz not null,
+        primary key (id)
     );
 
     create table "Tld" (

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

@@ -998,7 +998,7 @@ ALTER SEQUENCE public."SignedMarkRevocationList_revision_id_seq" OWNED BY public
 --
 
 CREATE TABLE public."SqlReplayCheckpoint" (
-    revision_id bigint NOT NULL,
+    id bigint NOT NULL,
     last_replay_time timestamp with time zone NOT NULL
 );
 
@@ -1415,7 +1415,7 @@ ALTER TABLE ONLY public."SignedMarkRevocationList"
 --
 
 ALTER TABLE ONLY public."SqlReplayCheckpoint"
-    ADD CONSTRAINT "SqlReplayCheckpoint_pkey" PRIMARY KEY (revision_id);
+    ADD CONSTRAINT "SqlReplayCheckpoint_pkey" PRIMARY KEY (id);
 
 
 --