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
2025-07-09 04:33:28 +02:00 · 2020-03-05 14:03:03 -05:00 · 2020-03-05 14:03:03 -05:00 · adafab60c4
commit adafab60c4
parent fec806ef8b
12 changed files with 660 additions and 95 deletions
--- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java
+++ b/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;
+    }
+  }
 }
--- a/core/src/test/java/google/registry/schema/registrar/RegistrarDaoTest.java
+++ b/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();
-    RegistrarDao.saveNew(testRegistrar);
-    assertThat(RegistrarDao.checkExists("registrarId")).isTrue();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isTrue();
  }

  @Test
  public void update_worksSuccessfully() {
-    RegistrarDao.saveNew(testRegistrar);
-    Registrar persisted = RegistrarDao.load("registrarId").get();
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();
    assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");
-    RegistrarDao.update(persisted.asBuilder().setRegistrarName("changedRegistrarName").build());
-    persisted = RegistrarDao.load("registrarId").get();
-    assertThat(persisted.getRegistrarName()).isEqualTo("changedRegistrarName");
+    jpaTm()
+        .transact(
+            () ->
+                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();
-    assertThrows(IllegalArgumentException.class, () -> RegistrarDao.update(testRegistrar));
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> jpaTm().transact(() -> jpaTm().update(testRegistrar)));
  }

  @Test
  public void load_worksSuccessfully() {
-    assertThat(RegistrarDao.checkExists("registrarId")).isFalse();
-    RegistrarDao.saveNew(testRegistrar);
-    Registrar persisted = RegistrarDao.load("registrarId").get();
+    assertThat(jpaTm().transact(() -> jpaTm().checkExists(testRegistrar))).isFalse();
+    jpaTm().transact(() -> jpaTm().saveNew(testRegistrar));
+    Registrar persisted = jpaTm().transact(() -> jpaTm().load(registrarKey)).get();

    assertThat(persisted.getClientId()).isEqualTo("registrarId");
    assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");
--- a/core/src/test/java/google/registry/tools/CreateRegistrarCommandTest.java
+++ b/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
--- a/core/src/test/java/google/registry/tools/UpdateRegistrarCommandTest.java
+++ b/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