diff --git a/core/src/main/java/google/registry/model/registry/Registries.java b/core/src/main/java/google/registry/model/registry/Registries.java index 599225428..5d7c10a4e 100644 --- a/core/src/main/java/google/registry/model/registry/Registries.java +++ b/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 tlds = - ofyOrJpaTm( - () -> - ofy() - .load() - .type(Registry.class) - .ancestor(getCrossTldKey()) - .keys() - .list() - .stream() - .map(Key::getName) - .collect(toImmutableSet()), - () -> - tm().loadAll(Registry.class).stream() - .map(Registry::getTldStr) - .collect(toImmutableSet())); + tm().isOfy() + ? ofy() + .load() + .type(Registry.class) + .ancestor(getCrossTldKey()) + .keys() + .list() + .stream() + .map(Key::getName) + .collect(toImmutableSet()) + : tm().loadAll(Registry.class).stream() + .map(Registry::getTldStr) + .collect(toImmutableSet()); return Registry.getAll(tlds).stream() .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType())) .collect(entriesToImmutableMap()); diff --git a/core/src/main/java/google/registry/model/registry/label/PremiumList.java b/core/src/main/java/google/registry/model/registry/label/PremiumList.java index 50f45daa1..2dc55b41f 100644 --- a/core/src/main/java/google/registry/model/registry/label/PremiumList.java +++ b/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 ofy().load().type(PremiumList.class).parent(getCrossTldKey()).id(name).now(), - () -> PremiumListDao.getLatestRevision(name).orElseThrow(NoSuchElementException::new)); + return tm().isOfy() + ? ofy().load().type(PremiumList.class).parent(getCrossTldKey()).id(name).now() + : PremiumListDao.getLatestRevision(name).orElseThrow(NoSuchElementException::new); } /** diff --git a/core/src/main/java/google/registry/model/registry/label/ReservedList.java b/core/src/main/java/google/registry/model/registry/label/ReservedList.java index c30caf3f8..27931fbae 100644 --- a/core/src/main/java/google/registry/model/registry/label/ReservedList.java +++ b/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() { @Override public ReservedList load(String listName) { - return ofyOrJpaTm( - () -> ReservedListDualWriteDao.getLatestRevision(listName).orElse(null), - () -> ReservedListSqlDao.getLatestRevision(listName).orElse(null)); + return tm().isOfy() + ? ReservedListDualWriteDao.getLatestRevision(listName).orElse(null) + : ReservedListSqlDao.getLatestRevision(listName).orElse(null); } }); diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java index edbc1b0b5..2d9db46f9 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java +++ b/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; + } } diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerUtil.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerUtil.java index 1028505b3..fc36a688c 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerUtil.java +++ b/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 ofyOrJpaTm(Supplier ofySupplier, Supplier 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. diff --git a/core/src/test/java/google/registry/testing/AppEngineExtension.java b/core/src/test/java/google/registry/testing/AppEngineExtension.java index 54f90cee1..bc22ca122 100644 --- a/core/src/test/java/google/registry/testing/AppEngineExtension.java +++ b/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 (withDatastore && !withoutCannedData) { - loadInitialData(); - } - }, - () -> { - if (withCloudSql && !withJpaUnitTest && !withoutCannedData) { - loadInitialData(); - } - }); + if (tm().isOfy()) { + if (withDatastore && !withoutCannedData) { + loadInitialData(); + } + } else { + if (withCloudSql && !withJpaUnitTest && !withoutCannedData) { + loadInitialData(); + } + } } /** diff --git a/core/src/test/java/google/registry/testing/DatastoreHelper.java b/core/src/test/java/google/registry/testing/DatastoreHelper.java index 23083423f..5fe631191 100644 --- a/core/src/test/java/google/registry/testing/DatastoreHelper.java +++ b/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( - () -> persistResource(reservedList), - () -> { - tm().transact(() -> tm().insert(reservedList)); - return reservedList; - }); + return tm().isOfy() + ? persistResource(reservedList) + : tm().transact( + () -> { + tm().insert(reservedList); + return reservedList; + }); } /** @@ -364,15 +364,15 @@ public class DatastoreHelper { .build(); PremiumListRevision revision = PremiumListRevision.create(premiumList, entries.keySet()); - ofyOrJpaTm( - () -> { - tm().putAllWithoutBackup( - ImmutableList.of( - premiumList.asBuilder().setRevision(Key.create(revision)).build(), revision)); - tm().putAllWithoutBackup( - parentPremiumListEntriesOnRevision(entries.values(), Key.create(revision))); - }, - () -> tm().transact(() -> tm().insert(premiumList))); + 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)); + } // 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 void saveResource(R resource, boolean wantBackup) { - ofyOrJpaTm( - () -> { - Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup(); - saver.entity(resource); - if (resource instanceof EppResource) { - EppResource eppResource = (EppResource) resource; - persistEppResourceExtras( - eppResource, EppResourceIndex.create(Key.create(eppResource)), saver); - } - }, - () -> tm().put(resource)); + if (tm().isOfy()) { + Saver saver = wantBackup ? ofy().save() : ofy().saveWithoutBackup(); + saver.entity(resource); + if (resource instanceof EppResource) { + EppResource eppResource = (EppResource) resource; + persistEppResourceExtras( + eppResource, EppResourceIndex.create(Key.create(eppResource)), saver); + } + } else { + tm().put(resource); + } } private static 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( - () -> { - Saver saver = ofy().save(); - saver.entity(resource); - persistEppResourceExtras(resource, eppResourceIndex, saver); - }, - () -> tm().put(resource))); + () -> { + if (tm().isOfy()) { + Saver saver = ofy().save(); + saver.entity(resource); + persistEppResourceExtras(resource, eppResourceIndex, saver); + } else { + 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 getHistoryEntries(EppResource resource) { - return ofyOrJpaTm( - () -> - ofy() - .load() - .type(HistoryEntry.class) - .ancestor(resource) - .order("modificationTime") - .list(), - () -> - tm().transact( - () -> { - ImmutableList unsorted = null; - if (resource instanceof ContactBase) { - unsorted = tm().loadAll(ContactHistory.class); - } else if (resource instanceof HostBase) { - unsorted = tm().loadAll(HostHistory.class); - } else if (resource instanceof DomainContent) { - unsorted = tm().loadAll(DomainHistory.class); - } else { - fail("Expected an EppResource instance, but got " + resource.getClass()); - } - ImmutableList filtered = - unsorted.stream() - .filter( - historyEntry -> - historyEntry - .getParent() - .getName() - .equals(resource.getRepoId())) - .collect(toImmutableList()); - return ImmutableList.sortedCopyOf(DateTimeComparator.getInstance(), filtered); - })); + return tm().isOfy() + ? ofy().load().type(HistoryEntry.class).ancestor(resource).order("modificationTime").list() + : tm().transact( + () -> { + ImmutableList unsorted = null; + if (resource instanceof ContactBase) { + unsorted = tm().loadAll(ContactHistory.class); + } else if (resource instanceof HostBase) { + unsorted = tm().loadAll(HostHistory.class); + } else if (resource instanceof DomainContent) { + unsorted = tm().loadAll(DomainHistory.class); + } else { + fail("Expected an EppResource instance, but got " + resource.getClass()); + } + ImmutableList 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 cloneAndSetAutoTimestamps(final R resource) { return tm().transact( - () -> - ofyOrJpaTm( - () -> ofy().load().fromEntity(ofy().save().toEntity(resource)), - () -> { - tm().put(resource); - return tm().load(resource); - })); + tm().isOfy() + ? () -> 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}. */