Remove some unnecessary Ofy key creation (#1212)

core/src/main/java/google/registry/flows/FlowUtils.java

@@ -15,6 +15,7 @@
 package google.registry.flows;
 
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.xml.ValidationMode.LENIENT;
 import static google.registry.xml.ValidationMode.STRICT;
@@ -33,7 +34,6 @@ import google.registry.model.eppcommon.EppXmlTransformer;
 import google.registry.model.eppinput.EppInput.WrongProtocolVersionException;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.host.InetAddressAdapter.IpVersionMismatchException;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.translators.CurrencyUnitAdapter.UnknownCurrencyException;
 import google.registry.xml.XmlException;
@@ -105,7 +105,7 @@ public final class FlowUtils {
 
   public static <H extends HistoryEntry> Key<H> createHistoryKey(
       EppResource parent, Class<H> clazz) {
-    return Key.create(Key.create(parent), clazz, ObjectifyService.allocateId());
+    return Key.create(Key.create(parent), clazz, allocateId());
   }
 
   /** Registrar is not logged in. */

core/src/main/java/google/registry/flows/contact/ContactCreateFlow.java

@@ -19,6 +19,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceDoesNotExist
 import static google.registry.flows.contact.ContactFlowUtils.validateAsciiPostalInfo;
 import static google.registry.flows.contact.ContactFlowUtils.validateContactAgainstPolicy;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -41,7 +42,6 @@ import google.registry.model.eppoutput.CreateData.ContactCreateData;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import javax.inject.Inject;
@@ -81,7 +81,7 @@ public final class ContactCreateFlow implements TransactionalFlow {
             .setAuthInfo(command.getAuthInfo())
             .setCreationClientId(clientId)
             .setPersistedCurrentSponsorClientId(clientId)
-            .setRepoId(createRepoId(ObjectifyService.allocateId(), roidSuffix))
+            .setRepoId(createRepoId(allocateId(), roidSuffix))
             .setFaxNumber(command.getFax())
             .setVoiceNumber(command.getVoice())
             .setDisclose(command.getDisclose())

core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java

@@ -42,6 +42,7 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNo
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.eppcommon.StatusValue.SERVER_HOLD;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.QUIET_PERIOD;
@@ -99,7 +100,6 @@ import google.registry.model.eppoutput.CreateData.DomainCreateData;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessage.Autorenew;
@@ -302,12 +302,9 @@ public class DomainCreateFlow implements TransactionalFlow {
     Optional<SecDnsCreateExtension> secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
     DateTime registrationExpirationTime = leapSafeAddYears(now, years);
-    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
+    String repoId = createDomainRepoId(allocateId(), registry.getTldStr());
     Key<DomainHistory> domainHistoryKey =
-        Key.create(
-            Key.create(DomainBase.class, repoId),
-            DomainHistory.class,
-            ObjectifyService.allocateId());
+        Key.create(Key.create(DomainBase.class, repoId), DomainHistory.class, allocateId());
     historyBuilder.setId(domainHistoryKey.getId());
     // Bill for the create.
     BillingEvent.OneTime createBillingEvent =

core/src/main/java/google/registry/flows/host/HostCreateFlow.java

@@ -21,6 +21,7 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 
@@ -49,7 +50,6 @@ import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
@@ -126,7 +126,7 @@ public final class HostCreateFlow implements TransactionalFlow {
             .setPersistedCurrentSponsorClientId(clientId)
             .setHostName(targetId)
             .setInetAddresses(command.getInetAddresses())
-            .setRepoId(createRepoId(ObjectifyService.allocateId(), roidSuffix))
+            .setRepoId(createRepoId(allocateId(), roidSuffix))
             .setSuperordinateDomain(superordinateDomain.map(DomainBase::createVKey).orElse(null))
             .build();
     historyBuilder.setType(HistoryEntry.Type.HOST_CREATE).setModificationTime(now).setHost(newHost);

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

@@ -16,9 +16,10 @@ package google.registry.model;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+import static google.registry.model.IdService.allocateId;
+import static google.registry.model.ModelUtils.getAllFields;
 
-import google.registry.model.ofy.ObjectifyService;
+import com.googlecode.objectify.annotation.Id;
 import google.registry.util.TypeUtils.TypeInstantiator;
 import java.lang.reflect.Field;
 import java.util.Optional;
@@ -54,25 +55,18 @@ public interface Buildable {
     /** Build the instance. */
     public S build() {
       try {
-        // If this object has a Long or long Objectify @Id field that is not set, set it now.
-        Field idField = null;
-        try {
-          idField =
-              ModelUtils.getAllFields(instance.getClass())
-                  .get(
-                      auditedOfy()
-                          .factory()
-                          .getMetadata(instance.getClass())
-                          .getKeyMetadata()
-                          .getIdFieldName());
-        } catch (Exception e) {
-          // Expected if the class is not registered with Objectify.
-        }
+        // If this object has a Long or long Objectify @Id field that is not set, set it now. For
+        // any entity it has one and only one @Id field in its class hierarchy.
+        Field idField =
+            getAllFields(instance.getClass()).values().stream()
+                .filter(field -> field.isAnnotationPresent(Id.class))
+                .findFirst()
+                .orElse(null);
         if (idField != null
             && !idField.getType().equals(String.class)
             && Optional.ofNullable((Long) ModelUtils.getFieldValue(instance, idField))
                 .orElse(0L) == 0) {
-          ModelUtils.setFieldValue(instance, idField, ObjectifyService.allocateId());
+          ModelUtils.setFieldValue(instance, idField, allocateId());
         }
         return instance;
       } finally {

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

@@ -0,0 +1,61 @@
+// 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;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.appengine.api.datastore.DatastoreServiceFactory;
+import com.google.common.annotations.VisibleForTesting;
+import google.registry.config.RegistryEnvironment;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Allocates a globally unique {@link Long} number to use as an Ofy {@code @Id}.
+ *
+ * <p>In non-test environments the Id is generated by Datastore, whereas in tests it's from an
+ * atomic long number that's incremented every time this method is called.
+ */
+public final class IdService {
+
+  /**
+   * A placeholder String passed into DatastoreService.allocateIds that ensures that all ids are
+   * initialized from the same id pool.
+   */
+  private static final String APP_WIDE_ALLOCATION_KIND = "common";
+
+  /** Counts of used ids for use in unit tests. Outside tests this is never used. */
+  private static final AtomicLong nextTestId = new AtomicLong(1); // ids cannot be zero
+
+  /** Allocates an id. */
+  public static long allocateId() {
+    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
+        ? nextTestId.getAndIncrement()
+        : DatastoreServiceFactory.getDatastoreService()
+            .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
+            .iterator()
+            .next()
+            .getId();
+  }
+
+  /** Resets the global test id counter (i.e. sets the next id to 1). */
+  @VisibleForTesting
+  public static void resetNextTestId() {
+    checkState(
+        RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get()),
+        "Can't call resetTestIdCounts() from RegistryEnvironment.%s",
+        RegistryEnvironment.get());
+    nextTestId.set(1); // ids cannot be zero
+  }
+}

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

@@ -14,11 +14,13 @@
 
 package google.registry.model.common;
 
+import com.google.apphosting.api.ApiProxy;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.model.BackupGroupRoot;
 import google.registry.schema.replay.DatastoreOnlyEntity;
+import javax.annotation.Nullable;
 
 /**
  * The root key for the entity group which is known as the cross-tld entity group for historical
@@ -42,7 +44,13 @@ public class EntityGroupRoot extends BackupGroupRoot implements DatastoreOnlyEnt
   private String id;
 
   /** The root key for cross-tld resources such as registrars. */
-  public static Key<EntityGroupRoot> getCrossTldKey() {
-    return Key.create(EntityGroupRoot.class, "cross-tld");
+  public static @Nullable Key<EntityGroupRoot> getCrossTldKey() {
+    // If we cannot get a current environment, calling Key.create() will fail. Instead we return a
+    // null in cases where this key is not actually needed (for example when loading an entity from
+    // SQL) to initialize an object, to avoid having to register a DatastoreEntityExtension in
+    // tests.
+    return ApiProxy.getCurrentEnvironment() == null
+        ? null
+        : Key.create(EntityGroupRoot.class, "cross-tld");
   }
 }

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

@@ -15,7 +15,7 @@
 package google.registry.model.common;
 
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.api.users.User;

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

@@ -15,6 +15,7 @@
 package google.registry.model.domain;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -22,7 +23,6 @@ import com.googlecode.objectify.annotation.Embed;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.BillingVKey.BillingEventVKey;
 import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
 import google.registry.persistence.VKey;
@@ -68,7 +68,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
         (billingEventRecurring != null) == GracePeriodStatus.AUTO_RENEW.equals(type),
         "Recurring billing events must be present on (and only on) autorenew grace periods");
     GracePeriod instance = new GracePeriod();
-    instance.gracePeriodId = gracePeriodId == null ? ObjectifyService.allocateId() : gracePeriodId;
+    instance.gracePeriodId = gracePeriodId == null ? allocateId() : gracePeriodId;
     instance.type = checkArgumentNotNull(type);
     instance.domainRepoId = checkArgumentNotNull(domainRepoId);
     instance.expirationTime = checkArgumentNotNull(expirationTime);
@@ -210,7 +210,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
 
     static GracePeriodHistory createFrom(long historyRevisionId, GracePeriod gracePeriod) {
       GracePeriodHistory instance = new GracePeriodHistory();
-      instance.gracePeriodHistoryRevisionId = ObjectifyService.allocateId();
+      instance.gracePeriodHistoryRevisionId = allocateId();
       instance.domainHistoryRevisionId = historyRevisionId;
       instance.gracePeriodId = gracePeriod.gracePeriodId;
       instance.type = gracePeriod.type;

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

@@ -14,8 +14,9 @@
 
 package google.registry.model.domain.secdns;
 
+import static google.registry.model.IdService.allocateId;
+
 import google.registry.model.domain.DomainHistory;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.schema.replay.SqlOnlyEntity;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
@@ -48,7 +49,7 @@ public class DomainDsDataHistory extends DomainDsDataBase implements SqlOnlyEnti
     instance.algorithm = dsData.getAlgorithm();
     instance.digestType = dsData.getDigestType();
     instance.digest = dsData.getDigest();
-    instance.dsDataHistoryRevisionId = ObjectifyService.allocateId();
+    instance.dsDataHistoryRevisionId = allocateId();
     return instance;
   }
 

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

@@ -21,8 +21,6 @@ import static google.registry.util.TypeUtils.hasAnnotation;
 
 import com.google.appengine.api.datastore.AsyncDatastoreService;
 import com.google.appengine.api.datastore.DatastoreServiceConfig;
-import com.google.appengine.api.datastore.DatastoreServiceFactory;
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
@@ -48,7 +46,6 @@ import google.registry.model.translators.InetAddressTranslatorFactory;
 import google.registry.model.translators.ReadableInstantUtcTranslatorFactory;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
 import google.registry.model.translators.VKeyTranslatorFactory;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * An instance of Ofy, obtained via {@code #auditedOfy()}, should be used to access all persistable
@@ -57,12 +54,6 @@ import java.util.concurrent.atomic.AtomicLong;
  */
 public class ObjectifyService {
 
-  /**
-   * A placeholder String passed into DatastoreService.allocateIds that ensures that all ids are
-   * initialized from the same id pool.
-   */
-  public static final String APP_WIDE_ALLOCATION_KIND = "common";
-
   /** A singleton instance of our Ofy wrapper. */
   private static final Ofy OFY = new Ofy(null);
 
@@ -101,21 +92,24 @@ public class ObjectifyService {
   private static void initOfyOnce() {
     // Set an ObjectifyFactory that uses our extended ObjectifyImpl.
     // The "false" argument means that we are not using the v5-style Objectify embedded entities.
-    com.googlecode.objectify.ObjectifyService.setFactory(new ObjectifyFactory(false) {
+    com.googlecode.objectify.ObjectifyService.setFactory(
+        new ObjectifyFactory(false) {
           @Override
           public Objectify begin() {
             return new SessionKeyExposingObjectify(this);
           }
 
           @Override
-      protected AsyncDatastoreService createRawAsyncDatastoreService(DatastoreServiceConfig cfg) {
-        // In the unit test environment, wrap the Datastore service in a proxy that can be used to
-        // examine the number of requests sent to Datastore.
+          protected AsyncDatastoreService createRawAsyncDatastoreService(
+              DatastoreServiceConfig cfg) {
+            // In the unit test environment, wrap the Datastore service in a proxy that can be used
+            // to examine the number of requests sent to Datastore.
             AsyncDatastoreService service = super.createRawAsyncDatastoreService(cfg);
             return RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
                 ? new RequestCapturingAsyncDatastoreService(service)
                 : service;
-      }});
+          }
+        });
 
     // Translators must be registered before any entities can be registered.
     registerTranslators();
@@ -168,9 +162,11 @@ public class ObjectifyService {
       } else if (clazz.isAnnotationPresent(EntitySubclass.class)) {
         // Ensure that any @EntitySubclass classes have also had their parent @Entity registered,
         // which Objectify nominally requires but doesn't enforce in 4.x (though it may in 5.x).
-        checkState(registered,
+        checkState(
+            registered,
             "No base entity for kind '%s' registered yet, cannot register new @EntitySubclass %s",
-            kind, clazz.getCanonicalName());
+            kind,
+            clazz.getCanonicalName());
       }
       com.googlecode.objectify.ObjectifyService.register(clazz);
       // Autogenerated ids make the commit log code very difficult since we won't always be able
@@ -186,27 +182,4 @@ public class ObjectifyService {
       }
     }
   }
-
-  /** Counts of used ids for use in unit tests. Outside tests this is never used. */
-  private static final AtomicLong nextTestId = new AtomicLong(1); // ids cannot be zero
-
-  /** Allocates an id. */
-  public static long allocateId() {
-    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
-      ? nextTestId.getAndIncrement()
-      : DatastoreServiceFactory.getDatastoreService()
-          .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
-          .iterator()
-          .next()
-          .getId();
-  }
-
-  /** Resets the global test id counter (i.e. sets the next id to 1). */
-  @VisibleForTesting
-  public static void resetNextTestId() {
-    checkState(RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get()),
-        "Can't call resetTestIdCounts() from RegistryEnvironment.%s",
-        RegistryEnvironment.get());
-    nextTestId.set(1); // ids cannot be zero
-  }
 }

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

@@ -19,7 +19,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.hash.Funnels.stringFunnel;
 import static com.google.common.hash.Funnels.unencodedCharsFunnel;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
@@ -53,6 +53,7 @@ import javax.persistence.Column;
 import javax.persistence.Index;
 import javax.persistence.PostLoad;
 import javax.persistence.PostPersist;
+import javax.persistence.PostUpdate;
 import javax.persistence.PrePersist;
 import javax.persistence.PreRemove;
 import javax.persistence.Table;

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

@@ -17,7 +17,7 @@ package google.registry.model.tmch;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -36,6 +36,7 @@ import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.annotations.VirtualEntity;
 import google.registry.model.common.CrossTldSingleton;
+import google.registry.model.registry.label.ReservedList.ReservedListEntry;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.schema.replay.NonReplicatedEntity;
 import java.util.Map;
@@ -45,6 +46,7 @@ import javax.persistence.Column;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
 import javax.persistence.PostPersist;
+import javax.persistence.PostUpdate;
 import javax.persistence.PreRemove;
 import javax.persistence.Table;
 import javax.persistence.Transient;

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

@@ -45,7 +45,6 @@ import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
-import google.registry.testing.SystemPropertyExtension;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.DateTime;
@@ -65,14 +64,6 @@ public class RegistryJpaReadTest {
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
 
-  // The pipeline runner on Kokoro sometimes mistakes the platform as appengine, resulting in
-  // a null thread factory. The cause is unknown but it may be due to the interaction with
-  // the DatastoreEntityExtension above. To work around the problem, we explicitly unset the
-  // relevant property before test starts.
-  @RegisterExtension
-  final transient SystemPropertyExtension systemPropertyExtension =
-      new SystemPropertyExtension().setProperty("com.google.appengine.runtime.environment", null);
-
   @RegisterExtension
   final transient JpaIntegrationTestExtension database =
       new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -18,6 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -49,7 +50,6 @@ import google.registry.model.reporting.Spec11ThreatMatchDao;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
 import google.registry.persistence.transaction.TransactionManager;
-import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
@@ -71,7 +71,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 import org.json.JSONObject;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -113,8 +112,6 @@ class Spec11PipelineTest {
           ThreatMatch.create("THREAT_TYPE_UNSPECIFIED", "no-eamil.com"),
           ThreatMatch.create("UNWANTED_SOFTWARE", "anti-anti-anti-virus.dev"));
 
-  // This extension is only needed because Spec11ThreatMatch uses Ofy to generate the ID. Can be
-  // removed after the SQL migration.
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
@@ -136,12 +133,9 @@ class Spec11PipelineTest {
   private PCollection<KV<Subdomain, ThreatMatch>> threatMatches;
 
   ImmutableSet<Spec11ThreatMatch> sqlThreatMatches;
-  TransactionManager tm;
 
   @BeforeEach
   void beforeEach() throws Exception {
-    tm = tm();
-    TransactionManagerFactory.setTm(jpaTm());
     reportingBucketUrl = Files.createDirectory(tmpDir.resolve(REPORTING_BUCKET_URL)).toFile();
     options.setDate(DATE);
     options.setSafeBrowsingApiKey(SAFE_BROWSING_API_KEY);
@@ -196,11 +190,6 @@ class Spec11PipelineTest {
                 .build());
   }
 
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.setTm(tm);
-  }
-
   @Test
   void testSuccess_fullSqlPipeline() throws Exception {
     setupCloudSql();
@@ -241,6 +230,8 @@ class Spec11PipelineTest {
   }
 
   private void setupCloudSql() {
+    TransactionManager originalTm = tm();
+    setTm(jpaTm());
     persistNewRegistrar("TheRegistrar");
     persistNewRegistrar("NewRegistrar");
     Registrar registrar1 =
@@ -280,6 +271,7 @@ class Spec11PipelineTest {
     persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2));
     persistResource(
         createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3));
+    setTm(originalTm);
   }
 
   private void verifySaveToGcs() throws Exception {

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

@@ -44,6 +44,7 @@ import com.google.common.collect.Sets;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.ObjectifyFilter;
+import google.registry.model.IdService;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
@@ -466,7 +467,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (withDatastore) {
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
-      ObjectifyService.resetNextTestId();
+      IdService.resetNextTestId();
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }

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

@@ -28,10 +28,10 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -90,7 +90,6 @@ import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.EppResourceIndexBucket;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.pricing.StaticPremiumListPricingEngine;
 import google.registry.model.registrar.Registrar;
@@ -957,12 +956,12 @@ public class DatabaseHelper {
 
   /** Returns a newly allocated, globally unique domain repoId of the format HEX-TLD. */
   public static String generateNewDomainRoid(String tld) {
-    return createDomainRepoId(ObjectifyService.allocateId(), tld);
+    return createDomainRepoId(allocateId(), tld);
   }
 
   /** Returns a newly allocated, globally unique contact/host repoId of the format HEX_TLD-ROID. */
   public static String generateNewContactHostRoid() {
-    return createRepoId(ObjectifyService.allocateId(), getContactAndHostRoidSuffix());
+    return createRepoId(allocateId(), getContactAndHostRoidSuffix());
   }
 
   /**

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

@@ -14,6 +14,8 @@
 
 package google.registry.testing;
 
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+
 import com.google.apphosting.api.ApiProxy;
 import com.google.apphosting.api.ApiProxy.Environment;
 import java.util.Map;
@@ -24,7 +26,11 @@ import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 /**
  * Allows instantiation of Datastore {@code Entity entities} without the heavyweight {@code
- * AppEngineRule}.
+ * AppEngineExtension}.
+ *
+ * <p>When an Ofy key is created, by calling the various Key.create() methods, whether the current
+ * executing thread is a GAE thread is checked, which this extension masquerades. This happens
+ * frequently when an entity which has the key of another entity as a field is instantiated.
  *
  * <p>When used together with {@code JpaIntegrationWithCoverageExtension} or @{@code
  * TestPipelineExtension}, this extension must be registered first. For consistency's sake, it is
@@ -41,6 +47,10 @@ public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCa
   @Override
   public void beforeEach(ExtensionContext context) {
     ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+    // In order to create keys for entities they must be registered with Ofy. Calling this method
+    // will load the ObjectifyService class, whose static initialization block registers all Ofy
+    // entities.
+    auditedOfy();
   }
 
   @Override