Replace ofyOrJpaTm with tm().isOfy() conditionals (#870)

* Replace ofyOrJpaTm with tm().isOfy() conditionals

Replace existing ofyOrJpaTm() calls with conditionals (either "if" statements
or ternary expressions) gated on tm().isOfy().

core/src/main/java/google/registry/model/registry/Registries.java

@@ -25,7 +25,6 @@ import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
 import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -60,21 +59,19 @@ public final class Registries {
             tm().doTransactionless(
                     () -> {
                       ImmutableSet<String> tlds =
-                          ofyOrJpaTm(
+                          tm().isOfy()
-                              () ->
+                              ? ofy()
-                                  ofy()
+                                  .load()
-                                      .load()
+                                  .type(Registry.class)
-                                      .type(Registry.class)
+                                  .ancestor(getCrossTldKey())
-                                      .ancestor(getCrossTldKey())
+                                  .keys()
-                                      .keys()
+                                  .list()
-                                      .list()
+                                  .stream()
-                                      .stream()
+                                  .map(Key::getName)
-                                      .map(Key::getName)
+                                  .collect(toImmutableSet())
-                                      .collect(toImmutableSet()),
+                              : tm().loadAll(Registry.class).stream()
-                              () ->
+                                  .map(Registry::getTldStr)
-                                  tm().loadAll(Registry.class).stream()
+                                  .collect(toImmutableSet());
-                                      .map(Registry::getTldStr)
-                                      .collect(toImmutableSet()));
                       return Registry.getAll(tlds).stream()
                           .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
                           .collect(entriesToImmutableMap());

core/src/main/java/google/registry/model/registry/label/PremiumList.java

@@ -25,7 +25,6 @@ import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Splitter;
@@ -209,9 +208,9 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   }
 
   private static PremiumList loadPremiumList(String name) {
-    return ofyOrJpaTm(
+    return tm().isOfy()
-        () -> ofy().load().type(PremiumList.class).parent(getCrossTldKey()).id(name).now(),
+        ? ofy().load().type(PremiumList.class).parent(getCrossTldKey()).id(name).now()
-        () -> PremiumListDao.getLatestRevision(name).orElseThrow(NoSuchElementException::new));
+        : PremiumListDao.getLatestRevision(name).orElseThrow(NoSuchElementException::new);
   }
 
   /**

core/src/main/java/google/registry/model/registry/label/ReservedList.java

@@ -19,7 +19,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
 import static google.registry.model.registry.label.ReservationType.FULLY_BLOCKED;
-import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static org.joda.time.DateTimeZone.UTC;
 
@@ -247,9 +247,9 @@ public final class ReservedList
               new CacheLoader<String, ReservedList>() {
                 @Override
                 public ReservedList load(String listName) {
-                  return ofyOrJpaTm(
+                  return tm().isOfy()
-                      () -> ReservedListDualWriteDao.getLatestRevision(listName).orElse(null),
+                      ? ReservedListDualWriteDao.getLatestRevision(listName).orElse(null)
-                      () -> ReservedListSqlDao.getLatestRevision(listName).orElse(null));
+                      : ReservedListSqlDao.getLatestRevision(listName).orElse(null);
                 }
               });
 

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

@@ -17,6 +17,7 @@ package google.registry.persistence.transaction;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.VKey;
 import java.util.NoSuchElementException;
 import java.util.Optional;
@@ -236,4 +237,9 @@ public interface TransactionManager {
 
   /** Clears the session cache if the underlying database is Datastore, otherwise it is a no-op. */
   void clearSessionCache();
+
+  /** Returns true if the transaction manager is DatastoreTransactionManager. */
+  default boolean isOfy() {
+    return this instanceof DatastoreTransactionManager;
+  }
 }

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

@@ -50,40 +50,6 @@ public class TransactionManagerUtil {
         });
   }
 
-  /**
-   * Executes either {@code ofyRunnable} if {@link TransactionManagerFactory#tm()} returns a {@link
-   * JpaTransactionManager} instance, or {@code jpaRunnable} if {@link
-   * TransactionManagerFactory#tm()} returns a {@link DatastoreTransactionManager} instance.
-   */
-  public static void ofyOrJpaTm(Runnable ofyRunnable, Runnable jpaRunnable) {
-    ofyOrJpaTm(
-        () -> {
-          ofyRunnable.run();
-          return null;
-        },
-        () -> {
-          jpaRunnable.run();
-          return null;
-        });
-  }
-
-  /**
-   * Returns the result from either {@code ofySupplier} if {@link TransactionManagerFactory#tm()}
-   * returns a {@link JpaTransactionManager} instance, or {@code jpaSupplier} if {@link
-   * TransactionManagerFactory#tm()} returns a {@link DatastoreTransactionManager} instance.
-   */
-  public static <T> T ofyOrJpaTm(Supplier<T> ofySupplier, Supplier<T> jpaSupplier) {
-    if (tm() instanceof DatastoreTransactionManager) {
-      return ofySupplier.get();
-    } else if (tm() instanceof JpaTransactionManager) {
-      return jpaSupplier.get();
-    } else {
-      throw new IllegalStateException(
-          "Expected tm() to be DatastoreTransactionManager or JpaTransactionManager, but got "
-              + tm().getClass());
-    }
-  }
-
   /**
    * Executes the given {@link Runnable} if {@link TransactionManagerFactory#tm()} returns a {@link
    * DatastoreTransactionManager} instance, otherwise does nothing.

core/src/test/java/google/registry/testing/AppEngineExtension.java

@@ -17,7 +17,7 @@ package google.registry.testing;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.io.Files.asCharSink;
 import static com.google.common.truth.Truth.assertWithMessage;
-import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.persistSimpleResources;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.injectTmForDualDatabaseTest;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.restoreTmAfterDualDatabaseTest;
@@ -385,17 +385,15 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (isWithDatastoreAndCloudSql()) {
       injectTmForDualDatabaseTest(context);
     }
-    ofyOrJpaTm(
+    if (tm().isOfy()) {
-        () -> {
+      if (withDatastore && !withoutCannedData) {
-          if (withDatastore && !withoutCannedData) {
+        loadInitialData();
-            loadInitialData();
+      }
-          }
+    } else {
-        },
+      if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
-        () -> {
+        loadInitialData();
-          if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
+      }
-            loadInitialData();
+    }
-          }
-        });
   }
 
   /**

core/src/test/java/google/registry/testing/DatastoreHelper.java

@@ -34,7 +34,6 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.label.PremiumListUtils.parentPremiumListEntriesOnRevision;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.persistence.transaction.TransactionManagerUtil.ofyOrJpaTm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
@@ -329,12 +328,13 @@ public class DatastoreHelper {
             .setShouldPublish(shouldPublish)
             .setLastUpdateTime(DateTime.now(DateTimeZone.UTC))
             .build();
-    return ofyOrJpaTm(
+    return tm().isOfy()
-        () -> persistResource(reservedList),
+        ? persistResource(reservedList)
-        () -> {
+        : tm().transact(
-          tm().transact(() -> tm().insert(reservedList));
+                () -> {
-          return reservedList;
+                  tm().insert(reservedList);
-        });
+                  return reservedList;
+                });
   }
 
   /**
@@ -364,15 +364,15 @@ public class DatastoreHelper {
             .build();
     PremiumListRevision revision = PremiumListRevision.create(premiumList, entries.keySet());
 
-    ofyOrJpaTm(
+    if (tm().isOfy()) {
-        () -> {
+      tm().putAllWithoutBackup(
-          tm().putAllWithoutBackup(
+              ImmutableList.of(
-                  ImmutableList.of(
+                  premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision));
-                      premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision));
+      tm().putAllWithoutBackup(
-          tm().putAllWithoutBackup(
+              parentPremiumListEntriesOnRevision(entries.values(), Key.create(revision)));
-                  parentPremiumListEntriesOnRevision(entries.values(), Key.create(revision)));
+    } else {
-        },
+      tm().transact(() -> tm().insert(premiumList));
-        () -> tm().transact(() -> tm().insert(premiumList)));
+    }
     // The above premiumList is in the session cache and it is different from the corresponding
     // entity stored in Datastore because it has some @Ignore fields set dedicated for SQL. This
     // breaks the assumption we have in our application code, see
@@ -933,17 +933,17 @@ public class DatastoreHelper {
   }
 
   private static <R> void saveResource(R resource, boolean wantBackup) {
-    ofyOrJpaTm(
+    if (tm().isOfy()) {
-        () -> {
+      Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup();
-          Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup();
+      saver.entity(resource);
-          saver.entity(resource);
+      if (resource instanceof EppResource) {
-          if (resource instanceof EppResource) {
+        EppResource eppResource = (EppResource) resource;
-            EppResource eppResource = (EppResource) resource;
+        persistEppResourceExtras(
-            persistEppResourceExtras(
+            eppResource, EppResourceIndex.create(Key.create(eppResource)), saver);
-                eppResource, EppResourceIndex.create(Key.create(eppResource)), saver);
+      }
-          }
+    } else {
-        },
+      tm().put(resource);
-        () -> tm().put(resource));
+    }
   }
 
   private static <R extends EppResource> void persistEppResourceExtras(
@@ -975,14 +975,15 @@ public class DatastoreHelper {
     final EppResourceIndex eppResourceIndex =
         EppResourceIndex.create(Key.create(EppResourceIndexBucket.class, 1), Key.create(resource));
     tm().transact(
-            () ->
+            () -> {
-                ofyOrJpaTm(
+              if (tm().isOfy()) {
-                    () -> {
+                Saver saver = ofy().save();
-                      Saver saver = ofy().save();
+                saver.entity(resource);
-                      saver.entity(resource);
+                persistEppResourceExtras(resource, eppResourceIndex, saver);
-                      persistEppResourceExtras(resource, eppResourceIndex, saver);
+              } else {
-                    },
+                tm().put(resource);
-                    () -> tm().put(resource)));
+              }
+            });
     tm().clearSessionCache();
     return transactIfJpaTm(() -> tm().load(resource));
   }
@@ -1040,38 +1041,28 @@ public class DatastoreHelper {
 
   /** Returns all of the history entries that are parented off the given EppResource. */
   public static List<? extends HistoryEntry> getHistoryEntries(EppResource resource) {
-    return ofyOrJpaTm(
+    return tm().isOfy()
-        () ->
+        ? ofy().load().type(HistoryEntry.class).ancestor(resource).order("modificationTime").list()
-            ofy()
+        : tm().transact(
-                .load()
+                () -> {
-                .type(HistoryEntry.class)
+                  ImmutableList<? extends HistoryEntry> unsorted = null;
-                .ancestor(resource)
+                  if (resource instanceof ContactBase) {
-                .order("modificationTime")
+                    unsorted = tm().loadAll(ContactHistory.class);
-                .list(),
+                  } else if (resource instanceof HostBase) {
-        () ->
+                    unsorted = tm().loadAll(HostHistory.class);
-            tm().transact(
+                  } else if (resource instanceof DomainContent) {
-                    () -> {
+                    unsorted = tm().loadAll(DomainHistory.class);
-                      ImmutableList<? extends HistoryEntry> unsorted = null;
+                  } else {
-                      if (resource instanceof ContactBase) {
+                    fail("Expected an EppResource instance, but got " + resource.getClass());
-                        unsorted = tm().loadAll(ContactHistory.class);
+                  }
-                      } else if (resource instanceof HostBase) {
+                  ImmutableList<? extends HistoryEntry> filtered =
-                        unsorted = tm().loadAll(HostHistory.class);
+                      unsorted.stream()
-                      } else if (resource instanceof DomainContent) {
+                          .filter(
-                        unsorted = tm().loadAll(DomainHistory.class);
+                              historyEntry ->
-                      } else {
+                                  historyEntry.getParent().getName().equals(resource.getRepoId()))
-                        fail("Expected an EppResource instance, but got " + resource.getClass());
+                          .collect(toImmutableList());
-                      }
+                  return ImmutableList.sortedCopyOf(DateTimeComparator.getInstance(), filtered);
-                      ImmutableList<? extends HistoryEntry> filtered =
+                });
-                          unsorted.stream()
-                              .filter(
-                                  historyEntry ->
-                                      historyEntry
-                                          .getParent()
-                                          .getName()
-                                          .equals(resource.getRepoId()))
-                              .collect(toImmutableList());
-                      return ImmutableList.sortedCopyOf(DateTimeComparator.getInstance(), filtered);
-                    }));
   }
 
   /**
@@ -1154,13 +1145,12 @@ public class DatastoreHelper {
   /** Force the create and update timestamps to get written into the resource. **/
   public static <R> R cloneAndSetAutoTimestamps(final R resource) {
     return tm().transact(
-            () ->
+            tm().isOfy()
-                ofyOrJpaTm(
+                ? () -> ofy().load().fromEntity(ofy().save().toEntity(resource))
-                    () -> ofy().load().fromEntity(ofy().save().toEntity(resource)),
+                : () -> {
-                    () -> {
+                  tm().put(resource);
-                      tm().put(resource);
+                  return tm().load(resource);
-                      return tm().load(resource);
+                });
-                    }));
   }
 
   /** Returns the entire map of {@link PremiumListEntry}s for the given {@link PremiumList}. */