Make TransactionManager.loadAllOf() smart w.r.t the cross-TLD entity group (#1040)

* Make TransactionManager.loadAllOf() smart w.r.t the cross-TLD entity group

The loadAllOf() method will now automatically append the cross-TLD entity group
ancestor query as necessary, iff the entity class being loaded is tagged with
the new @IsCrossTld annotation.

* Add tests

core/src/main/java/google/registry/export/AnnotatedEntities.java

@@ -21,12 +21,13 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityClasses;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.annotations.VirtualEntity;
 
 /** Constants related to export code. */
-public final class ExportConstants {
+public final class AnnotatedEntities {
 
   /** Returns the names of kinds to include in Datastore backups. */
   public static ImmutableSet<String> getBackupKinds() {
@@ -49,4 +50,13 @@ public final class ExportConstants {
         .map(Key::getKind)
         .collect(toImmutableSortedSet(Ordering.natural()));
   }
+
+  /** Returns the names of kinds that are in the cross-TLD entity group. */
+  public static ImmutableSet<String> getCrossTldKinds() {
+    return EntityClasses.ALL_CLASSES.stream()
+        .filter(hasAnnotation(InCrossTld.class))
+        .filter(hasAnnotation(VirtualEntity.class).negate())
+        .map(Key::getKind)
+        .collect(toImmutableSortedSet(Ordering.natural()));
+  }
 }

core/src/main/java/google/registry/export/BackupDatastoreAction.java

@@ -66,12 +66,13 @@ public class BackupDatastoreAction implements Runnable {
     try {
       Operation backup =
           datastoreAdmin
-              .export(RegistryConfig.getDatastoreBackupsBucket(), ExportConstants.getBackupKinds())
+              .export(
+                  RegistryConfig.getDatastoreBackupsBucket(), AnnotatedEntities.getBackupKinds())
               .execute();
 
       String backupName = backup.getName();
       // Enqueue a poll task to monitor the backup and load REPORTING-related kinds into bigquery.
-      enqueuePollTask(backupName, ExportConstants.getReportingKinds());
+      enqueuePollTask(backupName, AnnotatedEntities.getReportingKinds());
       String message =
           String.format(
               "Datastore backup started with name: %s\nSaving to %s",

core/src/main/java/google/registry/model/annotations/InCrossTld.java

@@ -0,0 +1,34 @@
+// 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.annotations;
+
+import com.googlecode.objectify.annotation.Entity;
+import google.registry.model.common.EntityGroupRoot;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation for an Objectify {@link Entity} to indicate that it is in the cross-TLD entity group.
+ *
+ * <p>This means that the entity's <code>@Parent</code> field has to have the value of {@link
+ * EntityGroupRoot#getCrossTldKey}.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE})
+@Inherited
+public @interface InCrossTld {}

core/src/main/java/google/registry/model/common/CrossTldSingleton.java

@@ -20,11 +20,13 @@ import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import javax.persistence.MappedSuperclass;
 import javax.persistence.Transient;
 
 /** A singleton entity in Datastore. */
 @MappedSuperclass
+@InCrossTld
 public abstract class CrossTldSingleton extends ImmutableObject {
 
   public static final long SINGLETON_ID = 1; // There is always exactly one of these.

core/src/main/java/google/registry/model/common/Cursor.java

@@ -29,6 +29,7 @@ import com.googlecode.objectify.annotation.OnLoad;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.common.Cursor.CursorId;
 import google.registry.model.registry.Registry;
 import google.registry.persistence.VKey;
@@ -51,6 +52,7 @@ import org.joda.time.DateTime;
 @Entity
 @javax.persistence.Entity
 @IdClass(CursorId.class)
+@InCrossTld
 public class Cursor extends ImmutableObject implements DatastoreAndSqlEntity {
 
   /** The scope of a global cursor. A global cursor is a cursor that is not specific to one tld. */

core/src/main/java/google/registry/model/common/DatabaseTransitionSchedule.java

@@ -32,6 +32,7 @@ import com.googlecode.objectify.annotation.Mapify;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.common.TimedTransitionProperty.TimeMapper;
 import google.registry.model.common.TimedTransitionProperty.TimedTransition;
 import google.registry.model.registry.label.PremiumList;
@@ -46,6 +47,7 @@ import org.joda.time.DateTime;
 
 @Entity
 @Immutable
+@InCrossTld
 public class DatabaseTransitionSchedule extends ImmutableObject implements DatastoreOnlyEntity {
 
   /**

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -17,6 +17,7 @@ package google.registry.model.ofy;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -29,6 +30,8 @@ import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
+import com.googlecode.objectify.cmd.Query;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.contact.ContactHistory;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.host.HostHistory;
@@ -251,7 +254,13 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
-    return ImmutableList.copyOf(getOfy().load().type(clazz));
+    Query<T> query = getOfy().load().type(clazz);
+    // If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
+    // query to give us strong transactional consistency.
+    if (clazz.isAnnotationPresent(InCrossTld.class)) {
+      query = query.ancestor(getCrossTldKey());
+    }
+    return ImmutableList.copyOf(query);
   }
 
   @Override

core/src/main/java/google/registry/model/registrar/Registrar.java

@@ -72,6 +72,7 @@ import google.registry.model.ImmutableObject;
 import google.registry.model.JsonMapBuilder;
 import google.registry.model.Jsonifiable;
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.registrar.Registrar.BillingAccountEntry.CurrencyMapper;
@@ -112,6 +113,7 @@ import org.joda.time.DateTime;
           columnList = "ianaIdentifier",
           name = "registrar_iana_identifier_idx"),
     })
+@InCrossTld
 public class Registrar extends ImmutableObject
     implements Buildable, DatastoreAndSqlEntity, Jsonifiable {
 
@@ -985,9 +987,7 @@ public class Registrar extends ImmutableObject
 
   /** Loads all registrar entities directly from Datastore. */
   public static Iterable<Registrar> loadAll() {
-    return tm().isOfy()
+    return transactIfJpaTm(() -> tm().loadAllOf(Registrar.class));
-        ? ImmutableList.copyOf(ofy().load().type(Registrar.class).ancestor(getCrossTldKey()))
-        : tm().transact(() -> tm().loadAllOf(Registrar.class));
   }
 
   /** Loads all registrar entities using an in-memory cache. */

core/src/main/java/google/registry/model/registrar/RegistrarContact.java

@@ -44,6 +44,7 @@ import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
 import google.registry.model.JsonMapBuilder;
 import google.registry.model.Jsonifiable;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.registrar.RegistrarContact.RegistrarPocId;
 import google.registry.persistence.VKey;
@@ -77,6 +78,7 @@ import javax.persistence.Transient;
       @javax.persistence.Index(columnList = "gaeUserId", name = "registrarpoc_gae_user_id_idx")
     })
 @IdClass(RegistrarPocId.class)
+@InCrossTld
 public class RegistrarContact extends ImmutableObject
     implements DatastoreAndSqlEntity, Jsonifiable {
 

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

@@ -52,6 +52,7 @@ import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.common.TimedTransitionProperty;
@@ -86,6 +87,7 @@ import org.joda.time.Duration;
 @ReportedOn
 @Entity
 @javax.persistence.Entity(name = "Tld")
+@InCrossTld
 public class Registry extends ImmutableObject implements Buildable, DatastoreAndSqlEntity {
 
   @Parent @Transient Key<EntityGroupRoot> parent = getCrossTldKey();

core/src/main/java/google/registry/model/registry/label/BaseDomainLabelList.java

@@ -35,6 +35,7 @@ import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.ReservedList.ReservedListEntry;
@@ -59,6 +60,7 @@ import org.joda.time.DateTime;
  *     must subclass {@link DomainLabelEntry}.
  */
 @MappedSuperclass
+@InCrossTld
 public abstract class BaseDomainLabelList<T extends Comparable<?>, R extends DomainLabelEntry<T, ?>>
     extends ImmutableObject implements Buildable {
 

core/src/main/java/google/registry/model/registry/label/PremiumList.java

@@ -31,6 +31,7 @@ import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.registry.Registry;
 import google.registry.schema.replay.DatastoreOnlyEntity;
@@ -96,6 +97,7 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   /** Virtual parent entity for premium list entry entities associated with a single revision. */
   @ReportedOn
   @Entity
+  @InCrossTld
   public static class PremiumListRevision extends ImmutableObject implements DatastoreOnlyEntity {
 
     @Parent Key<PremiumList> parent;
@@ -195,6 +197,7 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
    */
   @ReportedOn
   @Entity
+  @InCrossTld
   public static class PremiumListEntry extends DomainLabelEntry<Money, PremiumListEntry>
       implements Buildable, DatastoreOnlyEntity {
 

core/src/main/java/google/registry/model/server/KmsSecret.java

@@ -21,6 +21,7 @@ import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.schema.replay.DatastoreOnlyEntity;
@@ -28,6 +29,7 @@ import google.registry.schema.replay.DatastoreOnlyEntity;
 /** Pointer to the latest {@link KmsSecretRevision}. */
 @Entity
 @ReportedOn
+@InCrossTld
 public class KmsSecret extends ImmutableObject implements DatastoreOnlyEntity {
 
   /** The unique name of this {@link KmsSecret}. */

core/src/main/java/google/registry/model/server/KmsSecretRevision.java

@@ -26,6 +26,7 @@ import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.schema.replay.NonReplicatedEntity;
 import javax.persistence.Column;
@@ -58,6 +59,7 @@ import javax.persistence.Transient;
 @ReportedOn
 @javax.persistence.Entity(name = "KmsSecret")
 @Table(indexes = {@Index(columnList = "secretName")})
+@InCrossTld
 public class KmsSecretRevision extends ImmutableObject implements NonReplicatedEntity {
 
   /**

core/src/main/java/google/registry/model/smd/SignedMarkRevocationList.java

@@ -30,6 +30,7 @@ import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.OnSave;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.EntityGroupRoot;
@@ -67,6 +68,7 @@ import org.joda.time.DateTime;
 @Entity
 @javax.persistence.Entity
 @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
+@InCrossTld
 public class SignedMarkRevocationList extends ImmutableObject implements NonReplicatedEntity {
 
   @VisibleForTesting static final int SHARD_SIZE = 10000;

core/src/main/java/google/registry/model/tmch/ClaimsListShard.java

@@ -35,6 +35,7 @@ import com.googlecode.objectify.annotation.OnSave;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.annotations.VirtualEntity;
@@ -89,6 +90,7 @@ import org.joda.time.DateTime;
 @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
 @javax.persistence.Entity(name = "ClaimsList")
 @Table
+@InCrossTld
 public class ClaimsListShard extends ImmutableObject implements NonReplicatedEntity {
 
   /** The number of claims list entries to store per shard. */

core/src/main/java/google/registry/persistence/transaction/TransactionManager.java

@@ -17,6 +17,7 @@ package google.registry.persistence.transaction;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.persistence.VKey;
 import java.util.NoSuchElementException;
 import java.util.Optional;
@@ -239,7 +240,9 @@ public interface TransactionManager {
   /**
    * Returns a stream of all entities of the given type that exist in the database.
    *
-   * <p>The resulting stream is empty if there are no entities of this type.
+   * <p>The resulting stream is empty if there are no entities of this type. In Datastore mode, if
+   * the class is a member of the cross-TLD entity group (i.e. if it has the {@link InCrossTld}
+   * annotation, then the correct ancestor query will automatically be applied.
    */
   <T> ImmutableList<T> loadAllOf(Class<T> clazz);
 

core/src/main/java/google/registry/tools/LoadSnapshotCommand.java

@@ -25,7 +25,7 @@ import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import google.registry.bigquery.BigqueryUtils.SourceFormat;
-import google.registry.export.ExportConstants;
+import google.registry.export.AnnotatedEntities;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -48,7 +48,7 @@ final class LoadSnapshotCommand extends BigqueryCommand {
   @Parameter(
       names = "--kinds",
       description = "List of Datastore kinds for which to import snapshot data.")
-  private List<String> kindNames = new ArrayList<>(ExportConstants.getReportingKinds());
+  private List<String> kindNames = new ArrayList<>(AnnotatedEntities.getReportingKinds());
 
   /** Runs the main snapshot import logic. */
   @Override

core/src/test/java/google/registry/export/AnnotatedEntitiesTest.java

@@ -19,8 +19,9 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.io.Resources.getResource;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
-import static google.registry.export.ExportConstants.getBackupKinds;
+import static google.registry.export.AnnotatedEntities.getBackupKinds;
-import static google.registry.export.ExportConstants.getReportingKinds;
+import static google.registry.export.AnnotatedEntities.getCrossTldKinds;
+import static google.registry.export.AnnotatedEntities.getReportingKinds;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 
 import com.google.common.base.Joiner;
@@ -31,13 +32,15 @@ import com.google.common.collect.Streams;
 import com.google.re2j.Pattern;
 import org.junit.jupiter.api.Test;
 
-/** Unit tests for {@link ExportConstants}. */
+/** Unit tests for {@link AnnotatedEntities}. */
-class ExportConstantsTest {
+class AnnotatedEntitiesTest {
 
   private static final String GOLDEN_BACKUP_KINDS_FILENAME = "backup_kinds.txt";
 
   private static final String GOLDEN_REPORTING_KINDS_FILENAME = "reporting_kinds.txt";
 
+  private static final String GOLDEN_CROSSTLD_KINDS_FILENAME = "crosstld_kinds.txt";
+
   private static final String UPDATE_INSTRUCTIONS_TEMPLATE = Joiner.on('\n').join(
       "",
       repeat("-", 80),
@@ -50,15 +53,20 @@ class ExportConstantsTest {
       "");
 
   @Test
-  void testBackupKinds_matchGoldenBackupKindsFile() {
+  void testBackupKinds_matchGoldenFile() {
     checkKindsMatchGoldenFile("backed-up", GOLDEN_BACKUP_KINDS_FILENAME, getBackupKinds());
   }
 
   @Test
-  void testReportingKinds_matchGoldenReportingKindsFile() {
+  void testReportingKinds_matchGoldenFile() {
     checkKindsMatchGoldenFile("reporting", GOLDEN_REPORTING_KINDS_FILENAME, getReportingKinds());
   }
 
+  @Test
+  void testCrossTldKinds_matchGoldenFile() {
+    checkKindsMatchGoldenFile("crosstld", GOLDEN_CROSSTLD_KINDS_FILENAME, getCrossTldKinds());
+  }
+
   @Test
   void testReportingKinds_areSubsetOfBackupKinds() {
     assertThat(getBackupKinds()).containsAtLeastElementsIn(getReportingKinds());
@@ -70,7 +78,7 @@ class ExportConstantsTest {
         String.format(
             UPDATE_INSTRUCTIONS_TEMPLATE,
             kindsName,
-            getResource(ExportConstantsTest.class, goldenFilename).toString(),
+            getResource(AnnotatedEntitiesTest.class, goldenFilename).toString(),
             Joiner.on('\n').join(actualKinds));
     assertWithMessage(updateInstructions)
         .that(actualKinds)
@@ -85,7 +93,7 @@ class ExportConstantsTest {
    * @return ImmutableList<String>
    */
   private static ImmutableList<String> extractListFromFile(String filename) {
-    String fileContents = readResourceUtf8(ExportConstantsTest.class, filename);
+    String fileContents = readResourceUtf8(AnnotatedEntitiesTest.class, filename);
     final Pattern stripComments = Pattern.compile("\\s*#.*$");
     return Streams.stream(Splitter.on('\n').split(fileContents.trim()))
         .map(line -> stripComments.matcher(line).replaceFirst(""))

core/src/test/java/google/registry/export/BackupDatastoreActionTest.java

@@ -54,7 +54,7 @@ public class BackupDatastoreActionTest {
     action.response = response;
 
     when(datastoreAdmin.export(
-            "gs://registry-project-id-datastore-backups", ExportConstants.getBackupKinds()))
+            "gs://registry-project-id-datastore-backups", AnnotatedEntities.getBackupKinds()))
         .thenReturn(exportRequest);
     when(exportRequest.execute()).thenReturn(backupOperation);
     when(backupOperation.getName())
@@ -73,7 +73,7 @@ public class BackupDatastoreActionTest {
             .param(CHECK_BACKUP_NAME_PARAM, "projects/registry-project-id/operations/ASA1ODYwNjc")
             .param(
                 CHECK_BACKUP_KINDS_TO_LOAD_PARAM,
-                Joiner.on(",").join(ExportConstants.getReportingKinds()))
+                Joiner.on(",").join(AnnotatedEntities.getReportingKinds()))
             .method("POST"));
     assertThat(response.getPayload())
         .isEqualTo(

core/src/test/java/google/registry/model/ofy/DatastoreTransactionManagerTest.java

@@ -0,0 +1,72 @@
+// 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.ofy;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.testing.DatabaseHelper.persistResources;
+
+import com.google.common.collect.ImmutableList;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.Entity;
+import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Parent;
+import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
+import google.registry.model.common.EntityGroupRoot;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
+import google.registry.testing.AppEngineExtension;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link DatastoreTransactionManager}. */
+public class DatastoreTransactionManagerTest {
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder()
+          .withDatastore()
+          .withOfyTestEntities(InCrossTldTestEntity.class)
+          .build();
+
+  @Test
+  void test_loadAllOf_usesAncestorQuery() {
+    InCrossTldTestEntity foo = new InCrossTldTestEntity("foo");
+    InCrossTldTestEntity bar = new InCrossTldTestEntity("bar");
+    InCrossTldTestEntity baz = new InCrossTldTestEntity("baz");
+    baz.parent = null;
+    persistResources(ImmutableList.of(foo, bar, baz));
+    // baz is excluded by the cross-TLD ancestor query
+    assertThat(ofyTm().loadAllOf(InCrossTldTestEntity.class)).containsExactly(foo, bar);
+  }
+
+  @Entity
+  @EntityForTesting
+  @InCrossTld
+  private static class InCrossTldTestEntity extends ImmutableObject {
+
+    @Id String name;
+    @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
+
+    private InCrossTldTestEntity(String name) {
+      this.name = name;
+    }
+
+    // Needs to exist to make Objectify happy.
+    @SuppressWarnings("unused")
+    private InCrossTldTestEntity() {}
+  }
+}

core/src/test/resources/google/registry/export/crosstld_kinds.txt

@@ -0,0 +1,16 @@
+ClaimsListShard
+ClaimsListSingleton
+Cursor
+DatabaseTransitionSchedule
+KmsSecret
+KmsSecretRevision
+PremiumList
+PremiumListEntry
+PremiumListRevision
+Registrar
+RegistrarContact
+Registry
+ReservedList
+ServerSecret
+SignedMarkRevocationList
+TmchCrl