From bb5d2dcf0af1e5efbcee05b70b14e2cb854850a4 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Wed, 14 Jul 2021 13:05:01 -0400 Subject: [PATCH] Use the DatabaseMigrationSchedule to determine which TM to use (#1233) * Use the DatabaseMigrationSchedule to determine which TM to use We still allow the "manual" specification of a particular transaction manager, most useful in @DualDatabaseTest classes. If that isn't specified, we examine the migration schedule to see which to return. Notes: - This requires that any test that sets the migration schedule clean up after itself so that it won't affect future test runs of other classes (because the migration schedule cache is static) - One alternative would, instead of having a "test override" for the transaction manager, be to examine the registry environment and only override the transaction manager in the UNIT_TEST environment. This doesn't work because there are many instances in which tests simulate non-test environment. --- .../TransactionManagerFactory.java | 37 ++++++++++++++++--- .../ReplayCommitLogsToSqlActionTest.java | 6 +++ .../registry/beam/rde/RdePipelineTest.java | 6 +-- .../beam/spec11/Spec11PipelineTest.java | 6 +-- .../DatabaseMigrationStateScheduleTest.java | 17 +++++++++ .../ReplicateToDatastoreActionTest.java | 7 ++++ ...DatabaseTestInvocationContextProvider.java | 6 +-- .../GetDatabaseMigrationStateCommandTest.java | 6 +-- .../SetDatabaseMigrationStateCommandTest.java | 6 +++ 9 files changed, 78 insertions(+), 19 deletions(-) diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java index 8ed871803..bdeee15e1 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java +++ b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java @@ -16,17 +16,22 @@ package google.registry.persistence.transaction; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import static org.joda.time.DateTimeZone.UTC; import com.google.appengine.api.utils.SystemProperty; import com.google.appengine.api.utils.SystemProperty.Environment.Value; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Suppliers; import google.registry.config.RegistryEnvironment; +import google.registry.model.common.DatabaseMigrationStateSchedule; +import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase; import google.registry.model.ofy.DatastoreTransactionManager; import google.registry.persistence.DaggerPersistenceComponent; import google.registry.tools.RegistryToolEnvironment; import google.registry.util.NonFinalForTesting; +import java.util.Optional; import java.util.function.Supplier; +import org.joda.time.DateTime; /** Factory class to create {@link TransactionManager} instance. */ // TODO: Rename this to PersistenceFactory and move to persistence package. @@ -34,7 +39,8 @@ public class TransactionManagerFactory { private static final DatastoreTransactionManager ofyTm = createTransactionManager(); - @NonFinalForTesting private static TransactionManager tm = ofyTm; + /** Optional override to manually set the transaction manager per-test. */ + private static Optional tmForTest = Optional.empty(); /** Supplier for jpaTm so that it is initialized only once, upon first usage. */ @NonFinalForTesting @@ -70,9 +76,19 @@ public class TransactionManagerFactory { return SystemProperty.environment.value() == Value.Production; } - /** Returns {@link TransactionManager} instance. */ + /** + * Returns the {@link TransactionManager} instance. + * + *

Returns the {@link JpaTransactionManager} or {@link DatastoreTransactionManager} based on + * the migration schedule or the manually specified per-test transaction manager. + */ public static TransactionManager tm() { - return tm; + if (tmForTest.isPresent()) { + return tmForTest.get(); + } + PrimaryDatabase primaryDatabase = + DatabaseMigrationStateSchedule.getValueAtTime(DateTime.now(UTC)).getPrimaryDatabase(); + return primaryDatabase.equals(PrimaryDatabase.DATASTORE) ? ofyTm() : jpaTm(); } /** @@ -111,9 +127,18 @@ public class TransactionManagerFactory { jpaTm = Suppliers.memoize(jpaTmSupplier::get); } - /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */ + /** + * Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. + * + *

Used when overriding the per-test transaction manager for dual-database tests. + */ @VisibleForTesting - public static void setTm(TransactionManager newTm) { - tm = newTm; + public static void setTmForTest(TransactionManager newTm) { + tmForTest = Optional.of(newTm); + } + + /** Resets the overridden transaction manager post-test. */ + public static void removeTmOverrideForTest() { + tmForTest = Optional.empty(); } } diff --git a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java index 40c7a2db1..6ba146ba9 100644 --- a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java +++ b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java @@ -74,6 +74,7 @@ import google.registry.util.RequestStatusChecker; import java.io.IOException; import org.joda.time.DateTime; import org.joda.time.Duration; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -144,6 +145,11 @@ public class ReplayCommitLogsToSqlActionTest { TestObject.beforeSqlDeleteCallCount = 0; } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @Test void testReplay_multipleDiffFiles() throws Exception { jpaTm() diff --git a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java index 2610ef685..eba0ed1c6 100644 --- a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java +++ b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java @@ -23,7 +23,7 @@ import static google.registry.model.common.Cursor.CursorType.RDE_STAGING; import static google.registry.model.rde.RdeMode.FULL; import static google.registry.model.rde.RdeMode.THIN; import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm; -import static google.registry.persistence.transaction.TransactionManagerFactory.setTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.rde.RdeResourceType.CONTACT; import static google.registry.rde.RdeResourceType.DOMAIN; @@ -122,7 +122,7 @@ public class RdePipelineTest { @BeforeEach void beforeEach() throws Exception { originalTm = tm(); - setTm(jpaTm()); + setTmForTest(jpaTm()); loadInitialData(); // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The @@ -169,7 +169,7 @@ public class RdePipelineTest { @AfterEach void afterEach() { - setTm(originalTm); + setTmForTest(originalTm); } @Test diff --git a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java index 08f015649..1a9f92a61 100644 --- a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java +++ b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java @@ -18,7 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence; import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm; -import static google.registry.persistence.transaction.TransactionManagerFactory.setTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.testing.AppEngineExtension.makeRegistrar1; import static google.registry.testing.DatabaseHelper.createTld; @@ -231,7 +231,7 @@ class Spec11PipelineTest { private void setupCloudSql() { TransactionManager originalTm = tm(); - setTm(jpaTm()); + setTmForTest(jpaTm()); persistNewRegistrar("TheRegistrar"); persistNewRegistrar("NewRegistrar"); Registrar registrar1 = @@ -271,7 +271,7 @@ class Spec11PipelineTest { persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2)); persistResource( createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3)); - setTm(originalTm); + setTmForTest(originalTm); } private void verifySaveToGcs() throws Exception { diff --git a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java index 34b00a0e2..9906f21ea 100644 --- a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java +++ b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java @@ -15,6 +15,7 @@ package google.registry.model.common; import static com.google.common.truth.Truth.assertThat; +import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_ONLY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY_READ_ONLY; @@ -22,6 +23,7 @@ import static google.registry.model.common.DatabaseMigrationStateSchedule.Migrat import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY_READ_ONLY; import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.util.DateTimeUtils.START_OF_TIME; import static org.junit.Assert.assertThrows; @@ -30,6 +32,7 @@ import google.registry.model.EntityTestCase; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; import org.joda.time.DateTime; import org.joda.time.Duration; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -40,6 +43,11 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase { fakeClock.setAutoIncrementByOneMilli(); } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @Test void testEmpty_returnsDatastoreOnlyMap() { assertThat(DatabaseMigrationStateSchedule.getUncached()) @@ -134,6 +142,15 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase { .isEqualTo("Must be called in a transaction"); } + @Test + void testSuccess_factoryUsesSchedule() { + assertThat(tm().isOfy()).isTrue(); + // set the schedule to have converted to SQL_PRIMARY in the past + fakeClock.setTo(START_OF_TIME.plusDays(1)); + runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY); + assertThat(tm().isOfy()).isFalse(); + } + private void runValidTransition(MigrationState from, MigrationState to) { ImmutableSortedMap transitions = createMapEndingWithTransition(from, to); diff --git a/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java b/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java index e24e5cee8..6adf86b09 100644 --- a/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java +++ b/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java @@ -96,6 +96,13 @@ public class ReplicateToDatastoreActionTest { public void tearDown() { fakeClock.disableAutoIncrement(); RegistryConfig.overrideCloudSqlReplicateTransactions(false); + ofyTm() + .transact( + () -> + ofyTm() + .loadSingleton(DatabaseMigrationStateSchedule.class) + .ifPresent(ofyTm()::delete)); + DatabaseMigrationStateSchedule.CACHE.invalidateAll(); } @RetryingTest(4) diff --git a/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java b/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java index 05a3b4296..2c3413f85 100644 --- a/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java +++ b/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java @@ -154,15 +154,13 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm()); DatabaseType databaseType = (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY); - TransactionManagerFactory.setTm(databaseType.getTm()); + TransactionManagerFactory.setTmForTest(databaseType.getTm()); } } static void restoreTmAfterDualDatabaseTest(ExtensionContext context) { if (isDualDatabaseTest(context)) { - TransactionManager original = - (TransactionManager) context.getStore(NAMESPACE).get(ORIGINAL_TM_KEY); - TransactionManagerFactory.setTm(original); + TransactionManagerFactory.removeTmOverrideForTest(); } } diff --git a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java index 37a479c1e..5698c110c 100644 --- a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java @@ -24,15 +24,15 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat import google.registry.testing.DualDatabaseTest; import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; -import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.AfterEach; /** Tests for {@link GetDatabaseMigrationStateCommand}. */ @DualDatabaseTest public class GetDatabaseMigrationStateCommandTest extends CommandTestCase { - @BeforeEach - void beforeEach() { + @AfterEach + void afterEach() { ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); } diff --git a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java index 2be0c8890..5781e8998 100644 --- a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java @@ -28,6 +28,7 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat import google.registry.testing.DualDatabaseTest; import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; /** Tests for {@link SetDatabaseMigrationStateCommand}. */ @@ -47,6 +48,11 @@ public class SetDatabaseMigrationStateCommandTest DatabaseMigrationStateSchedule.CACHE.invalidateAll(); } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @TestOfyAndSql void testSuccess_setsBasicSchedule() throws Exception { assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);