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