Alt entity model for fast JPA bulk query (#1398)

* Alt entity model for fast JPA bulk query

Defined an alternative JPA entity model that allows fast bulk loading of
multi-level entities, DomainBase and DomainHistory. The idea is to bulk
the base table as well as the child tables separately, and assemble them
into the target entity in memory in a pipeline.

For DomainBase:

- Defined a DomainBaseLite class that models the "Domain" table only.

- Defined a DomainHost class that models the "DomainHost" table
  (nsHosts field).

- Exposed ID fields in GracePeriod so that they can be mapped to domains
  after being loaded into memory.

For DomainHistory:

- Defined a DomainHistoryLite class that models the "DomainHistory"
  table only.

- Defined a DomainHistoryHost class that models its namesake table.

- Exposed ID fields in GracePeriodHistory and DomainDsDataHistory
  classes so that they can be mapped to DomainHistory after being
  loaded into memory.

In PersistenceModule, provisioned a JpaTransactionManager that uses
the alternative entity model.

Also added a pipeline option that specifies which JpaTransactionManager
to use in a pipeline.

core/src/main/java/google/registry/beam/common/RegistryPipelineComponent.java

@@ -21,6 +21,7 @@ import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.persistence.PersistenceModule;
+import google.registry.persistence.PersistenceModule.BeamBulkQueryJpaTm;
 import google.registry.persistence.PersistenceModule.BeamJpaTm;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.persistence.transaction.JpaTransactionManager;
@@ -45,9 +46,19 @@ public interface RegistryPipelineComponent {
   @Config("projectId")
   String getProjectId();
 
+  /** Returns the regular {@link JpaTransactionManager} for general use. */
   @BeamJpaTm
   Lazy<JpaTransactionManager> getJpaTransactionManager();
 
+  /**
+   * Returns a {@link JpaTransactionManager} optimized for bulk loading multi-level JPA entities
+   * ({@link google.registry.model.domain.DomainBase} and {@link
+   * google.registry.model.domain.DomainHistory}). Please refer to {@link
+   * google.registry.model.bulkquery.BulkQueryEntities} for more information.
+   */
+  @BeamBulkQueryJpaTm
+  Lazy<JpaTransactionManager> getBulkQueryJpaTransactionManager();
+
   @Component.Builder
   interface Builder {
 

core/src/main/java/google/registry/beam/common/RegistryPipelineOptions.java

@@ -16,6 +16,7 @@ package google.registry.beam.common;
 
 import google.registry.beam.common.RegistryJpaIO.Write;
 import google.registry.config.RegistryEnvironment;
+import google.registry.persistence.PersistenceModule.JpaTransactionManagerType;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -44,6 +45,12 @@ public interface RegistryPipelineOptions extends GcpOptions {
 
   void setIsolationOverride(TransactionIsolationLevel isolationOverride);
 
+  @Description("The JPA Transaction Manager to use.")
+  @Default.Enum(value = "REGULAR")
+  JpaTransactionManagerType getJpaTransactionManagerType();
+
+  void setJpaTransactionManagerType(JpaTransactionManagerType jpaTransactionManagerType);
+
   @Description("The number of entities to write to the SQL database in one operation.")
   @Default.Integer(20)
   int getSqlWriteBatchSize();

core/src/main/java/google/registry/beam/common/RegistryPipelineWorkerInitializer.java

@@ -49,10 +49,19 @@ public class RegistryPipelineWorkerInitializer implements JvmInitializer {
     }
     logger.atInfo().log("Setting up RegistryEnvironment %s.", environment);
     environment.setup();
-    Lazy<JpaTransactionManager> transactionManagerLazy =
-        toRegistryPipelineComponent(registryOptions).getJpaTransactionManager();
+    RegistryPipelineComponent registryPipelineComponent =
+        toRegistryPipelineComponent(registryOptions);
+    Lazy<JpaTransactionManager> transactionManagerLazy;
+    switch (registryOptions.getJpaTransactionManagerType()) {
+      case BULK_QUERY:
+        transactionManagerLazy = registryPipelineComponent.getBulkQueryJpaTransactionManager();
+        break;
+      case REGULAR:
+      default:
+        transactionManagerLazy = registryPipelineComponent.getJpaTransactionManager();
+    }
     TransactionManagerFactory.setJpaTmOnBeamWorker(transactionManagerLazy::get);
-    // Masquarade all threads as App Engine threads so we can create Ofy keys in the pipeline. Also
+    // Masquerade all threads as App Engine threads so we can create Ofy keys in the pipeline. Also
     // loads all ofy entities.
     new AppEngineEnvironment("Beam").setEnvironmentForAllThreads();
     // Set the system property so that we can call IdService.allocateId() without access to

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

@@ -14,6 +14,7 @@
 
 package google.registry.model;
 
+import google.registry.util.PreconditionsUtils;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.AttributeOverride;
@@ -49,4 +50,14 @@ public abstract class BackupGroupRoot extends ImmutableObject {
   public UpdateAutoTimestamp getUpdateTimestamp() {
     return updateTimestamp;
   }
+
+  /**
+   * Copies {@link #updateTimestamp} from another entity.
+   *
+   * <p>This method is for the few cases when {@code updateTimestamp} is copied between different
+   * types of entities. Use {@link #clone} for same-type copying.
+   */
+  protected void copyUpdateTimestamp(BackupGroupRoot other) {
+    this.updateTimestamp = PreconditionsUtils.checkArgumentNotNull(other, "other").updateTimestamp;
+  }
 }

core/src/main/java/google/registry/model/bulkquery/BulkQueryEntities.java

@@ -0,0 +1,109 @@
+// 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.model.bulkquery;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.domain.GracePeriod;
+import google.registry.model.domain.GracePeriod.GracePeriodHistory;
+import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.domain.secdns.DomainDsDataHistory;
+import google.registry.model.host.HostResource;
+import google.registry.model.reporting.DomainTransactionRecord;
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTransactionManager;
+
+/**
+ * Utilities for managing an alternative JPA entity model optimized for bulk loading multi-level
+ * entities such as {@link DomainBase} and {@link DomainHistory}.
+ *
+ * <p>In a bulk query for a multi-level JPA entity type, the JPA framework only generates a bulk
+ * query (SELECT * FROM table) for the base table. Then, for each row in the base table, additional
+ * queries are issued to load associated rows in child tables. This can be very slow when an entity
+ * type has multiple child tables.
+ *
+ * <p>We have defined an alternative entity model for {@code DomainBase} and {@code DomainHistory},
+ * where the base table as well as the child tables are mapped to single-level entity types. The
+ * idea is to load each of these types using a bulk query, and assemble them into the target type in
+ * memory in a pipeline. The main use case is Datastore-Cloud SQL validation during the Registry
+ * database migration, where we will need the full database snapshots frequently.
+ */
+public class BulkQueryEntities {
+  /**
+   * The JPA entity classes in persistence.xml to replace when creating the {@link
+   * JpaTransactionManager} for bulk query.
+   */
+  public static final ImmutableMap<String, String> JPA_ENTITIES_REPLACEMENTS =
+      ImmutableMap.of(
+          DomainBase.class.getCanonicalName(),
+          DomainBaseLite.class.getCanonicalName(),
+          DomainHistory.class.getCanonicalName(),
+          DomainHistoryLite.class.getCanonicalName());
+
+  /* The JPA entity classes that are not included in persistence.xml and need to be added to
+   * the {@link JpaTransactionManager} for bulk query.*/
+  public static final ImmutableList<String> JPA_ENTITIES_NEW =
+      ImmutableList.of(
+          DomainHost.class.getCanonicalName(), DomainHistoryHost.class.getCanonicalName());
+
+  public static DomainBase assembleDomainBase(
+      DomainBaseLite domainBaseLite,
+      ImmutableSet<GracePeriod> gracePeriods,
+      ImmutableSet<DelegationSignerData> delegationSignerData,
+      ImmutableSet<VKey<HostResource>> nsHosts) {
+    DomainBase.Builder builder = new DomainBase.Builder();
+    builder.copyFrom(domainBaseLite);
+    builder.setGracePeriods(gracePeriods);
+    builder.setDsData(delegationSignerData);
+    builder.setNameservers(nsHosts);
+    return builder.build();
+  }
+
+  public static DomainHistory assembleDomainHistory(
+      DomainHistoryLite domainHistoryLite,
+      ImmutableSet<DomainDsDataHistory> dsDataHistories,
+      ImmutableSet<VKey<HostResource>> domainHistoryHosts,
+      ImmutableSet<GracePeriodHistory> gracePeriodHistories,
+      ImmutableSet<DomainTransactionRecord> transactionRecords) {
+    DomainHistory.Builder builder = new DomainHistory.Builder();
+    builder.copyFrom(domainHistoryLite);
+    DomainContent rawDomainContent = domainHistoryLite.domainContent;
+    if (rawDomainContent != null) {
+      DomainContent newDomainContent =
+          domainHistoryLite
+              .domainContent
+              .asBuilder()
+              .setNameservers(domainHistoryHosts)
+              .setGracePeriods(
+                  gracePeriodHistories.stream()
+                      .map(GracePeriod::createFromHistory)
+                      .collect(toImmutableSet()))
+              .setDsData(
+                  dsDataHistories.stream()
+                      .map(DelegationSignerData::create)
+                      .collect(toImmutableSet()))
+              .build();
+      builder.setDomain(newDomainContent);
+    }
+    return builder.buildAndAssemble(
+        dsDataHistories, domainHistoryHosts, gracePeriodHistories, transactionRecords);
+  }
+}

core/src/main/java/google/registry/model/bulkquery/DomainBaseLite.java

@@ -0,0 +1,49 @@
+// 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.model.bulkquery;
+
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
+import google.registry.model.replay.SqlOnlyEntity;
+import google.registry.persistence.VKey;
+import google.registry.persistence.WithStringVKey;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.Entity;
+
+/**
+ * A 'light' version of {@link DomainBase} with only base table ("Domain") attributes, which allows
+ * fast bulk loading. They are used in in-memory assembly of {@code DomainBase} instances along with
+ * bulk-loaded child entities ({@code GracePeriod} etc). The in-memory assembly achieves much higher
+ * performance than loading {@code DomainBase} directly.
+ *
+ * <p>Please refer to {@link BulkQueryEntities} for more information.
+ */
+@Entity(name = "Domain")
+@WithStringVKey
+@Access(AccessType.FIELD)
+public class DomainBaseLite extends DomainContent implements SqlOnlyEntity {
+
+  @Override
+  @javax.persistence.Id
+  @Access(AccessType.PROPERTY)
+  public String getRepoId() {
+    return super.getRepoId();
+  }
+
+  public static VKey<DomainBaseLite> createVKey(String repoId) {
+    return VKey.createSql(DomainBaseLite.class, repoId);
+  }
+}

core/src/main/java/google/registry/model/bulkquery/DomainHistoryHost.java

@@ -0,0 +1,50 @@
+// 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.model.bulkquery;
+
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
+import google.registry.model.host.HostResource;
+import google.registry.model.replay.SqlOnlyEntity;
+import google.registry.persistence.VKey;
+import java.io.Serializable;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+
+/**
+ * A name server host referenced by a {@link google.registry.model.domain.DomainHistory} record.
+ * Please refer to {@link BulkQueryEntities} for usage.
+ */
+@Entity
+@Access(AccessType.FIELD)
+@IdClass(DomainHistoryHost.class)
+public class DomainHistoryHost implements Serializable, SqlOnlyEntity {
+
+  @Id private Long domainHistoryHistoryRevisionId;
+  @Id private String domainHistoryDomainRepoId;
+  @Id private String hostRepoId;
+
+  private DomainHistoryHost() {}
+
+  public DomainHistoryId getDomainHistoryId() {
+    return new DomainHistoryId(domainHistoryDomainRepoId, domainHistoryHistoryRevisionId);
+  }
+
+  public VKey<HostResource> getHostVKey() {
+    return VKey.create(HostResource.class, hostRepoId);
+  }
+}

core/src/main/java/google/registry/model/bulkquery/DomainHistoryLite.java

@@ -0,0 +1,122 @@
+// 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.model.bulkquery;
+
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
+import google.registry.model.domain.Period;
+import google.registry.model.replay.SqlOnlyEntity;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import javax.annotation.Nullable;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.AttributeOverride;
+import javax.persistence.AttributeOverrides;
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+import javax.persistence.PostLoad;
+
+/**
+ * A 'light' version of {@link DomainHistory} with only base table ("DomainHistory") attributes,
+ * which allows fast bulk loading. They are used in in-memory assembly of {@code DomainHistory}
+ * instances along with bulk-loaded child entities ({@code GracePeriodHistory} etc). The in-memory
+ * assembly achieves much higher performance than loading {@code DomainHistory} directly.
+ *
+ * <p>Please refer to {@link BulkQueryEntities} for more information.
+ *
+ * <p>This class is adapted from {@link DomainHistory} by removing the {@code dsDataHistories},
+ * {@code gracePeriodHistories}, and {@code nsHosts} fields and associated methods.
+ */
+@Entity(name = "DomainHistory")
+@Access(AccessType.FIELD)
+@IdClass(DomainHistoryId.class)
+public class DomainHistoryLite extends HistoryEntry implements SqlOnlyEntity {
+
+  // Store DomainContent instead of DomainBase so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
+  @Nullable DomainContent domainContent;
+
+  @Id
+  @Access(AccessType.PROPERTY)
+  public String getDomainRepoId() {
+    // We need to handle null case here because Hibernate sometimes accesses this method before
+    // parent gets initialized
+    return parent == null ? null : parent.getName();
+  }
+
+  /** This method is private because it is only used by Hibernate. */
+  @SuppressWarnings("unused")
+  private void setDomainRepoId(String domainRepoId) {
+    parent = Key.create(DomainBase.class, domainRepoId);
+  }
+
+  @Override
+  @Nullable
+  @Access(AccessType.PROPERTY)
+  @AttributeOverrides({
+    @AttributeOverride(name = "unit", column = @Column(name = "historyPeriodUnit")),
+    @AttributeOverride(name = "value", column = @Column(name = "historyPeriodValue"))
+  })
+  public Period getPeriod() {
+    return super.getPeriod();
+  }
+
+  /**
+   * For transfers, the id of the other registrar.
+   *
+   * <p>For requests and cancels, the other registrar is the losing party (because the registrar
+   * sending the EPP transfer command is the gaining party). For approves and rejects, the other
+   * registrar is the gaining party.
+   */
+  @Nullable
+  @Access(AccessType.PROPERTY)
+  @Column(name = "historyOtherRegistrarId")
+  @Override
+  public String getOtherRegistrarId() {
+    return super.getOtherRegistrarId();
+  }
+
+  @Id
+  @Column(name = "historyRevisionId")
+  @Access(AccessType.PROPERTY)
+  @Override
+  public long getId() {
+    return super.getId();
+  }
+
+  /** The key to the {@link DomainBase} this is based off of. */
+  public VKey<DomainBase> getParentVKey() {
+    return VKey.create(DomainBase.class, getDomainRepoId());
+  }
+
+  @PostLoad
+  void postLoad() {
+    if (domainContent == null) {
+      return;
+    }
+    // See inline comments in DomainHistory.postLoad for reasons for the following lines.
+    if (domainContent.getDomainName() == null) {
+      domainContent = null;
+    } else if (domainContent.getRepoId() == null) {
+      domainContent.setRepoId(parent.getName());
+    }
+  }
+}

core/src/main/java/google/registry/model/bulkquery/DomainHost.java

@@ -0,0 +1,46 @@
+// 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.model.bulkquery;
+
+import google.registry.model.host.HostResource;
+import google.registry.model.replay.SqlOnlyEntity;
+import google.registry.persistence.VKey;
+import java.io.Serializable;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+
+/** A name server host of a domain. Please refer to {@link BulkQueryEntities} for usage. */
+@Entity
+@Access(AccessType.FIELD)
+@IdClass(DomainHost.class)
+public class DomainHost implements Serializable, SqlOnlyEntity {
+
+  @Id private String domainRepoId;
+
+  @Id private String hostRepoId;
+
+  DomainHost() {}
+
+  public String getDomainRepoId() {
+    return domainRepoId;
+  }
+
+  public VKey<HostResource> getHostVKey() {
+    return VKey.create(HostResource.class, hostRepoId);
+  }
+}

core/src/main/java/google/registry/model/domain/DomainBase.java

@@ -189,6 +189,7 @@ public class DomainBase extends DomainContent
     }
 
     public Builder copyFrom(DomainContent domainContent) {
+      this.getInstance().copyUpdateTimestamp(domainContent);
       return this.setAuthInfo(domainContent.getAuthInfo())
           .setAutorenewPollMessage(domainContent.getAutorenewPollMessage())
           .setAutorenewBillingEvent(domainContent.getAutorenewBillingEvent())

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -216,6 +216,10 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
     return super.getId();
   }
 
+  public DomainHistoryId getDomainHistoryId() {
+    return new DomainHistoryId(getDomainRepoId(), getId());
+  }
+
   /** Returns keys to the {@link HostResource} that are the nameservers for the domain. */
   public Set<VKey<HostResource>> getNsHosts() {
     return nsHosts;
@@ -314,6 +318,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
           nullToEmptyImmutableCopy(domainHistory.domainContent.getGracePeriods()).stream()
               .map(gracePeriod -> GracePeriodHistory.createFrom(domainHistory.id, gracePeriod))
               .collect(toImmutableSet());
+    } else {
+      domainHistory.nsHosts = ImmutableSet.of();
     }
   }
 
@@ -393,8 +399,16 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       if (domainContent == null) {
         return this;
       }
+      // TODO(b/203609982): if actual type of domainContent is DomainBase, convert to DomainContent
+      // Note: a DomainHistory fetched by JPA has DomainContent in this field. Allowing DomainBase
+      // in the setter makes equality checks messy.
       getInstance().domainContent = domainContent;
-      return super.setParent(domainContent);
+      if (domainContent instanceof DomainBase) {
+        super.setParent(domainContent);
+      } else {
+        super.setParent(Key.create(DomainBase.class, domainContent.getRepoId()));
+      }
+      return this;
     }
 
     public Builder setDomainRepoId(String domainRepoId) {
@@ -412,5 +426,19 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       fillAuxiliaryFieldsFromDomain(instance);
       return instance;
     }
+
+    public DomainHistory buildAndAssemble(
+        ImmutableSet<DomainDsDataHistory> dsDataHistories,
+        ImmutableSet<VKey<HostResource>> domainHistoryHosts,
+        ImmutableSet<GracePeriodHistory> gracePeriodHistories,
+        ImmutableSet<DomainTransactionRecord> transactionRecords) {
+      DomainHistory instance = super.build();
+      instance.dsDataHistories = dsDataHistories;
+      instance.nsHosts = domainHistoryHosts;
+      instance.gracePeriodHistories = gracePeriodHistories;
+      instance.domainTransactionRecords = transactionRecords;
+      instance.hashCode = null;
+      return instance;
+    }
   }
 }

core/src/main/java/google/registry/model/domain/GracePeriod.java

@@ -22,6 +22,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.googlecode.objectify.annotation.Embed;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.replay.DatastoreAndSqlEntity;
 import google.registry.model.replay.SqlOnlyEntity;
@@ -203,7 +204,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
   /** Entity class to represent a historic {@link GracePeriod}. */
   @Entity(name = "GracePeriodHistory")
   @Table(indexes = @Index(columnList = "domainRepoId"))
-  static class GracePeriodHistory extends GracePeriodBase implements SqlOnlyEntity {
+  public static class GracePeriodHistory extends GracePeriodBase implements SqlOnlyEntity {
     @Id Long gracePeriodHistoryRevisionId;
 
     /** ID for the associated {@link DomainHistory} entity. */
@@ -215,6 +216,10 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
       return super.getGracePeriodId();
     }
 
+    public DomainHistoryId getDomainHistoryId() {
+      return new DomainHistoryId(getDomainRepoId(), domainHistoryRevisionId);
+    }
+
     static GracePeriodHistory createFrom(long historyRevisionId, GracePeriod gracePeriod) {
       GracePeriodHistory instance = new GracePeriodHistory();
       instance.gracePeriodHistoryRevisionId = allocateId();

core/src/main/java/google/registry/model/domain/secdns/DomainDsDataHistory.java

@@ -17,6 +17,7 @@ package google.registry.model.domain.secdns;
 import static google.registry.model.IdService.allocateId;
 
 import google.registry.model.domain.DomainHistory;
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.replay.SqlOnlyEntity;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
@@ -53,6 +54,10 @@ public class DomainDsDataHistory extends DomainDsDataBase implements SqlOnlyEnti
     return instance;
   }
 
+  public DomainHistory.DomainHistoryId getDomainHistoryId() {
+    return new DomainHistoryId(getDomainRepoId(), domainHistoryRevisionId);
+  }
+
   @Override
   @Access(AccessType.PROPERTY)
   public String getDomainRepoId() {

core/src/main/java/google/registry/model/reporting/DomainTransactionRecord.java

@@ -22,6 +22,7 @@ import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.replay.DatastoreAndSqlEntity;
 import javax.persistence.Column;
 import javax.persistence.Entity;
@@ -49,7 +50,6 @@ public class DomainTransactionRecord extends ImmutableObject
 
   @Id
   @Ignore
-  @ImmutableObject.DoNotCompare
   @GeneratedValue(strategy = GenerationType.IDENTITY)
   @ImmutableObject.Insignificant
   Long id;
@@ -58,6 +58,14 @@ public class DomainTransactionRecord extends ImmutableObject
   @Column(nullable = false)
   String tld;
 
+  // The following two fields are exposed in this entity to support bulk-loading in Cloud SQL by the
+  // Datastore-SQL validation. They are excluded from equality check since they are not set in
+  // Datastore.
+  // TODO(b/203609782): post migration, decide whether to keep these two fields.
+  @Ignore @ImmutableObject.Insignificant String domainRepoId;
+
+  @Ignore @ImmutableObject.Insignificant Long historyRevisionId;
+
   /**
    * The time this Transaction takes effect (counting grace periods and other nuances).
    *
@@ -174,6 +182,10 @@ public class DomainTransactionRecord extends ImmutableObject
     }
   }
 
+  public DomainHistoryId getDomainHistoryId() {
+    return new DomainHistoryId(domainRepoId, historyRevisionId);
+  }
+
   public DateTime getReportingTime() {
     return reportingTime;
   }

core/src/main/java/google/registry/persistence/BulkQueryJpaFactory.java

@@ -0,0 +1,65 @@
+// 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.persistence;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
+import google.registry.model.bulkquery.BulkQueryEntities;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.persistence.transaction.JpaTransactionManagerImpl;
+import google.registry.util.Clock;
+import java.util.List;
+import javax.persistence.EntityManagerFactory;
+import org.hibernate.jpa.boot.internal.ParsedPersistenceXmlDescriptor;
+import org.hibernate.jpa.boot.spi.Bootstrap;
+
+/**
+ * Defines factory method for instantiating the bulk-query optimized {@link JpaTransactionManager}.
+ */
+public final class BulkQueryJpaFactory {
+
+  private BulkQueryJpaFactory() {}
+
+  static EntityManagerFactory createBulkQueryEntityManagerFactory(
+      ImmutableMap<String, String> cloudSqlConfigs) {
+    ParsedPersistenceXmlDescriptor descriptor =
+        PersistenceXmlUtility.getParsedPersistenceXmlDescriptor();
+
+    List<String> updatedManagedClasses =
+        Streams.concat(
+                descriptor.getManagedClassNames().stream(),
+                BulkQueryEntities.JPA_ENTITIES_NEW.stream())
+            .map(
+                name -> {
+                  if (BulkQueryEntities.JPA_ENTITIES_REPLACEMENTS.containsKey(name)) {
+                    return BulkQueryEntities.JPA_ENTITIES_REPLACEMENTS.get(name);
+                  }
+                  return name;
+                })
+            .collect(ImmutableList.toImmutableList());
+
+    descriptor.getManagedClassNames().clear();
+    descriptor.getManagedClassNames().addAll(updatedManagedClasses);
+
+    return Bootstrap.getEntityManagerFactoryBuilder(descriptor, cloudSqlConfigs).build();
+  }
+
+  public static JpaTransactionManager createBulkQueryJpaTransactionManager(
+      ImmutableMap<String, String> cloudSqlConfigs, Clock clock) {
+    return new JpaTransactionManagerImpl(
+        createBulkQueryEntityManagerFactory(cloudSqlConfigs), clock);
+  }
+}

core/src/main/java/google/registry/persistence/PersistenceModule.java

@@ -152,13 +152,36 @@ public abstract class PersistenceModule {
   @Singleton
   @BeamPipelineCloudSqlConfigs
   static ImmutableMap<String, String> provideBeamPipelineCloudSqlConfigs(
-      @Config("beamCloudSqlJdbcUrl") String jdbcUrl,
-      @Config("beamCloudSqlInstanceConnectionName") String instanceConnectionName,
-      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
+      SqlCredentialStore credentialStore,
+      @Config("instanceConnectionNameOverride")
+          Optional<Provider<String>> instanceConnectionNameOverride,
       @Config("beamIsolationOverride")
-          Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
-    return createPartialSqlConfigs(
-        jdbcUrl, instanceConnectionName, defaultConfigs, isolationOverride);
+          Optional<Provider<TransactionIsolationLevel>> isolationOverride,
+      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs) {
+    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
+    // TODO(b/175700623): make sql username configurable from config file.
+    SqlCredential credential = credentialStore.getCredential(new RobotUser(RobotId.NOMULUS));
+    overrides.put(Environment.USER, credential.login());
+    overrides.put(Environment.PASS, credential.password());
+    // Override the default minimum which is tuned for the Registry server. A worker VM should
+    // release all connections if it no longer interacts with the database.
+    overrides.put(HIKARI_MINIMUM_IDLE, "0");
+    /**
+     * Disable Hikari's maxPoolSize limit check by setting it to an absurdly large number. The
+     * effective (and desirable) limit is the number of pipeline threads on the pipeline worker,
+     * which can be configured using pipeline options. See {@link RegistryPipelineOptions} for more
+     * information.
+     */
+    overrides.put(HIKARI_MAXIMUM_POOL_SIZE, String.valueOf(Integer.MAX_VALUE));
+    instanceConnectionNameOverride
+        .map(Provider::get)
+        .ifPresent(
+            instanceConnectionName ->
+                overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, instanceConnectionName));
+    isolationOverride
+        .map(Provider::get)
+        .ifPresent(isolation -> overrides.put(Environment.ISOLATION, isolation.name()));
+    return ImmutableMap.copyOf(overrides);
   }
 
   @VisibleForTesting
@@ -230,37 +253,17 @@ public abstract class PersistenceModule {
   @Singleton
   @BeamJpaTm
   static JpaTransactionManager provideBeamJpaTm(
-      SqlCredentialStore credentialStore,
-      @Config("instanceConnectionNameOverride")
-          Optional<Provider<String>> instanceConnectionNameOverride,
-      @Config("beamIsolationOverride")
-          Optional<Provider<TransactionIsolationLevel>> isolationOverride,
-      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
-      Clock clock) {
-    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
-    // TODO(b/175700623): make sql username configurable from config file.
-    SqlCredential credential = credentialStore.getCredential(new RobotUser(RobotId.NOMULUS));
-    overrides.put(Environment.USER, credential.login());
-    overrides.put(Environment.PASS, credential.password());
-    // Override the default minimum which is tuned for the Registry server. A worker VM should
-    // release all connections if it no longer interacts with the database.
-    overrides.put(HIKARI_MINIMUM_IDLE, "0");
-    /**
-     * Disable Hikari's maxPoolSize limit check by setting it to an absurdly large number. The
-     * effective (and desirable) limit is the number of pipeline threads on the pipeline worker,
-     * which can be configured using pipeline options. See {@link RegistryPipelineOptions} for more
-     * information.
-     */
-    overrides.put(HIKARI_MAXIMUM_POOL_SIZE, String.valueOf(Integer.MAX_VALUE));
-    instanceConnectionNameOverride
-        .map(Provider::get)
-        .ifPresent(
-            instanceConnectionName ->
-                overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, instanceConnectionName));
-    isolationOverride
-        .map(Provider::get)
-        .ifPresent(isolation -> overrides.put(Environment.ISOLATION, isolation.name()));
-    return new JpaTransactionManagerImpl(create(overrides), clock);
+      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> beamCloudSqlConfigs, Clock clock) {
+    return new JpaTransactionManagerImpl(create(beamCloudSqlConfigs), clock);
+  }
+
+  @Provides
+  @Singleton
+  @BeamBulkQueryJpaTm
+  static JpaTransactionManager provideBeamBulkQueryJpaTm(
+      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> beamCloudSqlConfigs, Clock clock) {
+    return new JpaTransactionManagerImpl(
+        BulkQueryJpaFactory.createBulkQueryEntityManagerFactory(beamCloudSqlConfigs), clock);
   }
 
   @Provides
@@ -346,6 +349,17 @@ public abstract class PersistenceModule {
     }
   }
 
+  /** Types of {@link JpaTransactionManager JpaTransactionManagers}. */
+  public enum JpaTransactionManagerType {
+    /** The regular {@link JpaTransactionManager} for general use. */
+    REGULAR,
+    /**
+     * The {@link JpaTransactionManager} optimized for bulk loading multi-level JPA entities. Please
+     * see {@link google.registry.model.bulkquery.BulkQueryEntities} for more information.
+     */
+    BULK_QUERY
+  }
+
   /** Dagger qualifier for JDBC {@link Connection} with schema management privilege. */
   @Qualifier
   @Documented
@@ -357,11 +371,18 @@ public abstract class PersistenceModule {
   @interface AppEngineJpaTm {}
 
   /** Dagger qualifier for {@link JpaTransactionManager} used inside BEAM pipelines. */
-  // Note: @SocketFactoryJpaTm will be phased out in favor of this qualifier.
   @Qualifier
   @Documented
   public @interface BeamJpaTm {}
 
+  /**
+   * Dagger qualifier for {@link JpaTransactionManager} that uses an alternative entity model for
+   * faster bulk queries.
+   */
+  @Qualifier
+  @Documented
+  public @interface BeamBulkQueryJpaTm {}
+
   /** Dagger qualifier for {@link JpaTransactionManager} used for Nomulus tool. */
   @Qualifier
   @Documented