Add common CRUD operations to TransactionManager (#487)

* Add BasicDao

* Refactor RegistrarDao to extend BasicDao

* Introduce VKey and rewrite BasicDao

* Move CRUD methods to TransactionManager

* Refactor code to simplify the way to get id from entity and sqlKey

* Assert in transaction

* Fix broken test

* Change methods name

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

@@ -16,7 +16,11 @@ package google.registry.model.ofy;
 
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
+import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.TransactionManager;
+import java.util.Optional;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;
 
@@ -83,4 +87,64 @@ public class DatastoreTransactionManager implements TransactionManager {
   public DateTime getTransactionTime() {
     return getOfy().getTransactionTime();
   }
+
+  @Override
+  public void saveNew(Object entity) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public void saveAllNew(ImmutableCollection<?> entities) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public void saveNewOrUpdate(Object entity) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public void saveNewOrUpdateAll(ImmutableCollection<?> entities) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public void update(Object entity) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public void updateAll(ImmutableCollection<?> entities) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public boolean checkExists(Object entity) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> boolean checkExists(VKey<T> key) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> Optional<T> load(VKey<T> key) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> int delete(VKey<T> key) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
+
+  @Override
+  public <T> void assertDelete(VKey<T> key) {
+    throw new UnsupportedOperationException("Not available in the Datastore transaction manager");
+  }
 }

core/src/main/java/google/registry/persistence/VKey.java

@@ -0,0 +1,68 @@
+// 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;
+
+import google.registry.model.ImmutableObject;
+
+/**
+ * VKey is an abstraction that encapsulates the key concept.
+ *
+ * <p>A VKey instance must contain both the JPA primary key for the referenced entity class and the
+ * objectify key for the object.
+ */
+public class VKey<T> extends ImmutableObject {
+
+  // The primary key for the referenced entity.
+  private Object primaryKey;
+
+  // The objectify key for the referenced entity.
+  private com.googlecode.objectify.Key<T> ofyKey;
+
+  private Class<? extends T> kind;
+
+  private VKey(Class<? extends T> kind, com.googlecode.objectify.Key<T> ofyKey, Object primaryKey) {
+    this.kind = kind;
+    this.ofyKey = ofyKey;
+    this.primaryKey = primaryKey;
+  }
+
+  public static <T> VKey<T> create(
+      Class<? extends T> kind, com.googlecode.objectify.Key<T> ofyKey, Object primaryKey) {
+    return new VKey(kind, ofyKey, primaryKey);
+  }
+
+  public static <T> VKey<T> create(Class<? extends T> kind, Object primaryKey) {
+    return new VKey(kind, null, primaryKey);
+  }
+
+  public static <T> VKey<T> create(
+      Class<? extends T> kind, com.googlecode.objectify.Key<T> ofyKey) {
+    return new VKey(kind, ofyKey, null);
+  }
+
+  public Class<? extends T> getKind() {
+    return this.kind;
+  }
+
+  /** Returns the SQL primary key. */
+  public Object getSqlKey() {
+    return this.primaryKey;
+  }
+
+  /** Returns the objectify key. */
+  public com.googlecode.objectify.Key<T> getOfyKey() {
+    return this.ofyKey;
+  }
+}

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

@@ -18,6 +18,7 @@ import javax.persistence.EntityManager;
 
 /** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
 public interface JpaTransactionManager extends TransactionManager {
+
   /** Returns the {@link EntityManager} for the current request. */
   EntityManager getEntityManager();
 }

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

@@ -14,13 +14,28 @@
 
 package google.registry.persistence.transaction;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
+import static java.util.stream.Collectors.joining;
+
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
+import google.registry.persistence.VKey;
 import google.registry.util.Clock;
+import java.lang.reflect.Field;
+import java.util.Optional;
 import java.util.function.Supplier;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.EntityTransaction;
 import javax.persistence.PersistenceException;
+import javax.persistence.Query;
+import javax.persistence.TypedQuery;
+import javax.persistence.metamodel.EntityType;
+import javax.persistence.metamodel.SingularAttribute;
 import org.joda.time.DateTime;
 
 /** Implementation of {@link JpaTransactionManager} for JPA compatible database. */
@@ -142,6 +157,180 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     return txnInfo.transactionTime;
   }
 
+  @Override
+  public void saveNew(Object entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    getEntityManager().persist(entity);
+  }
+
+  @Override
+  public void saveAllNew(ImmutableCollection<?> entities) {
+    checkArgumentNotNull(entities, "entities must be specified");
+    assertInTransaction();
+    entities.forEach(this::saveNew);
+  }
+
+  @Override
+  public void saveNewOrUpdate(Object entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    getEntityManager().merge(entity);
+  }
+
+  @Override
+  public void saveNewOrUpdateAll(ImmutableCollection<?> entities) {
+    checkArgumentNotNull(entities, "entities must be specified");
+    assertInTransaction();
+    entities.forEach(this::saveNewOrUpdate);
+  }
+
+  @Override
+  public void update(Object entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    checkArgument(checkExists(entity), "Given entity does not exist");
+    getEntityManager().merge(entity);
+  }
+
+  @Override
+  public void updateAll(ImmutableCollection<?> entities) {
+    checkArgumentNotNull(entities, "entities must be specified");
+    assertInTransaction();
+    entities.forEach(this::update);
+  }
+
+  @Override
+  public <T> boolean checkExists(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    EntityType<?> entityType = getEntityType(key.getKind());
+    ImmutableSet<EntityId> entityIds = getEntityIdsFromSqlKey(entityType, key.getSqlKey());
+    return checkExists(entityType.getName(), entityIds);
+  }
+
+  @Override
+  public boolean checkExists(Object entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    EntityType<?> entityType = getEntityType(entity.getClass());
+    ImmutableSet<EntityId> entityIds = getEntityIdsFromEntity(entityType, entity);
+    return checkExists(entityType.getName(), entityIds);
+  }
+
+  private boolean checkExists(String entityName, ImmutableSet<EntityId> entityIds) {
+    assertInTransaction();
+    TypedQuery<Integer> query =
+        getEntityManager()
+            .createQuery(
+                String.format("SELECT 1 FROM %s WHERE %s", entityName, getAndClause(entityIds)),
+                Integer.class)
+            .setMaxResults(1);
+    entityIds.forEach(entityId -> query.setParameter(entityId.name, entityId.value));
+    return query.getResultList().size() > 0;
+  }
+
+  @Override
+  public <T> Optional<T> load(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
+    checkArgumentNotNull(clazz, "clazz must be specified");
+    assertInTransaction();
+    return ImmutableList.copyOf(
+        getEntityManager()
+            .createQuery(
+                String.format("SELECT entity FROM %s entity", getEntityType(clazz).getName()),
+                clazz)
+            .getResultList());
+  }
+
+  @Override
+  public <T> int delete(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    EntityType<?> entityType = getEntityType(key.getKind());
+    ImmutableSet<EntityId> entityIds = getEntityIdsFromSqlKey(entityType, key.getSqlKey());
+    String sql =
+        String.format("DELETE FROM %s WHERE %s", entityType.getName(), getAndClause(entityIds));
+    Query query = getEntityManager().createQuery(sql);
+    entityIds.forEach(entityId -> query.setParameter(entityId.name, entityId.value));
+    return query.executeUpdate();
+  }
+
+  @Override
+  public <T> void assertDelete(VKey<T> key) {
+    if (delete(key) != 1) {
+      throw new IllegalArgumentException(
+          String.format("Error deleting the entity of the key: %s", key.getSqlKey()));
+    }
+  }
+
+  private <T> EntityType<T> getEntityType(Class<T> clazz) {
+    return emf.getMetamodel().entity(clazz);
+  }
+
+  private static class EntityId {
+    private String name;
+    private Object value;
+
+    private EntityId(String name, Object value) {
+      this.name = name;
+      this.value = value;
+    }
+  }
+
+  private static ImmutableSet<EntityId> getEntityIdsFromEntity(
+      EntityType<?> entityType, Object entity) {
+    if (entityType.hasSingleIdAttribute()) {
+      String idName = entityType.getDeclaredId(entityType.getIdType().getJavaType()).getName();
+      Object idValue = getFieldValue(entity, idName);
+      return ImmutableSet.of(new EntityId(idName, idValue));
+    } else {
+      return getEntityIdsFromIdContainer(entityType, entity);
+    }
+  }
+
+  private static ImmutableSet<EntityId> getEntityIdsFromSqlKey(
+      EntityType<?> entityType, Object sqlKey) {
+    if (entityType.hasSingleIdAttribute()) {
+      String idName = entityType.getDeclaredId(entityType.getIdType().getJavaType()).getName();
+      return ImmutableSet.of(new EntityId(idName, sqlKey));
+    } else {
+      return getEntityIdsFromIdContainer(entityType, sqlKey);
+    }
+  }
+
+  private static ImmutableSet<EntityId> getEntityIdsFromIdContainer(
+      EntityType<?> entityType, Object idContainer) {
+    return entityType.getIdClassAttributes().stream()
+        .map(SingularAttribute::getName)
+        .map(
+            idName -> {
+              Object idValue = getFieldValue(idContainer, idName);
+              return new EntityId(idName, idValue);
+            })
+        .collect(toImmutableSet());
+  }
+
+  private String getAndClause(ImmutableSet<EntityId> entityIds) {
+    return entityIds.stream()
+        .map(entityId -> String.format("%s = :%s", entityId.name, entityId.name))
+        .collect(joining(" AND "));
+  }
+
+  private static Object getFieldValue(Object object, String fieldName) {
+    try {
+      Field field = object.getClass().getDeclaredField(fieldName);
+      field.setAccessible(true);
+      return field.get(object);
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      throw new IllegalArgumentException(e);
+    }
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;

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

@@ -14,6 +14,10 @@
 
 package google.registry.persistence.transaction;
 
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
+import google.registry.persistence.VKey;
+import java.util.Optional;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;
 
@@ -78,4 +82,40 @@ public interface TransactionManager {
 
   /** Returns the time associated with the start of this particular transaction attempt. */
   DateTime getTransactionTime();
+
+  /** Persists a new entity in the database, throws exception if the entity already exists. */
+  void saveNew(Object entity);
+
+  /** Persists all new entities in the database, throws exception if any entity already exists. */
+  void saveAllNew(ImmutableCollection<?> entities);
+
+  /** Persists a new entity or update the existing entity in the database. */
+  void saveNewOrUpdate(Object entity);
+
+  /** Persists all new entities or update the existing entities in the database. */
+  void saveNewOrUpdateAll(ImmutableCollection<?> entities);
+
+  /** Updates an entity in the database, throws exception if the entity does not exist. */
+  void update(Object entity);
+
+  /** Updates all entities in the database, throws exception if any entity does not exist. */
+  void updateAll(ImmutableCollection<?> entities);
+
+  /** Returns whether the given entity with same ID exists. */
+  boolean checkExists(Object entity);
+
+  /** Returns whether the entity of given key exists. */
+  <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);
+
+  /** Loads all entities of the given type, returns empty if there is no such entity. */
+  <T> ImmutableList<T> loadAll(Class<T> clazz);
+
+  /** Deletes the entity by its id, returns the number of deleted entity. */
+  <T> int delete(VKey<T> key);
+
+  /** Deletes the entity by its id, throws exception if the entity is not deleted. */
+  <T> void assertDelete(VKey<T> key);
 }

core/src/main/java/google/registry/schema/registrar/RegistrarDao.java

@@ -1,73 +0,0 @@
-// 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.schema.registrar;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
-
-import google.registry.model.registrar.Registrar;
-import java.util.Optional;
-
-/** Data access object for {@link Registrar}. */
-public class RegistrarDao {
-
-  private RegistrarDao() {}
-
-  /** Persists a new or updates an existing registrar in Cloud SQL. */
-  public static void saveNew(Registrar registrar) {
-    checkArgumentNotNull(registrar, "registrar must be specified");
-    jpaTm().transact(() -> jpaTm().getEntityManager().persist(registrar));
-  }
-
-  /** Updates an existing registrar in Cloud SQL, throws excpetion if it does not exist. */
-  public static void update(Registrar registrar) {
-    checkArgumentNotNull(registrar, "registrar must be specified");
-    jpaTm()
-        .transact(
-            () -> {
-              checkArgument(
-                  checkExists(registrar.getClientId()),
-                  "A registrar of this id does not exist: %s.",
-                  registrar.getClientId());
-              jpaTm().getEntityManager().merge(registrar);
-            });
-  }
-
-  /** Returns whether the registrar of the given id exists. */
-  public static boolean checkExists(String clientId) {
-    checkArgumentNotNull(clientId, "clientId must be specified");
-    return jpaTm()
-        .transact(
-            () ->
-                jpaTm()
-                        .getEntityManager()
-                        .createQuery(
-                            "SELECT 1 FROM Registrar WHERE clientIdentifier = :clientIdentifier",
-                            Integer.class)
-                        .setParameter("clientIdentifier", clientId)
-                        .setMaxResults(1)
-                        .getResultList()
-                        .size()
-                    > 0);
-  }
-
-  /** Loads the registrar by its id, returns empty if it doesn't exist. */
-  public static Optional<Registrar> load(String clientId) {
-    checkArgumentNotNull(clientId, "clientId must be specified");
-    return Optional.ofNullable(
-        jpaTm().transact(() -> jpaTm().getEntityManager().find(Registrar.class, clientId)));
-  }
-}

core/src/main/java/google/registry/tools/CreateRegistrarCommand.java

@@ -19,6 +19,7 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static google.registry.model.registrar.Registrar.State.ACTIVE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.tools.RegistryToolEnvironment.PRODUCTION;
 import static google.registry.tools.RegistryToolEnvironment.SANDBOX;
 import static google.registry.tools.RegistryToolEnvironment.UNITTEST;
@@ -32,7 +33,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.registrar.Registrar;
-import google.registry.schema.registrar.RegistrarDao;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -72,7 +72,7 @@ final class CreateRegistrarCommand extends CreateOrUpdateRegistrarCommand
 
   @Override
   void saveToCloudSql(Registrar registrar) {
-    RegistrarDao.saveNew(registrar);
+    jpaTm().saveNew(registrar);
   }
 
   @Nullable

core/src/main/java/google/registry/tools/UpdateRegistrarCommand.java

@@ -14,13 +14,13 @@
 
 package google.registry.tools;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
 
 import com.beust.jcommander.Parameters;
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.registrar.Registrar;
-import google.registry.schema.registrar.RegistrarDao;
 import javax.annotation.Nullable;
 
 /** Command to update a Registrar. */
@@ -53,6 +53,6 @@ final class UpdateRegistrarCommand extends CreateOrUpdateRegistrarCommand {
 
   @Override
   void saveToCloudSql(Registrar registrar) {
-    RegistrarDao.update(registrar);
+    jpaTm().update(registrar);
   }
 }

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

@@ -19,11 +19,19 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.Assert.assertThrows;
 
+import com.google.common.collect.ImmutableList;
+import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestRule;
 import google.registry.testing.FakeClock;
+import java.io.Serializable;
 import java.math.BigInteger;
+import javax.persistence.Entity;
 import javax.persistence.EntityManager;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
 import javax.persistence.PersistenceException;
+import javax.persistence.RollbackException;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -34,12 +42,24 @@ import org.junit.runners.JUnit4;
 public class JpaTransactionManagerImplTest {
 
   private final FakeClock fakeClock = new FakeClock();
+  private final TestEntity theEntity = new TestEntity("theEntity", "foo");
+  private final VKey<TestEntity> theEntityKey = VKey.create(TestEntity.class, "theEntity");
+  private final TestCompoundIdEntity compoundIdEntity =
+      new TestCompoundIdEntity("compoundIdEntity", 10, "foo");
+  private final VKey<TestCompoundIdEntity> compoundIdEntityKey =
+      VKey.create(TestCompoundIdEntity.class, new CompoundId("compoundIdEntity", 10));
+  private final ImmutableList<TestEntity> moreEntities =
+      ImmutableList.of(
+          new TestEntity("entity1", "foo"),
+          new TestEntity("entity2", "bar"),
+          new TestEntity("entity3", "qux"));
 
   @Rule
   public final JpaUnitTestRule jpaRule =
       new JpaTestRules.Builder()
           .withInitScript(fileClassPath(getClass(), "test_schema.sql"))
           .withClock(fakeClock)
+          .withEntityClass(TestEntity.class, TestCompoundIdEntity.class)
           .buildUnitTestRule();
 
   @Test
@@ -122,6 +142,203 @@ public class JpaTransactionManagerImplTest {
     assertCompanyExist("Bar");
   }
 
+  @Test
+  public void saveNew_succeeds() {
+    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);
+  }
+
+  @Test
+  public void saveNew_throwsExceptionIfEntityExists() {
+    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);
+    assertThrows(RollbackException.class, () -> jpaTm().transact(() -> jpaTm().saveNew(theEntity)));
+  }
+
+  @Test
+  public void createCompoundIdEntity_succeeds() {
+    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()))
+        .isEqualTo(compoundIdEntity);
+  }
+
+  @Test
+  public void saveAllNew_succeeds() {
+    moreEntities.forEach(
+        entity -> assertThat(jpaTm().transact(() -> jpaTm().checkExists(entity))).isFalse());
+    jpaTm().transact(() -> jpaTm().saveAllNew(moreEntities));
+    moreEntities.forEach(
+        entity -> assertThat(jpaTm().transact(() -> jpaTm().checkExists(entity))).isTrue());
+    assertThat(jpaTm().transact(() -> jpaTm().loadAll(TestEntity.class)))
+        .containsExactlyElementsIn(moreEntities);
+  }
+
+  @Test
+  public void saveAllNew_rollsBackWhenFailure() {
+    moreEntities.forEach(
+        entity -> assertThat(jpaTm().transact(() -> jpaTm().checkExists(entity))).isFalse());
+    jpaTm().transact(() -> jpaTm().saveNew(moreEntities.get(0)));
+    assertThrows(
+        RollbackException.class, () -> jpaTm().transact(() -> jpaTm().saveAllNew(moreEntities)));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(moreEntities.get(0)))).isTrue();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(moreEntities.get(1)))).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(moreEntities.get(2)))).isFalse();
+  }
+
+  @Test
+  public void saveNewOrUpdate_persistsNewEntity() {
+    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);
+  }
+
+  @Test
+  public void saveNewOrUpdate_updatesExistingEntity() {
+    jpaTm().transact(() -> jpaTm().saveNew(theEntity));
+    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    assertThat(persisted.data).isEqualTo("foo");
+    theEntity.data = "bar";
+    jpaTm().transact(() -> jpaTm().saveNewOrUpdate(theEntity));
+    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    assertThat(persisted.data).isEqualTo("bar");
+  }
+
+  @Test
+  public void saveNewOrUpdateAll_succeeds() {
+    moreEntities.forEach(
+        entity -> assertThat(jpaTm().transact(() -> jpaTm().checkExists(entity))).isFalse());
+    jpaTm().transact(() -> jpaTm().saveNewOrUpdateAll(moreEntities));
+    moreEntities.forEach(
+        entity -> assertThat(jpaTm().transact(() -> jpaTm().checkExists(entity))).isTrue());
+    assertThat(jpaTm().transact(() -> jpaTm().loadAll(TestEntity.class)))
+        .containsExactlyElementsIn(moreEntities);
+  }
+
+  @Test
+  public void update_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveNew(theEntity));
+    TestEntity persisted =
+        jpaTm().transact(() -> jpaTm().load(VKey.create(TestEntity.class, "theEntity"))).get();
+    assertThat(persisted.data).isEqualTo("foo");
+    theEntity.data = "bar";
+    jpaTm().transact(() -> jpaTm().update(theEntity));
+    persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    assertThat(persisted.data).isEqualTo("bar");
+  }
+
+  @Test
+  public void updateCompoundIdEntity_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
+    TestCompoundIdEntity persisted =
+        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());
+    assertThat(persisted.data).isEqualTo("bar");
+  }
+
+  @Test
+  public void update_throwsExceptionWhenEntityDoesNotExist() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    assertThrows(
+        IllegalArgumentException.class, () -> jpaTm().transact(() -> jpaTm().update(theEntity)));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+  }
+
+  @Test
+  public void updateAll_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveAllNew(moreEntities));
+    ImmutableList<TestEntity> updated =
+        ImmutableList.of(
+            new TestEntity("entity1", "foo_updated"),
+            new TestEntity("entity2", "bar_updated"),
+            new TestEntity("entity3", "qux_updated"));
+    jpaTm().transact(() -> jpaTm().updateAll(updated));
+    assertThat(jpaTm().transact(() -> jpaTm().loadAll(TestEntity.class)))
+        .containsExactlyElementsIn(updated);
+  }
+
+  @Test
+  public void updateAll_rollsBackWhenFailure() {
+    jpaTm().transact(() -> jpaTm().saveAllNew(moreEntities));
+    ImmutableList<TestEntity> updated =
+        ImmutableList.of(
+            new TestEntity("entity1", "foo_updated"),
+            new TestEntity("entity2", "bar_updated"),
+            new TestEntity("entity3", "qux_updated"),
+            theEntity);
+    assertThrows(
+        IllegalArgumentException.class, () -> jpaTm().transact(() -> jpaTm().updateAll(updated)));
+    assertThat(jpaTm().transact(() -> jpaTm().loadAll(TestEntity.class)))
+        .containsExactlyElementsIn(moreEntities);
+  }
+
+  @Test
+  public void load_succeeds() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    jpaTm().transact(() -> jpaTm().saveNew(theEntity));
+    TestEntity persisted = jpaTm().transact(() -> jpaTm().load(theEntityKey)).get();
+    assertThat(persisted.name).isEqualTo("theEntity");
+    assertThat(persisted.data).isEqualTo("foo");
+  }
+
+  @Test
+  public void loadCompoundIdEntity_succeeds() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isFalse();
+    jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
+    TestCompoundIdEntity persisted =
+        jpaTm().transact(() -> jpaTm().load(compoundIdEntityKey)).get();
+    assertThat(persisted.name).isEqualTo("compoundIdEntity");
+    assertThat(persisted.age).isEqualTo(10);
+    assertThat(persisted.data).isEqualTo("foo");
+  }
+
+  @Test
+  public void loadAll_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveAllNew(moreEntities));
+    ImmutableList<TestEntity> persisted = jpaTm().transact(() -> jpaTm().loadAll(TestEntity.class));
+    assertThat(persisted).containsExactlyElementsIn(moreEntities);
+  }
+
+  @Test
+  public void delete_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveNew(theEntity));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isTrue();
+    assertThat(jpaTm().transact(() -> jpaTm().delete(theEntityKey))).isEqualTo(1);
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+  }
+
+  @Test
+  public void delete_returnsZeroWhenNoEntity() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().delete(theEntityKey))).isEqualTo(0);
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+  }
+
+  @Test
+  public void deleteCompoundIdEntity_succeeds() {
+    jpaTm().transact(() -> jpaTm().saveNew(compoundIdEntity));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isTrue();
+    jpaTm().transact(() -> jpaTm().delete(compoundIdEntityKey));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(compoundIdEntity))).isFalse();
+  }
+
+  @Test
+  public void assertDelete_throwsExceptionWhenEntityNotDeleted() {
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> jpaTm().transact(() -> jpaTm().assertDelete(theEntityKey)));
+  }
+
   private void insertPerson(int age) {
     jpaTm()
         .getEntityManager()
@@ -194,4 +411,47 @@ public class JpaTransactionManagerImplTest {
               return colCount.intValue();
             });
   }
+
+  @Entity(name = "TestEntity")
+  private static class TestEntity extends ImmutableObject {
+    @Id private String name;
+
+    private String data;
+
+    private TestEntity() {}
+
+    private TestEntity(String name, String data) {
+      this.name = name;
+      this.data = data;
+    }
+  }
+
+  @Entity(name = "TestCompoundIdEntity")
+  @IdClass(CompoundId.class)
+  private static class TestCompoundIdEntity extends ImmutableObject {
+    @Id private String name;
+    @Id private int age;
+
+    private String data;
+
+    private TestCompoundIdEntity() {}
+
+    private TestCompoundIdEntity(String name, int age, String data) {
+      this.name = name;
+      this.age = age;
+      this.data = data;
+    }
+  }
+
+  private static class CompoundId implements Serializable {
+    String name;
+    int age;
+
+    private CompoundId() {}
+
+    private CompoundId(String name, int age) {
+      this.name = name;
+      this.age = age;
+    }
+  }
 }

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

@@ -15,12 +15,14 @@
 package google.registry.schema.registrar;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.model.EntityTestCase;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarAddress;
+import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageRule;
 import google.registry.testing.FakeClock;
@@ -33,6 +35,7 @@ import org.junit.runners.JUnit4;
 /** Unit tests for {@link RegistrarDao}. */
 @RunWith(JUnit4.class)
 public class RegistrarDaoTest extends EntityTestCase {
+  private final VKey<Registrar> registrarKey = VKey.create(Registrar.class, "registrarId");
   private final FakeClock fakeClock = new FakeClock();
 
   @Rule
@@ -61,32 +64,39 @@ public class RegistrarDaoTest extends EntityTestCase {
 
   @Test
   public void saveNew_worksSuccessfully() {
-    assertThat(RegistrarDao.checkExists("registrarId")).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
-    RegistrarDao.saveNew(testRegistrar);
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
-    assertThat(RegistrarDao.checkExists("registrarId")).isTrue();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isTrue();
   }
 
   @Test
   public void update_worksSuccessfully() {
-    RegistrarDao.saveNew(testRegistrar);
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
-    Registrar persisted = RegistrarDao.load("registrarId").get();
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");
-    RegistrarDao.update(persisted.asBuilder().setRegistrarName("changedRegistrarName").build());
+    jpaTm()
-    persisted = RegistrarDao.load("registrarId").get();
+        .transact(
-    assertThat(persisted.getRegistrarName()).isEqualTo("changedRegistrarName");
+            () ->
+                jpaTm()
+                    .update(
+                        persisted.asBuilder().setRegistrarName("changedRegistrarName").build()));
+    Registrar updated = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
+    assertThat(updated.getRegistrarName()).isEqualTo("changedRegistrarName");
   }
 
   @Test
   public void update_throwsExceptionWhenEntityDoesNotExist() {
-    assertThat(RegistrarDao.checkExists("registrarId")).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
-    assertThrows(IllegalArgumentException.class, () -> RegistrarDao.update(testRegistrar));
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> jpaTm().transact(() -> jpaTm().update(testRegistrar)));
   }
 
   @Test
   public void load_worksSuccessfully() {
-    assertThat(RegistrarDao.checkExists("registrarId")).isFalse();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
-    RegistrarDao.saveNew(testRegistrar);
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
-    Registrar persisted = RegistrarDao.load("registrarId").get();
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
 
     assertThat(persisted.getClientId()).isEqualTo("registrarId");
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");

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

@@ -17,6 +17,7 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT_HASH;
 import static google.registry.testing.DatastoreHelper.createTlds;
@@ -35,7 +36,6 @@ import com.google.common.net.MediaType;
 import google.registry.model.registrar.Registrar;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageRule;
-import google.registry.schema.registrar.RegistrarDao;
 import google.registry.testing.CertificateSamples;
 import google.registry.testing.FakeClock;
 import java.io.IOException;
@@ -130,7 +130,7 @@ public class CreateRegistrarCommandTest extends CommandTestCase<CreateRegistrarC
     Optional<Registrar> registrar = Registrar.loadByClientId("clientz");
     assertThat(registrar).isPresent();
     assertThat(registrar.get().verifyPassword("some_password")).isTrue();
-    assertThat(RegistrarDao.checkExists("clientz")).isTrue();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(registrar.get()))).isTrue();
   }
 
   @Test

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

@@ -17,6 +17,7 @@ package google.registry.tools;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT_HASH;
 import static google.registry.testing.DatastoreHelper.createTlds;
@@ -32,9 +33,9 @@ import com.google.common.collect.ImmutableSet;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import google.registry.model.registrar.Registrar.Type;
+import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageRule;
-import google.registry.schema.registrar.RegistrarDao;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.FakeClock;
 import google.registry.util.CidrAddressBlock;
@@ -56,10 +57,15 @@ public class UpdateRegistrarCommandTest extends CommandTestCase<UpdateRegistrarC
   @Test
   public void testSuccess_alsoUpdateInCloudSql() throws Exception {
     assertThat(loadRegistrar("NewRegistrar").verifyPassword("some_password")).isFalse();
-    RegistrarDao.saveNew(loadRegistrar("NewRegistrar"));
+    jpaTm().transact(() -> jpaTm().saveNew(loadRegistrar("NewRegistrar")));
     runCommand("--password=some_password", "--force", "NewRegistrar");
     assertThat(loadRegistrar("NewRegistrar").verifyPassword("some_password")).isTrue();
-    assertThat(RegistrarDao.load("NewRegistrar").get().verifyPassword("some_password")).isTrue();
+    assertThat(
+            jpaTm()
+                .transact(() -> jpaTm().load(VKey.create(Registrar.class, "NewRegistrar")))
+                .get()
+                .verifyPassword("some_password"))
+        .isTrue();
   }
 
   @Test