Validate SQL with Datastore being Primary (#1436)

* Validate SQL with Datastore being primary

Validates the data asynchronously replicated from Datastore to SQL.
This is a short term tool optimized for the current production database.

Tested in production.

core/src/test/java/google/registry/beam/common/RegistryJpaWriteTest.java

@@ -55,7 +55,8 @@ class RegistryJpaWriteTest implements Serializable {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension final transient InjectExtension injectExtension = new InjectExtension();
 

core/src/test/java/google/registry/beam/comparedb/DatastoreSnapshotsTest.java

@@ -0,0 +1,93 @@
+// Copyright 2021 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.beam.comparedb;
+
+import com.google.common.collect.ImmutableSet;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.beam.initsql.DatastoreSetupHelper;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.replay.SqlEntity;
+import google.registry.persistence.transaction.JpaTestExtensions;
+import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectExtension;
+import java.io.Serializable;
+import java.nio.file.Path;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Unit tests for {@link DatastoreSnapshots}. */
+class DatastoreSnapshotsTest {
+  static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
+
+  private FakeClock fakeClock = new FakeClock(START_TIME);
+
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
+
+  @RegisterExtension final transient InjectExtension injectExtension = new InjectExtension();
+
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  @RegisterExtension
+  final transient JpaIntegrationTestExtension database =
+      new JpaTestExtensions.Builder().withClock(fakeClock).buildIntegrationTestExtension();
+
+  DatastoreSetupHelper setupHelper;
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
+    setupHelper = new DatastoreSetupHelper(tmpDir, fakeClock).initializeData();
+    testPipeline
+        .getCoderRegistry()
+        .registerCoderForClass(SqlEntity.class, SerializableCoder.of(Serializable.class));
+  }
+
+  @Test
+  void loadDatastoreSnapshotByKind() {
+    PCollectionTuple tuple =
+        DatastoreSnapshots.loadDatastoreSnapshotByKind(
+            testPipeline,
+            setupHelper.exportDir.getAbsolutePath(),
+            setupHelper.commitLogDir.getAbsolutePath(),
+            START_TIME,
+            fakeClock.nowUtc().plusMillis(1),
+            ImmutableSet.copyOf(DatastoreSetupHelper.ALL_KINDS));
+    PAssert.that(tuple.get(ValidateSqlUtils.createSqlEntityTupleTag(Registrar.class)))
+        .containsInAnyOrder(setupHelper.registrar1, setupHelper.registrar2);
+    PAssert.that(tuple.get(ValidateSqlUtils.createSqlEntityTupleTag(DomainHistory.class)))
+        .containsInAnyOrder(setupHelper.historyEntry);
+    testPipeline.run();
+  }
+}

core/src/test/java/google/registry/beam/comparedb/SqlSnapshotsTest.java

@@ -0,0 +1,100 @@
+// Copyright 2021 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.beam.comparedb;
+
+import static google.registry.beam.comparedb.ValidateSqlUtils.createSqlEntityTupleTag;
+import static org.joda.time.DateTimeZone.UTC;
+
+import com.google.common.collect.ImmutableSet;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.model.bulkquery.TestSetupHelper;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.host.HostResource;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.replay.SqlEntity;
+import google.registry.model.tld.Registry;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import java.io.Serializable;
+import java.util.Optional;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link SqlSnapshots}. */
+class SqlSnapshotsTest {
+
+  protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  private final TestSetupHelper setupHelper = new TestSetupHelper(fakeClock);
+
+  @BeforeEach
+  void setUp() {
+    testPipeline
+        .getCoderRegistry()
+        .registerCoderForClass(SqlEntity.class, SerializableCoder.of(Serializable.class));
+    setupHelper.initializeAllEntities();
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+  }
+
+  @AfterEach
+  void afterEach() {
+    setupHelper.tearDownBulkQueryJpaTm();
+  }
+
+  @Test
+  void loadCloudSqlSnapshotByType() {
+    PCollectionTuple sqlSnapshot =
+        SqlSnapshots.loadCloudSqlSnapshotByType(
+            testPipeline,
+            ImmutableSet.of(
+                Registry.class,
+                Registrar.class,
+                DomainBase.class,
+                DomainHistory.class,
+                ContactResource.class,
+                HostResource.class),
+            Optional.empty());
+    PAssert.that(sqlSnapshot.get(createSqlEntityTupleTag(Registry.class)))
+        .containsInAnyOrder(setupHelper.registry);
+    PAssert.that(sqlSnapshot.get(createSqlEntityTupleTag(DomainBase.class)))
+        .containsInAnyOrder(setupHelper.domain);
+    PAssert.that(sqlSnapshot.get(createSqlEntityTupleTag(DomainHistory.class)))
+        .containsInAnyOrder(setupHelper.domainHistory);
+    testPipeline.run();
+  }
+}

core/src/test/java/google/registry/beam/comparedb/ValidateSqlUtilsTest.java

@@ -0,0 +1,63 @@
+// Copyright 2021 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.beam.comparedb;
+
+import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.beam.comparedb.ValidateSqlUtils.getMedianIdForHistoryTable;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static org.joda.time.DateTimeZone.UTC;
+
+import com.google.common.truth.Truth;
+import google.registry.model.bulkquery.TestSetupHelper;
+import google.registry.model.domain.DomainHistory;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link ValidateSqlUtils}. */
+class ValidateSqlUtilsTest {
+
+  private final FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  private final TestSetupHelper setupHelper = new TestSetupHelper(fakeClock);
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+
+  @Test
+  void getMedianIdForHistoryTable_emptyTable() {
+    assertThat(getMedianIdForHistoryTable("DomainHistory")).isEmpty();
+  }
+
+  @Test
+  void getMedianIdForHistoryTable_oneRow() {
+    setupHelper.initializeAllEntities();
+    Truth.assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(DomainHistory.class))).hasSize(1);
+    assertThat(getMedianIdForHistoryTable("DomainHistory"))
+        .hasValue(setupHelper.domainHistory.getId());
+  }
+
+  @Test
+  void getMedianIdForHistoryTable_twoRows() {
+    setupHelper.initializeAllEntities();
+    setupHelper.applyChangeToDomainAndHistory();
+    Truth.assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(DomainHistory.class))).hasSize(2);
+    assertThat(getMedianIdForHistoryTable("DomainHistory"))
+        .hasValue(setupHelper.domainHistory.getId());
+  }
+}

core/src/test/java/google/registry/beam/initsql/DatastoreSetupHelper.java

@@ -0,0 +1,295 @@
+// Copyright 2021 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.beam.initsql;
+
+import static google.registry.model.common.Cursor.CursorType.BRDA;
+import static google.registry.model.common.Cursor.CursorType.RECURRING_BILLING;
+import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
+import static google.registry.testing.DatabaseHelper.newRegistry;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.DatabaseHelper.persistSimpleResource;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.flows.domain.DomainFlowUtils;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Flag;
+import google.registry.model.billing.BillingEvent.Reason;
+import google.registry.model.common.Cursor;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DesignatedContact;
+import google.registry.model.domain.DomainAuthInfo;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.domain.GracePeriod;
+import google.registry.model.domain.launch.LaunchNotice;
+import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
+import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostResource;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.RegistrarContact;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.tld.Registry;
+import google.registry.model.transfer.DomainTransferData;
+import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import java.io.File;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import org.joda.money.Money;
+
+/**
+ * Sets up a test scenario in Datastore.
+ *
+ * <p>The {@link #initializeData} populates Datastore with test data, including {@link DomainBase},
+ * {@link DomainHistory}, and commit logs. The up-to-date version of the relevant entities are saved
+ * in public instance variables (e.g., {@link #domain} for easy access.
+ */
+public class DatastoreSetupHelper {
+
+  /**
+   * All kinds of entities to be set up in the Datastore. Must contain all kinds known to {@link
+   * InitSqlPipeline}.
+   */
+  public static final ImmutableList<Class<?>> ALL_KINDS =
+      ImmutableList.of(
+          Registry.class,
+          Cursor.class,
+          Registrar.class,
+          ContactResource.class,
+          RegistrarContact.class,
+          DomainBase.class,
+          HostResource.class,
+          HistoryEntry.class,
+          AllocationToken.class,
+          BillingEvent.Recurring.class,
+          BillingEvent.OneTime.class,
+          BillingEvent.Cancellation.class,
+          PollMessage.class);
+
+  private final Path tmpDir;
+  private final FakeClock fakeClock;
+
+  public File exportRootDir;
+  public File exportDir;
+  public File commitLogDir;
+
+  public Registrar registrar1;
+  public Registrar registrar2;
+  public DomainBase domain;
+  public ContactResource contact1;
+  public ContactResource contact2;
+  public HostResource hostResource;
+
+  public DomainHistory historyEntry;
+
+  public Cursor globalCursor;
+  public Cursor tldCursor;
+
+  public DatastoreSetupHelper(Path tempDir, FakeClock fakeClock) {
+    this.tmpDir = tempDir;
+    this.fakeClock = fakeClock;
+  }
+
+  public DatastoreSetupHelper initializeData() throws Exception {
+    try (BackupTestStore store = new BackupTestStore(fakeClock)) {
+      exportRootDir = Files.createDirectory(tmpDir.resolve("exports")).toFile();
+
+      persistResource(newRegistry("com", "COM"));
+      registrar1 = persistResource(AppEngineExtension.makeRegistrar1());
+      registrar2 = persistResource(AppEngineExtension.makeRegistrar2());
+      Key<DomainBase> domainKey = Key.create(null, DomainBase.class, "4-COM");
+      hostResource =
+          persistResource(
+              new HostResource.Builder()
+                  .setHostName("ns1.example.com")
+                  .setSuperordinateDomain(VKey.from(domainKey))
+                  .setRepoId("1-COM")
+                  .setCreationRegistrarId(registrar1.getRegistrarId())
+                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
+                  .build());
+      contact1 =
+          persistResource(
+              new ContactResource.Builder()
+                  .setContactId("contact_id1")
+                  .setRepoId("2-COM")
+                  .setCreationRegistrarId(registrar1.getRegistrarId())
+                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
+                  .build());
+      contact2 =
+          persistResource(
+              new ContactResource.Builder()
+                  .setContactId("contact_id2")
+                  .setRepoId("3-COM")
+                  .setCreationRegistrarId(registrar1.getRegistrarId())
+                  .setPersistedCurrentSponsorRegistrarId(registrar1.getRegistrarId())
+                  .build());
+      persistSimpleResource(
+          new RegistrarContact.Builder()
+              .setParent(registrar1)
+              .setName("John Abused")
+              .setEmailAddress("johnabuse@example.com")
+              .setVisibleInWhoisAsAdmin(true)
+              .setVisibleInWhoisAsTech(false)
+              .setPhoneNumber("+1.2125551213")
+              .setFaxNumber("+1.2125551213")
+              .setTypes(ImmutableSet.of(RegistrarContact.Type.ABUSE, RegistrarContact.Type.ADMIN))
+              .build());
+      historyEntry =
+          persistResource(
+              new DomainHistory.Builder()
+                  .setDomainRepoId(domainKey.getName())
+                  .setModificationTime(fakeClock.nowUtc())
+                  .setRegistrarId(registrar1.getRegistrarId())
+                  .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                  .build());
+      persistResource(
+          new AllocationToken.Builder().setToken("abc123").setTokenType(SINGLE_USE).build());
+      Key<DomainHistory> historyEntryKey = Key.create(historyEntry);
+      BillingEvent.OneTime onetimeBillEvent =
+          new BillingEvent.OneTime.Builder()
+              .setId(1)
+              .setReason(Reason.RENEW)
+              .setTargetId("example.com")
+              .setRegistrarId("TheRegistrar")
+              .setCost(Money.parse("USD 44.00"))
+              .setPeriodYears(4)
+              .setEventTime(fakeClock.nowUtc())
+              .setBillingTime(fakeClock.nowUtc())
+              .setParent(historyEntryKey)
+              .build();
+      persistResource(onetimeBillEvent);
+      Key<BillingEvent.OneTime> oneTimeBillKey = Key.create(onetimeBillEvent);
+      BillingEvent.Recurring recurringBillEvent =
+          new BillingEvent.Recurring.Builder()
+              .setId(2)
+              .setReason(Reason.RENEW)
+              .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+              .setTargetId("example.com")
+              .setRegistrarId("TheRegistrar")
+              .setEventTime(fakeClock.nowUtc())
+              .setRecurrenceEndTime(END_OF_TIME)
+              .setParent(historyEntryKey)
+              .build();
+      persistResource(recurringBillEvent);
+      VKey<BillingEvent.Recurring> recurringBillKey = recurringBillEvent.createVKey();
+      PollMessage.Autorenew autorenewPollMessage =
+          new PollMessage.Autorenew.Builder()
+              .setId(3L)
+              .setTargetId("example.com")
+              .setRegistrarId("TheRegistrar")
+              .setEventTime(fakeClock.nowUtc())
+              .setMsg("Domain was auto-renewed.")
+              .setParent(historyEntry)
+              .build();
+      persistResource(autorenewPollMessage);
+      VKey<PollMessage.Autorenew> autorenewPollKey = autorenewPollMessage.createVKey();
+      PollMessage.OneTime oneTimePollMessage =
+          new PollMessage.OneTime.Builder()
+              .setId(1L)
+              .setParent(historyEntry)
+              .setEventTime(fakeClock.nowUtc())
+              .setRegistrarId("TheRegistrar")
+              .setMsg(DomainFlowUtils.COLLISION_MESSAGE)
+              .build();
+      persistResource(oneTimePollMessage);
+      VKey<PollMessage.OneTime> onetimePollKey = oneTimePollMessage.createVKey();
+      domain =
+          persistResource(
+              new DomainBase.Builder()
+                  .setDomainName("example.com")
+                  .setRepoId("4-COM")
+                  .setCreationRegistrarId(registrar1.getRegistrarId())
+                  .setLastEppUpdateTime(fakeClock.nowUtc())
+                  .setLastEppUpdateRegistrarId(registrar2.getRegistrarId())
+                  .setLastTransferTime(fakeClock.nowUtc())
+                  .setStatusValues(
+                      ImmutableSet.of(
+                          StatusValue.CLIENT_DELETE_PROHIBITED,
+                          StatusValue.SERVER_DELETE_PROHIBITED,
+                          StatusValue.SERVER_TRANSFER_PROHIBITED,
+                          StatusValue.SERVER_UPDATE_PROHIBITED,
+                          StatusValue.SERVER_RENEW_PROHIBITED,
+                          StatusValue.SERVER_HOLD))
+                  .setRegistrant(contact1.createVKey())
+                  .setContacts(
+                      ImmutableSet.of(
+                          DesignatedContact.create(
+                              DesignatedContact.Type.ADMIN, contact2.createVKey())))
+                  .setNameservers(ImmutableSet.of(hostResource.createVKey()))
+                  .setSubordinateHosts(ImmutableSet.of("ns1.example.com"))
+                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
+                  .setRegistrationExpirationTime(fakeClock.nowUtc().plusYears(1))
+                  .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("password")))
+                  .setDsData(
+                      ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {0, 1, 2})))
+                  .setLaunchNotice(
+                      LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
+                  .setTransferData(
+                      new DomainTransferData.Builder()
+                          .setGainingRegistrarId(registrar1.getRegistrarId())
+                          .setLosingRegistrarId(registrar2.getRegistrarId())
+                          .setPendingTransferExpirationTime(fakeClock.nowUtc())
+                          .setServerApproveEntities(
+                              ImmutableSet.of(
+                                  VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
+                          .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
+                          .setServerApproveAutorenewEvent(recurringBillKey)
+                          .setServerApproveAutorenewPollMessage(autorenewPollKey)
+                          .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
+                          .setTransferStatus(TransferStatus.SERVER_APPROVED)
+                          .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
+                          .build())
+                  .setDeletePollMessage(onetimePollKey)
+                  .setAutorenewBillingEvent(recurringBillKey)
+                  .setAutorenewPollMessage(autorenewPollKey)
+                  .setSmdId("smdid")
+                  .addGracePeriod(
+                      GracePeriod.create(
+                          GracePeriodStatus.ADD,
+                          "4-COM",
+                          fakeClock.nowUtc().plusDays(1),
+                          "TheRegistrar",
+                          null))
+                  .build());
+      persistResource(
+          new BillingEvent.Cancellation.Builder()
+              .setReason(Reason.RENEW)
+              .setTargetId(domain.getDomainName())
+              .setRegistrarId(domain.getCurrentSponsorRegistrarId())
+              .setEventTime(fakeClock.nowUtc())
+              .setBillingTime(fakeClock.nowUtc())
+              .setRecurringEventKey(recurringBillEvent.createVKey())
+              .setParent(historyEntryKey)
+              .build());
+      globalCursor = persistResource(Cursor.createGlobal(RECURRING_BILLING, fakeClock.nowUtc()));
+      tldCursor = persistResource(Cursor.create(BRDA, fakeClock.nowUtc(), Registry.get("com")));
+      exportDir = store.export(exportRootDir.getAbsolutePath(), ALL_KINDS, ImmutableSet.of());
+      commitLogDir = Files.createDirectory(tmpDir.resolve("commits")).toFile();
+      fakeClock.advanceOneMilli();
+    }
+    return this;
+  }
+}

core/src/test/java/google/registry/beam/initsql/ExportLoadingTransformsTest.java

@@ -79,7 +79,7 @@ class ExportLoadingTransformsTest implements Serializable {
   @RegisterExtension
   @Order(value = 1)
   final transient DatastoreEntityExtension datastoreEntityExtension =
-      new DatastoreEntityExtension();
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient TestPipelineExtension testPipeline =

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -17,59 +17,22 @@ package google.registry.beam.initsql;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
-import static google.registry.model.common.Cursor.CursorType.BRDA;
-import static google.registry.model.common.Cursor.CursorType.RECURRING_BILLING;
-import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.testing.DatabaseHelper.newRegistry;
-import static google.registry.testing.DatabaseHelper.persistResource;
-import static google.registry.testing.DatabaseHelper.persistSimpleResource;
-import static google.registry.util.DateTimeUtils.END_OF_TIME;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.beam.TestPipelineExtension;
-import google.registry.flows.domain.DomainFlowUtils;
-import google.registry.model.billing.BillingEvent;
-import google.registry.model.billing.BillingEvent.Flag;
-import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.common.Cursor;
 import google.registry.model.contact.ContactResource;
-import google.registry.model.domain.DesignatedContact;
-import google.registry.model.domain.DomainAuthInfo;
 import google.registry.model.domain.DomainBase;
-import google.registry.model.domain.DomainHistory;
-import google.registry.model.domain.GracePeriod;
-import google.registry.model.domain.launch.LaunchNotice;
-import google.registry.model.domain.rgp.GracePeriodStatus;
-import google.registry.model.domain.secdns.DelegationSignerData;
-import google.registry.model.domain.token.AllocationToken;
-import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
-import google.registry.model.eppcommon.StatusValue;
-import google.registry.model.eppcommon.Trid;
 import google.registry.model.host.HostResource;
 import google.registry.model.ofy.Ofy;
-import google.registry.model.poll.PollMessage;
 import google.registry.model.registrar.Registrar;
-import google.registry.model.registrar.RegistrarContact;
-import google.registry.model.reporting.HistoryEntry;
-import google.registry.model.tld.Registry;
-import google.registry.model.transfer.DomainTransferData;
-import google.registry.model.transfer.TransferStatus;
-import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
-import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
-import java.io.File;
-import java.nio.file.Files;
 import java.nio.file.Path;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
@@ -81,27 +44,7 @@ import org.junit.jupiter.api.io.TempDir;
 class InitSqlPipelineTest {
   private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
 
-  /**
-   * All kinds of entities to be set up in the Datastore. Must contain all kinds known to {@link
-   * InitSqlPipeline}.
-   */
-  private static final ImmutableList<Class<?>> ALL_KINDS =
-      ImmutableList.of(
-          Registry.class,
-          Cursor.class,
-          Registrar.class,
-          ContactResource.class,
-          RegistrarContact.class,
-          DomainBase.class,
-          HostResource.class,
-          HistoryEntry.class,
-          AllocationToken.class,
-          BillingEvent.Recurring.class,
-          BillingEvent.OneTime.class,
-          BillingEvent.Cancellation.class,
-          PollMessage.class);
-
-  private transient FakeClock fakeClock = new FakeClock(START_TIME);
+  private FakeClock fakeClock = new FakeClock(START_TIME);
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
@@ -122,201 +65,12 @@ class InitSqlPipelineTest {
   final transient JpaIntegrationTestExtension database =
       new JpaTestExtensions.Builder().withClock(fakeClock).buildIntegrationTestExtension();
 
-  private File exportRootDir;
-  private File exportDir;
-  private File commitLogDir;
-
-  private transient Registrar registrar1;
-  private transient Registrar registrar2;
-  private transient DomainBase domain;
-  private transient ContactResource contact1;
-  private transient ContactResource contact2;
-  private transient HostResource hostResource;
-
-  private transient DomainHistory historyEntry;
-
-  private transient Cursor globalCursor;
-  private transient Cursor tldCursor;
+  DatastoreSetupHelper setupHelper;
 
   @BeforeEach
   void beforeEach() throws Exception {
-    try (BackupTestStore store = new BackupTestStore(fakeClock)) {
-      injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
-      exportRootDir = Files.createDirectory(tmpDir.resolve("exports")).toFile();
-
-      persistResource(newRegistry("com", "COM"));
-      registrar1 = persistResource(AppEngineExtension.makeRegistrar1());
-      registrar2 = persistResource(AppEngineExtension.makeRegistrar2());
-      Key<DomainBase> domainKey = Key.create(null, DomainBase.class, "4-COM");
-      hostResource =
-          persistResource(
-              new HostResource.Builder()
-                  .setHostName("ns1.example.com")
-                  .setSuperordinateDomain(VKey.from(domainKey))
-                  .setRepoId("1-COM")
-                  .setCreationRegistrarId(registrar1.getRegistrarId())
-                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
-                  .build());
-      contact1 =
-          persistResource(
-              new ContactResource.Builder()
-                  .setContactId("contact_id1")
-                  .setRepoId("2-COM")
-                  .setCreationRegistrarId(registrar1.getRegistrarId())
-                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
-                  .build());
-      contact2 =
-          persistResource(
-              new ContactResource.Builder()
-                  .setContactId("contact_id2")
-                  .setRepoId("3-COM")
-                  .setCreationRegistrarId(registrar1.getRegistrarId())
-                  .setPersistedCurrentSponsorRegistrarId(registrar1.getRegistrarId())
-                  .build());
-      persistSimpleResource(
-          new RegistrarContact.Builder()
-              .setParent(registrar1)
-              .setName("John Abused")
-              .setEmailAddress("johnabuse@example.com")
-              .setVisibleInWhoisAsAdmin(true)
-              .setVisibleInWhoisAsTech(false)
-              .setPhoneNumber("+1.2125551213")
-              .setFaxNumber("+1.2125551213")
-              .setTypes(ImmutableSet.of(RegistrarContact.Type.ABUSE, RegistrarContact.Type.ADMIN))
-              .build());
-      historyEntry =
-          persistResource(
-              new DomainHistory.Builder()
-                  .setDomainRepoId(domainKey.getName())
-                  .setModificationTime(fakeClock.nowUtc())
-                  .setRegistrarId(registrar1.getRegistrarId())
-                  .setType(HistoryEntry.Type.DOMAIN_CREATE)
-                  .build());
-      persistResource(
-          new AllocationToken.Builder().setToken("abc123").setTokenType(SINGLE_USE).build());
-      Key<DomainHistory> historyEntryKey = Key.create(historyEntry);
-      BillingEvent.OneTime onetimeBillEvent =
-          new BillingEvent.OneTime.Builder()
-              .setId(1)
-              .setReason(Reason.RENEW)
-              .setTargetId("example.com")
-              .setRegistrarId("TheRegistrar")
-              .setCost(Money.parse("USD 44.00"))
-              .setPeriodYears(4)
-              .setEventTime(fakeClock.nowUtc())
-              .setBillingTime(fakeClock.nowUtc())
-              .setParent(historyEntryKey)
-              .build();
-      persistResource(onetimeBillEvent);
-      Key<BillingEvent.OneTime> oneTimeBillKey = Key.create(onetimeBillEvent);
-      BillingEvent.Recurring recurringBillEvent =
-          new BillingEvent.Recurring.Builder()
-              .setId(2)
-              .setReason(Reason.RENEW)
-              .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-              .setTargetId("example.com")
-              .setRegistrarId("TheRegistrar")
-              .setEventTime(fakeClock.nowUtc())
-              .setRecurrenceEndTime(END_OF_TIME)
-              .setParent(historyEntryKey)
-              .build();
-      persistResource(recurringBillEvent);
-      VKey<BillingEvent.Recurring> recurringBillKey = recurringBillEvent.createVKey();
-      PollMessage.Autorenew autorenewPollMessage =
-          new PollMessage.Autorenew.Builder()
-              .setId(3L)
-              .setTargetId("example.com")
-              .setRegistrarId("TheRegistrar")
-              .setEventTime(fakeClock.nowUtc())
-              .setMsg("Domain was auto-renewed.")
-              .setParent(historyEntry)
-              .build();
-      persistResource(autorenewPollMessage);
-      VKey<PollMessage.Autorenew> autorenewPollKey = autorenewPollMessage.createVKey();
-      PollMessage.OneTime oneTimePollMessage =
-          new PollMessage.OneTime.Builder()
-              .setId(1L)
-              .setParent(historyEntry)
-              .setEventTime(fakeClock.nowUtc())
-              .setRegistrarId("TheRegistrar")
-              .setMsg(DomainFlowUtils.COLLISION_MESSAGE)
-              .build();
-      persistResource(oneTimePollMessage);
-      VKey<PollMessage.OneTime> onetimePollKey = oneTimePollMessage.createVKey();
-      domain =
-          persistResource(
-              new DomainBase.Builder()
-                  .setDomainName("example.com")
-                  .setRepoId("4-COM")
-                  .setCreationRegistrarId(registrar1.getRegistrarId())
-                  .setLastEppUpdateTime(fakeClock.nowUtc())
-                  .setLastEppUpdateRegistrarId(registrar2.getRegistrarId())
-                  .setLastTransferTime(fakeClock.nowUtc())
-                  .setStatusValues(
-                      ImmutableSet.of(
-                          StatusValue.CLIENT_DELETE_PROHIBITED,
-                          StatusValue.SERVER_DELETE_PROHIBITED,
-                          StatusValue.SERVER_TRANSFER_PROHIBITED,
-                          StatusValue.SERVER_UPDATE_PROHIBITED,
-                          StatusValue.SERVER_RENEW_PROHIBITED,
-                          StatusValue.SERVER_HOLD))
-                  .setRegistrant(contact1.createVKey())
-                  .setContacts(
-                      ImmutableSet.of(
-                          DesignatedContact.create(
-                              DesignatedContact.Type.ADMIN, contact2.createVKey())))
-                  .setNameservers(ImmutableSet.of(hostResource.createVKey()))
-                  .setSubordinateHosts(ImmutableSet.of("ns1.example.com"))
-                  .setPersistedCurrentSponsorRegistrarId(registrar2.getRegistrarId())
-                  .setRegistrationExpirationTime(fakeClock.nowUtc().plusYears(1))
-                  .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("password")))
-                  .setDsData(
-                      ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {0, 1, 2})))
-                  .setLaunchNotice(
-                      LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
-                  .setTransferData(
-                      new DomainTransferData.Builder()
-                          .setGainingRegistrarId(registrar1.getRegistrarId())
-                          .setLosingRegistrarId(registrar2.getRegistrarId())
-                          .setPendingTransferExpirationTime(fakeClock.nowUtc())
-                          .setServerApproveEntities(
-                              ImmutableSet.of(
-                                  VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
-                          .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
-                          .setServerApproveAutorenewEvent(recurringBillKey)
-                          .setServerApproveAutorenewPollMessage(autorenewPollKey)
-                          .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
-                          .setTransferStatus(TransferStatus.SERVER_APPROVED)
-                          .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
-                          .build())
-                  .setDeletePollMessage(onetimePollKey)
-                  .setAutorenewBillingEvent(recurringBillKey)
-                  .setAutorenewPollMessage(autorenewPollKey)
-                  .setSmdId("smdid")
-                  .addGracePeriod(
-                      GracePeriod.create(
-                          GracePeriodStatus.ADD,
-                          "4-COM",
-                          fakeClock.nowUtc().plusDays(1),
-                          "TheRegistrar",
-                          null))
-                  .build());
-      persistResource(
-          new BillingEvent.Cancellation.Builder()
-              .setReason(Reason.RENEW)
-              .setTargetId(domain.getDomainName())
-              .setRegistrarId(domain.getCurrentSponsorRegistrarId())
-              .setEventTime(fakeClock.nowUtc())
-              .setBillingTime(fakeClock.nowUtc())
-              .setRecurringEventKey(recurringBillEvent.createVKey())
-              .setParent(historyEntryKey)
-              .build());
-      globalCursor = persistResource(Cursor.createGlobal(RECURRING_BILLING, fakeClock.nowUtc()));
-      tldCursor = persistResource(Cursor.create(BRDA, fakeClock.nowUtc(), Registry.get("com")));
-      exportDir = store.export(exportRootDir.getAbsolutePath(), ALL_KINDS, ImmutableSet.of());
-      commitLogDir = Files.createDirectory(tmpDir.resolve("commits")).toFile();
-      fakeClock.advanceOneMilli();
-    }
+    injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
+    setupHelper = new DatastoreSetupHelper(tmpDir, fakeClock).initializeData();
   }
 
   @Test
@@ -325,24 +79,27 @@ class InitSqlPipelineTest {
         PipelineOptionsFactory.fromArgs(
                 "--commitLogStartTimestamp=" + START_TIME,
                 "--commitLogEndTimestamp=" + fakeClock.nowUtc().plusMillis(1),
-                "--datastoreExportDir=" + exportDir.getAbsolutePath(),
-                "--commitLogDir=" + commitLogDir.getAbsolutePath())
+                "--datastoreExportDir=" + setupHelper.exportDir.getAbsolutePath(),
+                "--commitLogDir=" + setupHelper.commitLogDir.getAbsolutePath())
             .withValidation()
             .as(InitSqlPipelineOptions.class);
     InitSqlPipeline initSqlPipeline = new InitSqlPipeline(options);
     initSqlPipeline.run(testPipeline).waitUntilFinish();
     assertHostResourceEquals(
-        jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
+        jpaTm().transact(() -> jpaTm().loadByKey(setupHelper.hostResource.createVKey())),
+        setupHelper.hostResource);
     assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
         .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
-        .containsExactly(registrar1, registrar2);
+        .containsExactly(setupHelper.registrar1, setupHelper.registrar2);
     assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
         .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
-        .containsExactly(contact1, contact2);
-    assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
+        .containsExactly(setupHelper.contact1, setupHelper.contact2);
+    assertDomainEquals(
+        jpaTm().transact(() -> jpaTm().loadByKey(setupHelper.domain.createVKey())),
+        setupHelper.domain);
     assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Cursor.class)))
         .comparingElementsUsing(immutableObjectCorrespondence())
-        .containsExactly(globalCursor, tldCursor);
+        .containsExactly(setupHelper.globalCursor, setupHelper.tldCursor);
   }
 
   private static void assertHostResourceEquals(HostResource actual, HostResource expected) {

core/src/test/java/google/registry/model/bulkquery/TestSetupHelper.java

@@ -50,8 +50,18 @@ import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
 
-/** Entity creation utilities for domain-related tests. */
-class TestSetupHelper {
+/**
+ * Utilities for managing domain-related SQL test scenarios.
+ *
+ * <p>The {@link #initializeAllEntities} method initializes the database, and the {@link
+ * #applyChangeToDomainAndHistory} makes one change to the domain, generating an additional history
+ * event. The most up-to-date values of the relevant entities are saved in public instance
+ * variables, {@link #domain} etc.
+ *
+ * <p>This class makes use of {@link DatabaseHelper}, which requires Datastore. Tests that use this
+ * class should use {@link AppEngineExtension}.
+ */
+public final class TestSetupHelper {
 
   public static final String TLD = "tld";
   public static final String DOMAIN_REPO_ID = "4-TLD";
@@ -60,21 +70,21 @@ class TestSetupHelper {
 
   private final FakeClock fakeClock;
 
-  Registry registry;
-  Registrar registrar;
-  ContactResource contact;
-  DomainBase domain;
-  DomainHistory domainHistory;
-  HostResource host;
+  public Registry registry;
+  public Registrar registrar;
+  public ContactResource contact;
+  public DomainBase domain;
+  public DomainHistory domainHistory;
+  public HostResource host;
 
   private JpaTransactionManager originalJpaTm;
   private JpaTransactionManager bulkQueryJpaTm;
 
-  TestSetupHelper(FakeClock fakeClock) {
+  public TestSetupHelper(FakeClock fakeClock) {
     this.fakeClock = fakeClock;
   }
 
-  void initializeAllEntities() {
+  public void initializeAllEntities() {
     registry = putInDb(DatabaseHelper.newRegistry(TLD, Ascii.toUpperCase(TLD)));
     registrar = saveRegistrar(REGISTRAR_ID);
     contact = putInDb(createContact(DOMAIN_REPO_ID, REGISTRAR_ID));
@@ -83,12 +93,12 @@ class TestSetupHelper {
     host = putInDb(createHost());
   }
 
-  void applyChangeToDomainAndHistory() {
+  public void applyChangeToDomainAndHistory() {
     domain = putInDb(createFullDomain(contact, host, fakeClock));
     domainHistory = putInDb(createFullHistory(domain, fakeClock));
   }
 
-  void setupBulkQueryJpaTm(AppEngineExtension appEngineExtension) {
+  public void setupBulkQueryJpaTm(AppEngineExtension appEngineExtension) {
     bulkQueryJpaTm =
         BulkQueryJpaFactory.createBulkQueryJpaTransactionManager(
             appEngineExtension
@@ -101,7 +111,7 @@ class TestSetupHelper {
     TransactionManagerFactory.setJpaTm(() -> bulkQueryJpaTm);
   }
 
-  void tearDownBulkQueryJpaTm() {
+  public void tearDownBulkQueryJpaTm() {
     if (bulkQueryJpaTm != null) {
       bulkQueryJpaTm.teardown();
       TransactionManagerFactory.setJpaTm(() -> originalJpaTm);