Convert TmchCrl and ServerSecret to cleaner tm() impls (#1068)

* Convert TmchCrl and ServerSecret to cleaner tm() impls

When I implemented this originally I knew a lot less than I know now
about how we'll be storing and retrieving these singletons from SQL. The
optimal way here is to use the single SINGLETON_ID as the primary key,
that way we always know how to create the key that we can use in the
tm() retrieval.

This allows us to use generic tm() methods and to remove the handcrafted
SQL queries.

core/src/main/java/google/registry/model/common/CrossTldSingleton.java

@@ -31,7 +31,7 @@ public abstract class CrossTldSingleton extends ImmutableObject {
 
   public static final long SINGLETON_ID = 1; // There is always exactly one of these.
 
-  @Id @Transient long id = SINGLETON_ID;
+  @Id @javax.persistence.Id long id = SINGLETON_ID;
 
   @Transient @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
 }

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

@@ -15,8 +15,6 @@
 package google.registry.model.server;
 
 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.google.common.annotations.VisibleForTesting;
@@ -35,10 +33,10 @@ 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;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import javax.persistence.Column;
-import javax.persistence.Id;
 import javax.persistence.PostLoad;
 import javax.persistence.Transient;
 
@@ -67,22 +65,28 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
               });
 
   private static ServerSecret retrieveAndSaveSecret() {
-    VKey<ServerSecret> key =
+    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);
+      if (secretWithoutTransaction.isPresent()) {
+        return secretWithoutTransaction.get();
+      }
+    }
     return tm().transact(
             () -> {
-              // transactionally create a new ServerSecret (once per app setup) if necessary.
-              // return the ofy() result during Datastore-primary phase
-              ServerSecret secret =
-                  ofyTm().loadByKeyIfPresent(key).orElseGet(() -> create(UUID.randomUUID()));
-              // During a dual-write period, write it to both Datastore and SQL
-              // even if we didn't have to retrieve it from the DB
-              ofyTm().transact(() -> ofyTm().putWithoutBackup(secret));
-              jpaTm().transact(() -> jpaTm().putWithoutBackup(secret));
-              return secret;
+              // 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);
+              if (!secret.isPresent()) {
+                secret = Optional.of(create(UUID.randomUUID()));
+                tm().insertWithoutBackup(secret.get());
+              }
+              return secret.get();
             });
   }
 
@@ -102,7 +106,6 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
   @Transient long leastSignificant;
 
   /** The UUID value itself. */
-  @Id
   @Column(columnDefinition = "uuid")
   @Ignore
   UUID secret;

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

@@ -25,15 +25,11 @@ 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.model.tmch.TmchCrl.TmchCrlId;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.NonReplicatedEntity;
-import java.io.Serializable;
 import java.util.Optional;
 import javax.annotation.concurrent.Immutable;
 import javax.persistence.Column;
-import javax.persistence.Id;
-import javax.persistence.IdClass;
 import org.joda.time.DateTime;
 
 /** Datastore singleton for ICANN's TMCH CA certificate revocation list (CRL). */
@@ -41,22 +37,26 @@ import org.joda.time.DateTime;
 @javax.persistence.Entity
 @Immutable
 @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
-@IdClass(TmchCrlId.class)
 public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEntity {
 
-  @Id String crl;
+  @Column(name = "certificateRevocations", nullable = false)
+  String crl;
 
-  @Id DateTime updated;
+  @Column(name = "updateTimestamp", nullable = false)
+  DateTime updated;
 
-  @Id String url;
+  @Column(nullable = false)
+  String url;
 
   /** Returns the singleton instance of this entity, without memoization. */
   public static Optional<TmchCrl> get() {
-    VKey<TmchCrl> key =
-        VKey.create(
-            TmchCrl.class, SINGLETON_ID, Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID));
-    // return the ofy() result during Datastore-primary phase
-    return ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(key));
+    return tm().transact(
+            () ->
+                tm().loadByKeyIfPresent(
+                        VKey.create(
+                            TmchCrl.class,
+                            SINGLETON_ID,
+                            Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID))));
   }
 
   /**
@@ -75,13 +75,7 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
               tmchCrl.crl = checkNotNull(crl, "crl");
               tmchCrl.url = checkNotNull(url, "url");
               ofyTm().transactNew(() -> ofyTm().putWithoutBackup(tmchCrl));
-              jpaTm()
-                  .transactNew(
-                      () -> {
-                        // Delete the old one and insert the new one
-                        jpaTm().query("DELETE FROM TmchCrl").executeUpdate();
-                        jpaTm().putWithoutBackup(tmchCrl);
-                      });
+              jpaTm().transactNew(() -> jpaTm().putWithoutBackup(tmchCrl));
             });
   }
 
@@ -99,26 +93,4 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
   public final DateTime getUpdated() {
     return updated;
   }
-
-  static class TmchCrlId implements Serializable {
-
-    @Column(name = "certificateRevocations")
-    String crl;
-
-    @Column(name = "updateTimestamp")
-    DateTime updated;
-
-    String url;
-
-    /** Hibernate requires this default constructor. */
-    private TmchCrlId() {}
-
-    static TmchCrlId create(String crl, DateTime updated, String url) {
-      TmchCrlId result = new TmchCrlId();
-      result.crl = crl;
-      result.updated = updated;
-      result.url = url;
-      return result;
-    }
-  }
 }

core/src/test/java/google/registry/model/server/ServerSecretTest.java

@@ -15,16 +15,19 @@
 package google.registry.model.server;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
+import static google.registry.testing.DatabaseHelper.persistResource;
 
 import google.registry.model.EntityTestCase;
 import google.registry.model.ofy.RequestCapturingAsyncDatastoreService;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import java.util.UUID;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link ServerSecret}. */
+@DualDatabaseTest
 public class ServerSecretTest extends EntityTestCase {
 
   ServerSecretTest() {
@@ -36,24 +39,22 @@ public class ServerSecretTest extends EntityTestCase {
     ServerSecret.resetCache();
   }
 
-  @Test
+  @TestOfyAndSql
   void testGet_bootstrapping_savesSecretToDatastore() {
     ServerSecret secret = ServerSecret.get();
     assertThat(secret).isNotNull();
-    assertThat(ofy().load().entity(new ServerSecret()).now()).isEqualTo(secret);
-    assertThat(loadFromSql()).isEqualTo(secret);
+    assertThat(loadByEntity(new ServerSecret())).isEqualTo(secret);
   }
 
-  @Test
+  @TestOfyAndSql
   void testGet_existingSecret_returned() {
     ServerSecret secret = ServerSecret.create(new UUID(123, 456));
-    ofy().saveWithoutBackup().entity(secret).now();
+    persistResource(secret);
     assertThat(ServerSecret.get()).isEqualTo(secret);
-    assertThat(ofy().load().entity(new ServerSecret()).now()).isEqualTo(secret);
-    assertThat(loadFromSql()).isEqualTo(secret);
+    assertThat(loadByEntity(new ServerSecret())).isEqualTo(secret);
   }
 
-  @Test
+  @TestOfyOnly // relies on request-capturing datastore
   void testGet_cachedSecret() {
     int numInitialReads = RequestCapturingAsyncDatastoreService.getReads().size();
     ServerSecret secret = ServerSecret.get();
@@ -63,21 +64,9 @@ public class ServerSecretTest extends EntityTestCase {
     assertThat(RequestCapturingAsyncDatastoreService.getReads()).hasSize(numReads);
   }
 
-  @Test
+  @TestOfyAndSql
   void testAsBytes() {
     byte[] bytes = ServerSecret.create(new UUID(123, 0x456)).asBytes();
     assertThat(bytes).isEqualTo(new byte[] {0, 0, 0, 0, 0, 0, 0, 123, 0, 0, 0, 0, 0, 0, 0x4, 0x56});
   }
-
-  private static ServerSecret loadFromSql() {
-    return jpaTm()
-        .transact(
-            () ->
-                jpaTm()
-                    .query("FROM ServerSecret", ServerSecret.class)
-                    .setMaxResults(1)
-                    .getResultStream()
-                    .findFirst()
-                    .get());
-  }
 }

core/src/test/java/google/registry/model/tmch/TmchCrlTest.java

@@ -15,61 +15,43 @@
 package google.registry.model.tmch;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 
 import google.registry.model.EntityTestCase;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import java.util.Optional;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link TmchCrl}. */
+@DualDatabaseTest
 public class TmchCrlTest extends EntityTestCase {
 
   TmchCrlTest() {
     super(JpaEntityCoverageCheck.ENABLED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() {
     assertThat(TmchCrl.get()).isEqualTo(Optional.empty());
     TmchCrl.set("lolcat", "https://lol.cat");
     assertThat(TmchCrl.get().get().getCrl()).isEqualTo("lolcat");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDualWrite() {
     TmchCrl expected = new TmchCrl();
     expected.crl = "lolcat";
     expected.url = "https://lol.cat";
     expected.updated = fakeClock.nowUtc();
     TmchCrl.set("lolcat", "https://lol.cat");
-    assertThat(ofy().load().entity(new TmchCrl()).now()).isEqualTo(expected);
-    assertThat(loadFromSql()).isEqualTo(expected);
+    assertThat(loadByEntity(new TmchCrl())).isEqualTo(expected);
   }
 
-  @Test
+  @TestOfyAndSql
   void testMultipleWrites() {
     TmchCrl.set("first", "https://first.cat");
     assertThat(TmchCrl.get().get().getCrl()).isEqualTo("first");
     TmchCrl.set("second", "https://second.cat");
     assertThat(TmchCrl.get().get().getCrl()).isEqualTo("second");
-    jpaTm()
-        .transact(
-            () ->
-                assertThat(
-                        jpaTm().query("SELECT COUNT(*) FROM TmchCrl", Long.class).getSingleResult())
-                    .isEqualTo(1L));
-  }
-
-  private static TmchCrl loadFromSql() {
-    return jpaTm()
-        .transact(
-            () ->
-                jpaTm()
-                    .query("FROM TmchCrl", TmchCrl.class)
-                    .setMaxResults(1)
-                    .getResultStream()
-                    .findFirst()
-                    .get());
   }
 }