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 =
+    RegistryPipelineComponent registryPipelineComponent =
-        toRegistryPipelineComponent(registryOptions).getJpaTransactionManager();
+        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,
+      SqlCredentialStore credentialStore,
-      @Config("beamCloudSqlInstanceConnectionName") String instanceConnectionName,
+      @Config("instanceConnectionNameOverride")
-      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
+          Optional<Provider<String>> instanceConnectionNameOverride,
       @Config("beamIsolationOverride")
-          Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
+          Optional<Provider<TransactionIsolationLevel>> isolationOverride,
-    return createPartialSqlConfigs(
+      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs) {
-        jdbcUrl, instanceConnectionName, defaultConfigs, isolationOverride);
+    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,
+      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> beamCloudSqlConfigs, Clock clock) {
-      @Config("instanceConnectionNameOverride")
+    return new JpaTransactionManagerImpl(create(beamCloudSqlConfigs), clock);
-          Optional<Provider<String>> instanceConnectionNameOverride,
+  }
-      @Config("beamIsolationOverride")
+
-          Optional<Provider<TransactionIsolationLevel>> isolationOverride,
+  @Provides
-      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
+  @Singleton
-      Clock clock) {
+  @BeamBulkQueryJpaTm
-    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
+  static JpaTransactionManager provideBeamBulkQueryJpaTm(
-    // TODO(b/175700623): make sql username configurable from config file.
+      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> beamCloudSqlConfigs, Clock clock) {
-    SqlCredential credential = credentialStore.getCredential(new RobotUser(RobotId.NOMULUS));
+    return new JpaTransactionManagerImpl(
-    overrides.put(Environment.USER, credential.login());
+        BulkQueryJpaFactory.createBulkQueryEntityManagerFactory(beamCloudSqlConfigs), clock);
-    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);
   }
 
   @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

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

@@ -19,6 +19,7 @@ import static google.registry.beam.common.RegistryPipelineOptions.validateRegist
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import google.registry.config.RegistryEnvironment;
+import google.registry.persistence.PersistenceModule.JpaTransactionManagerType;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.testing.SystemPropertyExtension;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
@@ -123,4 +124,37 @@ class RegistryPipelineOptionsTest {
     validateRegistryPipelineOptions(options);
     assertThat(options.getProject()).isEqualTo("some-project");
   }
+
+  @Test
+  void jpaTransactionManagerType_default() {
+    RegistryPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(
+                "--registryEnvironment=" + RegistryEnvironment.UNITTEST.name())
+            .withValidation()
+            .as(RegistryPipelineOptions.class);
+    assertThat(options.getJpaTransactionManagerType()).isEqualTo(JpaTransactionManagerType.REGULAR);
+  }
+
+  @Test
+  void jpaTransactionManagerType_regularJpa() {
+    RegistryPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(
+                "--registryEnvironment=" + RegistryEnvironment.UNITTEST.name(),
+                "--jpaTransactionManagerType=REGULAR")
+            .withValidation()
+            .as(RegistryPipelineOptions.class);
+    assertThat(options.getJpaTransactionManagerType()).isEqualTo(JpaTransactionManagerType.REGULAR);
+  }
+
+  @Test
+  void jpaTransactionManagerType_bulkQueryJpa() {
+    RegistryPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(
+                "--registryEnvironment=" + RegistryEnvironment.UNITTEST.name(),
+                "--jpaTransactionManagerType=BULK_QUERY")
+            .withValidation()
+            .as(RegistryPipelineOptions.class);
+    assertThat(options.getJpaTransactionManagerType())
+        .isEqualTo(JpaTransactionManagerType.BULK_QUERY);
+  }
 }

core/src/test/java/google/registry/model/ImmutableObjectSubject.java

@@ -412,6 +412,10 @@ public final class ImmutableObjectSubject extends Subject {
     // don't use ImmutableMap or a stream->collect model since we can have nulls
     Map<Field, Object> result = new LinkedHashMap<>();
     for (Map.Entry<Field, Object> entry : originalFields.entrySet()) {
+      // TODO(b/203685960): filter by @DoNotCompare instead.
+      if (entry.getKey().isAnnotationPresent(ImmutableObject.Insignificant.class)) {
+        continue;
+      }
       if (!ignoredFieldSet.contains(entry.getKey().getName())) {
         result.put(entry.getKey(), entry.getValue());
       }
@@ -426,7 +430,9 @@ public final class ImmutableObjectSubject extends Subject {
     // don't use ImmutableMap or a stream->collect model since we can have nulls
     Map<Field, Object> result = new LinkedHashMap<>();
     for (Map.Entry<Field, Object> entry : originalFields.entrySet()) {
-      if (!entry.getKey().isAnnotationPresent(annotation)) {
+      // TODO(b/203685960): filter by @DoNotCompare instead.
+      if (!entry.getKey().isAnnotationPresent(annotation)
+          && !entry.getKey().isAnnotationPresent(ImmutableObject.Insignificant.class)) {
 
         // Perform any necessary substitutions.
         if (entry.getKey().isAnnotationPresent(ImmutableObject.EmptySetToNull.class)

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

@@ -0,0 +1,87 @@
+// 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 static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.domain.DomainHistory.DomainHistoryId;
+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.reporting.DomainTransactionRecord;
+import google.registry.persistence.VKey;
+
+/**
+ * Helpers for bulk-loading {@link google.registry.model.domain.DomainBase} and {@link
+ * google.registry.model.domain.DomainHistory} entities in <em>tests</em>.
+ */
+public class BulkQueryHelper {
+
+  static DomainBase loadAndAssembleDomainBase(String domainRepoId) {
+    return jpaTm()
+        .transact(
+            () ->
+                BulkQueryEntities.assembleDomainBase(
+                    jpaTm().loadByKey(DomainBaseLite.createVKey(domainRepoId)),
+                    jpaTm()
+                        .loadAllOfStream(GracePeriod.class)
+                        .filter(gracePeriod -> gracePeriod.getDomainRepoId().equals(domainRepoId))
+                        .collect(toImmutableSet()),
+                    jpaTm()
+                        .loadAllOfStream(DelegationSignerData.class)
+                        .filter(dsData -> dsData.getDomainRepoId().equals(domainRepoId))
+                        .collect(toImmutableSet()),
+                    jpaTm()
+                        .loadAllOfStream(DomainHost.class)
+                        .filter(domainHost -> domainHost.getDomainRepoId().equals(domainRepoId))
+                        .map(DomainHost::getHostVKey)
+                        .collect(toImmutableSet())));
+  }
+
+  static DomainHistory loadAndAssembleDomainHistory(DomainHistoryId domainHistoryId) {
+    return jpaTm()
+        .transact(
+            () ->
+                BulkQueryEntities.assembleDomainHistory(
+                    jpaTm().loadByKey(VKey.createSql(DomainHistoryLite.class, domainHistoryId)),
+                    jpaTm()
+                        .loadAllOfStream(DomainDsDataHistory.class)
+                        .filter(
+                            domainDsDataHistory ->
+                                domainDsDataHistory.getDomainHistoryId().equals(domainHistoryId))
+                        .collect(toImmutableSet()),
+                    jpaTm()
+                        .loadAllOfStream(DomainHistoryHost.class)
+                        .filter(
+                            domainHistoryHost ->
+                                domainHistoryHost.getDomainHistoryId().equals(domainHistoryId))
+                        .map(DomainHistoryHost::getHostVKey)
+                        .collect(toImmutableSet()),
+                    jpaTm()
+                        .loadAllOfStream(GracePeriodHistory.class)
+                        .filter(
+                            gracePeriodHistory ->
+                                gracePeriodHistory.getDomainHistoryId().equals(domainHistoryId))
+                        .collect(toImmutableSet()),
+                    jpaTm()
+                        .loadAllOfStream(DomainTransactionRecord.class)
+                        .filter(x -> true)
+                        .collect(toImmutableSet())));
+  }
+}

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

@@ -0,0 +1,117 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static org.joda.time.DateTimeZone.UTC;
+
+import com.google.common.collect.Sets;
+import com.google.common.collect.Sets.SetView;
+import com.google.common.truth.Truth8;
+import google.registry.model.domain.DomainBase;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.persistence.metamodel.Attribute;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for reading {@link DomainBaseLite}. */
+class DomainBaseLiteTest {
+
+  protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+
+  private final TestSetupHelper setupHelper = new TestSetupHelper(fakeClock);
+
+  @BeforeEach
+  void setUp() {
+    setupHelper.initializeAllEntities();
+  }
+
+  @AfterEach
+  void afterEach() {
+    setupHelper.tearDownBulkQueryJpaTm();
+  }
+
+  @Test
+  void readDomainHost() {
+    setupHelper.applyChangeToDomainAndHistory();
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    Truth8.assertThat(
+            jpaTm().transact(() -> jpaTm().loadAllOf(DomainHost.class)).stream()
+                .map(DomainHost::getHostVKey))
+        .containsExactly(setupHelper.host.createVKey());
+  }
+
+  @Test
+  void domainBaseLiteAttributes_versusDomainBase() {
+    Set<String> domainBaseAttributes =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .getEntityManager()
+                        .getMetamodel()
+                        .entity(DomainBase.class)
+                        .getAttributes())
+            .stream()
+            .map(Attribute::getName)
+            .collect(Collectors.toSet());
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    Set<String> domainBaseLiteAttributes =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .getEntityManager()
+                        .getMetamodel()
+                        .entity(DomainBaseLite.class)
+                        .getAttributes())
+            .stream()
+            .map(Attribute::getName)
+            .collect(Collectors.toSet());
+
+    assertThat(domainBaseAttributes).containsAtLeastElementsIn(domainBaseLiteAttributes);
+
+    SetView<?> excludedFromDomainBase =
+        Sets.difference(domainBaseAttributes, domainBaseLiteAttributes);
+    assertThat(excludedFromDomainBase)
+        .containsExactly("internalDelegationSignerData", "internalGracePeriods", "nsHosts");
+  }
+
+  @Test
+  void readDomainBaseLite_simple() {
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    assertThat(BulkQueryHelper.loadAndAssembleDomainBase(TestSetupHelper.DOMAIN_REPO_ID))
+        .isEqualTo(setupHelper.domain);
+  }
+
+  @Test
+  void readDomainBaseLite_full() {
+    setupHelper.applyChangeToDomainAndHistory();
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    assertThat(BulkQueryHelper.loadAndAssembleDomainBase(TestSetupHelper.DOMAIN_REPO_ID))
+        .isEqualTo(setupHelper.domain);
+  }
+}

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

@@ -0,0 +1,125 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static org.joda.time.DateTimeZone.UTC;
+
+import com.google.common.collect.Sets;
+import com.google.common.collect.Sets.SetView;
+import com.google.common.truth.Truth8;
+import google.registry.model.domain.DomainHistory;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.persistence.metamodel.Attribute;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link DomainHistoryLite}. */
+public class DomainHistoryLiteTest {
+
+  protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+
+  private final TestSetupHelper setupHelper = new TestSetupHelper(fakeClock);
+
+  @BeforeEach
+  void setUp() {
+    setupHelper.initializeAllEntities();
+  }
+
+  @AfterEach
+  void afterEach() {
+    setupHelper.tearDownBulkQueryJpaTm();
+  }
+
+  @Test
+  void readDomainHistoryHost() {
+    setupHelper.applyChangeToDomainAndHistory();
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    Truth8.assertThat(
+            jpaTm().transact(() -> jpaTm().loadAllOf(DomainHistoryHost.class)).stream()
+                .map(DomainHistoryHost::getHostVKey))
+        .containsExactly(setupHelper.host.createVKey());
+  }
+
+  @Test
+  void domainHistoryLiteAttributes_versusDomainHistory() {
+    Set<String> domainHistoryAttributes =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .getEntityManager()
+                        .getMetamodel()
+                        .entity(DomainHistory.class)
+                        .getAttributes())
+            .stream()
+            .map(Attribute::getName)
+            .collect(Collectors.toSet());
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    Set<String> domainHistoryLiteAttributes =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .getEntityManager()
+                        .getMetamodel()
+                        .entity(DomainHistoryLite.class)
+                        .getAttributes())
+            .stream()
+            .map(Attribute::getName)
+            .collect(Collectors.toSet());
+
+    assertThat(domainHistoryAttributes).containsAtLeastElementsIn(domainHistoryLiteAttributes);
+
+    SetView<?> excludedFromDomainHistory =
+        Sets.difference(domainHistoryAttributes, domainHistoryLiteAttributes);
+    assertThat(excludedFromDomainHistory)
+        .containsExactly(
+            "dsDataHistories",
+            "gracePeriodHistories",
+            "internalDomainTransactionRecords",
+            "nsHosts");
+  }
+
+  @Test
+  void readDomainHistory_noContent() {
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    assertThat(
+            BulkQueryHelper.loadAndAssembleDomainHistory(
+                setupHelper.domainHistory.getDomainHistoryId()))
+        .isEqualTo(setupHelper.domainHistory);
+  }
+
+  @Test
+  void readDomainHistory_full() {
+    setupHelper.applyChangeToDomainAndHistory();
+    setupHelper.setupBulkQueryJpaTm(appEngine);
+    assertThat(
+            BulkQueryHelper.loadAndAssembleDomainHistory(
+                setupHelper.domainHistory.getDomainHistoryId()))
+        .isEqualTo(setupHelper.domainHistory);
+  }
+}

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

@@ -0,0 +1,210 @@
+// 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 google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.SqlHelper.saveRegistrar;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.common.base.Ascii;
+import com.google.common.collect.ImmutableSet;
+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.Period;
+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.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.registrar.Registrar;
+import google.registry.model.reporting.DomainTransactionRecord;
+import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.tld.Registry;
+import google.registry.model.transfer.ContactTransferData;
+import google.registry.persistence.BulkQueryJpaFactory;
+import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
+import google.registry.testing.FakeClock;
+
+/** Entity creation utilities for domain-related tests. */
+class TestSetupHelper {
+
+  public static final String TLD = "tld";
+  public static final String DOMAIN_REPO_ID = "4-TLD";
+  public static final String DOMAIN_NAME = "example.tld";
+  public static final String REGISTRAR_ID = "AnRegistrar";
+
+  private final FakeClock fakeClock;
+
+  Registry registry;
+  Registrar registrar;
+  ContactResource contact;
+  DomainBase domain;
+  DomainHistory domainHistory;
+  HostResource host;
+
+  private JpaTransactionManager originalJpaTm;
+  private JpaTransactionManager bulkQueryJpaTm;
+
+  TestSetupHelper(FakeClock fakeClock) {
+    this.fakeClock = fakeClock;
+  }
+
+  void initializeAllEntities() {
+    registry = putInDb(DatabaseHelper.newRegistry(TLD, Ascii.toUpperCase(TLD)));
+    registrar = saveRegistrar(REGISTRAR_ID);
+    contact = putInDb(createContact(DOMAIN_REPO_ID, REGISTRAR_ID));
+    domain = putInDb(createSimpleDomain(contact));
+    domainHistory = putInDb(createHistoryWithoutContent(domain, fakeClock));
+    host = putInDb(createHost());
+  }
+
+  void applyChangeToDomainAndHistory() {
+    domain = putInDb(createFullDomain(contact, host, fakeClock));
+    domainHistory = putInDb(createFullHistory(domain, fakeClock));
+  }
+
+  void setupBulkQueryJpaTm(AppEngineExtension appEngineExtension) {
+    bulkQueryJpaTm =
+        BulkQueryJpaFactory.createBulkQueryJpaTransactionManager(
+            appEngineExtension
+                .getJpaIntegrationTestExtension()
+                .map(JpaIntegrationTestExtension::getJpaProperties)
+                .orElseThrow(
+                    () -> new IllegalStateException("Expecting JpaIntegrationTestExtension.")),
+            fakeClock);
+    originalJpaTm = TransactionManagerFactory.jpaTm();
+    TransactionManagerFactory.setJpaTm(() -> bulkQueryJpaTm);
+  }
+
+  void tearDownBulkQueryJpaTm() {
+    if (bulkQueryJpaTm != null) {
+      bulkQueryJpaTm.teardown();
+      TransactionManagerFactory.setJpaTm(() -> originalJpaTm);
+    }
+  }
+
+  static ContactResource createContact(String repoId, String registrarId) {
+    return new ContactResource.Builder()
+        .setRepoId(repoId)
+        .setCreationRegistrarId(registrarId)
+        .setTransferData(new ContactTransferData.Builder().build())
+        .setPersistedCurrentSponsorRegistrarId(registrarId)
+        .build();
+  }
+
+  static DomainBase createSimpleDomain(ContactResource contact) {
+    return DatabaseHelper.newDomainBase(DOMAIN_NAME, DOMAIN_REPO_ID, contact)
+        .asBuilder()
+        .setCreationRegistrarId(REGISTRAR_ID)
+        .setPersistedCurrentSponsorRegistrarId(REGISTRAR_ID)
+        .build();
+  }
+
+  static DomainBase createFullDomain(
+      ContactResource contact, HostResource host, FakeClock fakeClock) {
+    return createSimpleDomain(contact)
+        .asBuilder()
+        .setDomainName(DOMAIN_NAME)
+        .setRepoId(DOMAIN_REPO_ID)
+        .setCreationRegistrarId(REGISTRAR_ID)
+        .setLastEppUpdateTime(fakeClock.nowUtc())
+        .setLastEppUpdateRegistrarId(REGISTRAR_ID)
+        .setLastTransferTime(fakeClock.nowUtc())
+        .setNameservers(host.createVKey())
+        .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))
+        .setContacts(
+            ImmutableSet.of(
+                DesignatedContact.create(DesignatedContact.Type.ADMIN, contact.createVKey())))
+        .setSubordinateHosts(ImmutableSet.of("ns1.example.com"))
+        .setPersistedCurrentSponsorRegistrarId(REGISTRAR_ID)
+        .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))
+        .setSmdId("smdid")
+        .addGracePeriod(
+            GracePeriod.create(
+                GracePeriodStatus.ADD, DOMAIN_REPO_ID, END_OF_TIME, REGISTRAR_ID, null, 100L))
+        .build();
+  }
+
+  static HostResource createHost() {
+    return new HostResource.Builder()
+        .setRepoId("host1")
+        .setHostName("ns1.example.com")
+        .setCreationRegistrarId(REGISTRAR_ID)
+        .setPersistedCurrentSponsorRegistrarId(REGISTRAR_ID)
+        .build();
+  }
+
+  static DomainTransactionRecord createDomainTransactionRecord(FakeClock fakeClock) {
+    return new DomainTransactionRecord.Builder()
+        .setTld(TLD)
+        .setReportingTime(fakeClock.nowUtc())
+        .setReportField(TransactionReportField.NET_ADDS_1_YR)
+        .setReportAmount(1)
+        .build();
+  }
+
+  static DomainHistory createHistoryWithoutContent(DomainBase domain, FakeClock fakeClock) {
+    return new DomainHistory.Builder()
+        .setType(HistoryEntry.Type.DOMAIN_CREATE)
+        .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+        .setModificationTime(fakeClock.nowUtc())
+        .setRegistrarId(REGISTRAR_ID)
+        .setTrid(Trid.create("ABC-123", "server-trid"))
+        .setBySuperuser(false)
+        .setReason("reason")
+        .setRequestedByRegistrar(true)
+        .setDomainRepoId(domain.getRepoId())
+        .setOtherRegistrarId("otherClient")
+        .setPeriod(Period.create(1, Period.Unit.YEARS))
+        .build();
+  }
+
+  static DomainHistory createFullHistory(DomainBase domain, FakeClock fakeClock) {
+    return createHistoryWithoutContent(domain, fakeClock)
+        .asBuilder()
+        .setType(HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE)
+        .setDomain(domain)
+        .setDomainTransactionRecords(ImmutableSet.of(createDomainTransactionRecord(fakeClock)))
+        .build();
+  }
+
+  static <T> T putInDb(T entity) {
+    jpaTm().transact(() -> jpaTm().put(entity));
+    return jpaTm().transact(() -> jpaTm().loadByEntity(entity));
+  }
+}

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java

@@ -45,7 +45,6 @@ import java.sql.Driver;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -172,24 +171,39 @@ abstract class JpaTransactionManagerExtension implements BeforeEachCallback, Aft
       exporter.export(getTestEntities(), tempSqlFile);
       executeSql(new String(Files.readAllBytes(tempSqlFile.toPath()), StandardCharsets.UTF_8));
     }
+    assertReasonableNumDbConnections();
+    emf = createEntityManagerFactory(getJpaProperties());
+    emfEntityHash = entityHash;
+  }
 
-    ImmutableMap<String, String> properties = PersistenceModule.provideDefaultDatabaseConfigs();
+  /**
+   * Returns the full set of properties for setting up JPA {@link EntityManagerFactory} to the test
+   * database. This allows creation of customized JPA by individual tests.
+   *
+   * <p>Test that create {@code EntityManagerFactory} instances are reponsible for tearing them
+   * down.
+   */
+  public ImmutableMap<String, String> getJpaProperties() {
+    Map<String, String> mergedProperties =
+        Maps.newHashMap(PersistenceModule.provideDefaultDatabaseConfigs());
     if (!userProperties.isEmpty()) {
-      // If there are user properties, create a new properties object with these added.
-      Map<String, String> mergedProperties = Maps.newHashMap();
-      mergedProperties.putAll(properties);
       mergedProperties.putAll(userProperties);
-      properties = ImmutableMap.copyOf(mergedProperties);
     }
+    mergedProperties.put(Environment.URL, getJdbcUrl());
+    mergedProperties.put(Environment.USER, database.getUsername());
+    mergedProperties.put(Environment.PASS, database.getPassword());
+    // Tell Postgresql JDBC driver to retry on errors caused by out-of-band schema change between
+    // tests while the connection pool stays open (e.g., "cached plan must not change result type").
+    // We don't set this property in production since it has performance impact, and production
+    // schema is always compatible with the binary (enforced by our release process).
+    mergedProperties.put("hibernate.hikari.dataSource.autosave", "conservative");
+
     // Forbid Hibernate push to stay consistent with flyway-based schema management.
     checkState(
-        Objects.equals(properties.get(Environment.HBM2DDL_AUTO), "none"),
+        Objects.equals(mergedProperties.get(Environment.HBM2DDL_AUTO), "none"),
         "The HBM2DDL_AUTO property must be 'none'.");
-    assertReasonableNumDbConnections();
+
-    emf =
+    return ImmutableMap.copyOf(mergedProperties);
-        createEntityManagerFactory(
-            getJdbcUrl(), database.getUsername(), database.getPassword(), properties);
-    emfEntityHash = entityHash;
   }
 
   @Override
@@ -307,15 +321,7 @@ abstract class JpaTransactionManagerExtension implements BeforeEachCallback, Aft
   }
 
   /** Constructs the {@link EntityManagerFactory} instance. */
-  private EntityManagerFactory createEntityManagerFactory(
+  private EntityManagerFactory createEntityManagerFactory(ImmutableMap<String, String> properties) {
-      String jdbcUrl, String username, String password, ImmutableMap<String, String> configs) {
-    HashMap<String, String> properties = Maps.newHashMap(configs);
-    properties.put(Environment.URL, jdbcUrl);
-    properties.put(Environment.USER, username);
-    properties.put(Environment.PASS, password);
-    // Tell Postgresql JDBC driver to expect out-of-band schema change.
-    properties.put("hibernate.hikari.dataSource.autosave", "conservative");
-
     ParsedPersistenceXmlDescriptor descriptor =
         PersistenceXmlUtility.getParsedPersistenceXmlDescriptor();
 

core/src/test/java/google/registry/testing/AppEngineExtension.java

@@ -146,6 +146,10 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
   private ImmutableList<Class<?>> ofyTestEntities;
   private ImmutableList<Class<?>> jpaTestEntities;
 
+  public Optional<JpaIntegrationTestExtension> getJpaIntegrationTestExtension() {
+    return Optional.ofNullable(jpaIntegrationTestExtension);
+  }
+
   /** Builder for {@link AppEngineExtension}. */
   public static class Builder {
 

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -412,12 +412,12 @@
 
     create table "DomainTransactionRecord" (
        id  bigserial not null,
+        domain_repo_id text,
+        history_revision_id int8,
         report_amount int4 not null,
         report_field text not null,
         reporting_time timestamptz not null,
         tld text not null,
-        domain_repo_id text,
-        history_revision_id int8,
         primary key (id)
     );