Improve return value semantices for tm().load() (#576)

Since we rarely (if ever) want to check the result of a single element load,
make TransactionManager.load(VKey) return non-optional, non-nullable and just
throw a NoSuchElementException if the entity is not found.

Also add a maybeLoad() that does return an optional in case we ever want to do
this (exists() should suffice for an existence check).

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

@@ -22,6 +22,7 @@ import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.TransactionManager;
 import java.util.Iterator;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.StreamSupport;
@@ -136,10 +137,19 @@ public class DatastoreTransactionManager implements TransactionManager {
   // VKey instead of by ofy Key.  But ideally, there should be one set of TransactionManager
   // interface tests that are applied to both the datastore and SQL implementations.
   @Override
-  public <T> Optional<T> load(VKey<T> key) {
+  public <T> Optional<T> maybeLoad(VKey<T> key) {
     return Optional.of(getOfy().load().key(key.getOfyKey()).now());
   }
 
+  @Override
+  public <T> T load(VKey<T> key) {
+    T result = getOfy().load().key(key.getOfyKey()).now();
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
   @Override
   public <T> ImmutableList<T> load(Iterable<VKey<T>> keys) {
     Iterator<Key<T>> iter =

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

@@ -232,12 +232,23 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public <T> Optional<T> load(VKey<T> key) {
+  public <T> Optional<T> maybeLoad(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
     return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
   }
 
+  @Override
+  public <T> T load(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
   @Override
   public <T> ImmutableList<T> load(Iterable<VKey<T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");

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

@@ -108,7 +108,10 @@ public interface TransactionManager {
   <T> boolean checkExists(VKey<T> key);
 
   /** Loads the entity by its id, returns empty if the entity doesn't exist. */
-  <T> Optional<T> load(VKey<T> key);
+  <T> Optional<T> maybeLoad(VKey<T> key);
+
+  /** Loads the entity by its id, throws NoSuchElementException if it doesn't exist. */
+  <T> T load(VKey<T> key);
 
   /**
    * Leads the set of entities by their key id.

core/src/test/java/google/registry/model/contact/ContactResourceTest.java

@@ -140,8 +140,7 @@ public class ContactResourceTest extends EntityTestCase {
             .transact(
                 () ->
                     jpaTm()
-                        .load(VKey.createSql(ContactResource.class, originalContact.getRepoId())))
+                        .load(VKey.createSql(ContactResource.class, originalContact.getRepoId())));
-            .get();
     // TODO(b/153378849): Remove the hard code for postal info after resolving the issue that
     // @PostLoad doesn't work in Address
     ContactResource fixed =

core/src/test/java/google/registry/persistence/EntityCallbacksListenerTest.java

@@ -69,14 +69,14 @@ public class EntityCallbacksListenerTest {
     checkAll(updated, 0, 1, 0, 1);
 
     TestEntity testLoad =
-        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestEntity.class, "id"))).get();
+        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestEntity.class, "id")));
     checkAll(testLoad, 0, 0, 0, 1);
 
     TestEntity testRemove =
         jpaTm()
             .transact(
                 () -> {
-                  TestEntity removed = jpaTm().load(VKey.createSql(TestEntity.class, "id")).get();
+                  TestEntity removed = jpaTm().load(VKey.createSql(TestEntity.class, "id"));
                   jpaTm().getEntityManager().remove(removed);
                   return removed;
                 });

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

@@ -26,6 +26,7 @@ import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
 import google.registry.testing.FakeClock;
 import java.io.Serializable;
 import java.math.BigInteger;
+import java.util.NoSuchElementException;
 import javax.persistence.Entity;
 import javax.persistence.EntityManager;
 import javax.persistence.Id;
@@ -147,7 +148,7 @@ public class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(theEntity));
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey).get())).isEqualTo(theEntity);
+    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey))).isEqualTo(theEntity);
   }
 
   @Test
@@ -155,7 +156,7 @@ public class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(theEntity));
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey).get())).isEqualTo(theEntity);
+    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey))).isEqualTo(theEntity);
     assertThrows(RollbackException.class, () -> jpaTm().transact(() -> jpaTm().saveNew(theEntity)));
   }
 
@@ -164,7 +165,7 @@ public class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey).get()))
+    assertThat(jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey)))
         .isEqualTo(compoundIdEntity);
   }
 
@@ -196,17 +197,17 @@ public class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNewOrUpdate(theEntity));
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey).get())).isEqualTo(theEntity);
+    assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey))).isEqualTo(theEntity);
   }
 
   @Test
   public void saveNewOrUpdate_updatesExistingEntity() {
     jpaTm().transact(() -> jpaTm().saveNew(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey));
     assertThat(persisted.data).isEqualTo("foo");
     theEntity.data = "bar";
     jpaTm().transact(() -> jpaTm().saveNewOrUpdate(theEntity));
-    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey));
     assertThat(persisted.data).isEqualTo("bar");
   }
 
@@ -225,23 +226,22 @@ public class JpaTransactionManagerImplTest {
   public void update_succeeds() {
     jpaTm().transact(() -> jpaTm().saveNew(theEntity));
     TestEntity persisted =
-        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestEntity.class, "theEntity"))).get();
+        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestEntity.class, "theEntity")));
     assertThat(persisted.data).isEqualTo("foo");
     theEntity.data = "bar";
     jpaTm().transact(() -> jpaTm().update(theEntity));
-    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey));
     assertThat(persisted.data).isEqualTo("bar");
   }
 
   @Test
   public void updateCompoundIdEntity_succeeds() {
     jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
-    TestCompoundIdEntity persisted =
+    TestCompoundIdEntity persisted = jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey));
-        jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey)).get();
     assertThat(persisted.data).isEqualTo("foo");
     compoundIdEntity.data = "bar";
     jpaTm().transact(() -> jpaTm().update(compoundIdEntity));
-    persisted = jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey).get());
+    persisted = jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey));
     assertThat(persisted.data).isEqualTo("bar");
   }
 
@@ -285,17 +285,38 @@ public class JpaTransactionManagerImplTest {
   public void load_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
 
+  @Test
+  public void load_throwsOnMissingElement() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    assertThrows(
+        NoSuchElementException.class, () -> jpaTm().transact(() -> jpaTm().load(theEntityKey)));
+  }
+
+  @Test
+  public void maybeLoad_succeeds() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    jpaTm().transact(() -> jpaTm().saveNew(theEntity));
+    TestEntity persisted = jpaTm().transact(() -> jpaTm().maybeLoad(theEntityKey).get());
+    assertThat(persisted.name).isEqualTo("theEntity");
+    assertThat(persisted.data).isEqualTo("foo");
+  }
+
+  @Test
+  public void maybeLoad_nonExistentObject() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().maybeLoad(theEntityKey)).isPresent()).isFalse();
+  }
+
   @Test
   public void loadCompoundIdEntity_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
-    TestCompoundIdEntity persisted =
+    TestCompoundIdEntity persisted = jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey));
-        jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey)).get();
     assertThat(persisted.name).isEqualTo("compoundIdEntity");
     assertThat(persisted.age).isEqualTo(10);
     assertThat(persisted.data).isEqualTo("foo");

core/src/test/java/google/registry/schema/registrar/RegistrarDaoTest.java

@@ -78,7 +78,7 @@ public class RegistrarDaoTest {
   @Test
   public void update_worksSuccessfully() {
     jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
-    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey));
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");
     jpaTm()
         .transact(
@@ -86,7 +86,7 @@ public class RegistrarDaoTest {
                 jpaTm()
                     .update(
                         persisted.asBuilder().setRegistrarName("changedRegistrarName").build()));
-    Registrar updated = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
+    Registrar updated = jpaTm().transact(() -> jpaTm().load(registrarKey));
     assertThat(updated.getRegistrarName()).isEqualTo("changedRegistrarName");
   }
 
@@ -102,7 +102,7 @@ public class RegistrarDaoTest {
   public void load_worksSuccessfully() {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
     jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
-    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey));
 
     assertThat(persisted.getClientId()).isEqualTo("registrarId");
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");

core/src/test/java/google/registry/tools/UpdateRegistrarCommandTest.java

@@ -53,7 +53,6 @@ public class UpdateRegistrarCommandTest extends CommandTestCase<UpdateRegistrarC
     assertThat(
             jpaTm()
                 .transact(() -> jpaTm().load(VKey.createSql(Registrar.class, "NewRegistrar")))
-                .get()
                 .verifyPassword("some_password"))
         .isTrue();
   }