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,9 +59,8 @@ public final class Registries {
             tm().doTransactionless(
                     () -> {
                       ImmutableSet<String> tlds =
-                          ofyOrJpaTm(
+                          tm().isOfy()
-                              () ->
+                              ? ofy()
-                                  ofy()
                                   .load()
                                   .type(Registry.class)
                                   .ancestor(getCrossTldKey())
@@ -70,11 +68,10 @@ public final class Registries {
                                   .list()
                                   .stream()
                                   .map(Key::getName)
-                                      .collect(toImmutableSet()),
+                                  .collect(toImmutableSet())
-                              () ->
+                              : tm().loadAll(Registry.class).stream()
-                                  tm().loadAll(Registry.class).stream()
                                   .map(Registry::getTldStr)
-                                      .collect(toImmutableSet()));
+                                  .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) {
         loadInitialData();
       }
-        },
+    } else {
-        () -> {
       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,10 +328,11 @@ 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));
+                  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(
               ImmutableList.of(
                   premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision));
       tm().putAllWithoutBackup(
               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,8 +933,7 @@ public class DatastoreHelper {
   }
 
   private static <R> void saveResource(R resource, boolean wantBackup) {
-    ofyOrJpaTm(
+    if (tm().isOfy()) {
-        () -> {
       Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup();
       saver.entity(resource);
       if (resource instanceof EppResource) {
@@ -942,8 +941,9 @@ public class DatastoreHelper {
         persistEppResourceExtras(
             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.entity(resource);
                 persistEppResourceExtras(resource, eppResourceIndex, saver);
-                    },
+              } else {
-                    () -> tm().put(resource)));
+                tm().put(resource);
+              }
+            });
     tm().clearSessionCache();
     return transactIfJpaTm(() -> tm().load(resource));
   }
@@ -1040,16 +1041,9 @@ 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)
-                .ancestor(resource)
-                .order("modificationTime")
-                .list(),
-        () ->
-            tm().transact(
                 () -> {
                   ImmutableList<? extends HistoryEntry> unsorted = null;
                   if (resource instanceof ContactBase) {
@@ -1065,13 +1059,10 @@ public class DatastoreHelper {
                       unsorted.stream()
                           .filter(
                               historyEntry ->
-                                      historyEntry
+                                  historyEntry.getParent().getName().equals(resource.getRepoId()))
-                                          .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);
                   return tm().load(resource);
-                    }));
+                });
   }
 
   /** Returns the entire map of {@link PremiumListEntry}s for the given {@link PremiumList}. */