Modify the way we load resources via foreign keys

core/src/main/java/google/registry/bsa/BsaValidateAction.java

@@ -28,7 +28,6 @@
 import static google.registry.bsa.persistence.Queries.queryUnblockableDomainByLabels;
 import static google.registry.model.tld.Tld.isEnrolledWithBsa;
 import static google.registry.model.tld.Tlds.getTldEntitiesOfType;
-import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.BatchedStreams.toBatches;
@@ -53,7 +52,6 @@
 import google.registry.model.domain.Domain;
 import google.registry.model.tld.Tld;
 import google.registry.model.tld.Tld.TldType;
-import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Action.GaeService;
 import google.registry.request.Response;
@@ -185,8 +183,8 @@ ImmutableList<String> checkWronglyReportedUnblockableDomains() {
     ImmutableList<UnblockableDomain> batch;
     do {
       batch = Queries.batchReadUnblockableDomains(lastRead, transactionBatchSize);
-      ImmutableMap<String, VKey<Domain>> activeDomains =
-          ForeignKeyUtils.load(
+      ImmutableMap<String, Domain> activeDomains =
+          ForeignKeyUtils.loadResources(
               Domain.class,
               batch.stream().map(UnblockableDomain::domainName).collect(toImmutableList()),
               clock.nowUtc());
@@ -201,7 +199,7 @@ ImmutableList<String> checkWronglyReportedUnblockableDomains() {
   }
 
   Optional<String> verifyDomainStillUnblockableWithReason(
-      UnblockableDomain domain, ImmutableMap<String, VKey<Domain>> activeDomains) {
+      UnblockableDomain domain, ImmutableMap<String, Domain> activeDomains) {
     DateTime now = clock.nowUtc();
     boolean isRegistered = activeDomains.containsKey(domain.domainName());
     boolean isReserved = isReservedDomain(domain.domainName(), now);
@@ -230,10 +228,8 @@ Optional<String> verifyDomainStillUnblockableWithReason(
             domain.reason()));
   }
 
-  boolean isStalenessAllowed(VKey<Domain> domainVKey) {
-    Domain domain = bsaQuery(() -> replicaTm().loadByKey(domainVKey));
-    var now = clock.nowUtc();
-    return domain.getCreationTime().plus(maxStaleness).isAfter(now);
+  boolean isStalenessAllowed(Domain domain) {
+    return domain.getCreationTime().plus(maxStaleness).isAfter(clock.nowUtc());
   }
 
   /** Returns unique labels across all block lists in the download specified by {@code jobName}. */

core/src/main/java/google/registry/bsa/persistence/DomainsRefresher.java

@@ -156,7 +156,7 @@ ImmutableSet<UnblockableDomainChange> recheckStaleDomainsBatch(
             .collect(toImmutableSet());
     ImmutableSet<String> currRegistered =
         ImmutableSet.copyOf(
-            ForeignKeyUtils.load(Domain.class, nameToEntity.keySet(), now).keySet());
+            ForeignKeyUtils.loadKeys(Domain.class, nameToEntity.keySet(), now).keySet());
     SetView<String> noLongerRegistered = Sets.difference(prevRegistered, currRegistered);
     SetView<String> newlyRegistered = Sets.difference(currRegistered, prevRegistered);
 

core/src/main/java/google/registry/bsa/persistence/LabelDiffUpdates.java

@@ -145,11 +145,10 @@ static ImmutableList<UnblockableDomain> tallyUnblockableDomainsForNewLabels(
 
     ImmutableSet<String> validDomainNames =
         labels.stream()
-            .map(label -> validDomainNamesForLabel(label, idnChecker))
-            .flatMap(x -> x)
+            .flatMap(label -> validDomainNamesForLabel(label, idnChecker))
             .collect(toImmutableSet());
     ImmutableSet<String> registeredDomainNames =
-        ImmutableSet.copyOf(ForeignKeyUtils.load(Domain.class, validDomainNames, now).keySet());
+        ForeignKeyUtils.loadKeys(Domain.class, validDomainNames, now).keySet();
     for (String domain : registeredDomainNames) {
       nonBlockedDomains.add(new UnblockableDomain(domain, Reason.REGISTERED));
       tm().put(BsaUnblockableDomain.of(domain, BsaUnblockableDomain.Reason.REGISTERED));

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

@@ -38,7 +38,6 @@
 import static google.registry.util.DomainNameUtils.canonicalizeHostname;
 import static org.json.simple.JSONValue.toJSONString;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
@@ -185,8 +184,7 @@ private Map<String, Object> checkDomainName(InternetDomainName domainName) {
   }
 
   private boolean checkExists(String domainString, DateTime now) {
-    return !ForeignKeyUtils.loadByCache(Domain.class, ImmutableList.of(domainString), now)
-        .isEmpty();
+    return ForeignKeyUtils.loadKeyByCache(Domain.class, domainString, now).isPresent();
   }
 
   private Optional<String> checkReserved(InternetDomainName domainName) {

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

@@ -71,10 +71,9 @@ public static void verifyResourceOwnership(String myRegistrarId, EppResource res
    */
   public static <R extends EppResource> void checkLinkedDomains(
       final String targetId, final DateTime now, final Class<R> resourceClass) throws EppException {
-    VKey<R> key = ForeignKeyUtils.load(resourceClass, targetId, now);
-    if (key == null) {
-      throw new ResourceDoesNotExistException(resourceClass, targetId);
-    }
+    VKey<R> key =
+        ForeignKeyUtils.loadKey(resourceClass, targetId, now)
+            .orElseThrow(() -> new ResourceDoesNotExistException(resourceClass, targetId));
     if (isLinked(key, now)) {
       throw new ResourceToDeleteIsReferencedException();
     }
@@ -106,11 +105,10 @@ public static <R extends EppResource> R verifyExistence(
 
   public static <R extends EppResource> void verifyResourceDoesNotExist(
       Class<R> clazz, String targetId, DateTime now, String registrarId) throws EppException {
-    VKey<R> key = ForeignKeyUtils.load(clazz, targetId, now);
-    if (key != null) {
-      R resource = tm().loadByKey(key);
+    Optional<R> resource = ForeignKeyUtils.loadResource(clazz, targetId, now);
+    if (resource.isPresent()) {
       // These are similar exceptions, but we can track them internally as log-based metrics.
-      if (Objects.equals(registrarId, resource.getPersistedCurrentSponsorRegistrarId())) {
+      if (Objects.equals(registrarId, resource.get().getPersistedCurrentSponsorRegistrarId())) {
         throw new ResourceAlreadyExistsForThisClientException(targetId);
       } else {
         throw new ResourceCreateContentionException(targetId);

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

@@ -16,7 +16,6 @@
 
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.verifyTargetIdCount;
-import static google.registry.model.EppResourceUtils.checkResourcesExist;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -26,6 +25,7 @@
 import google.registry.flows.FlowModule.RegistrarId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.contact.Contact;
 import google.registry.model.contact.ContactCommand.Check;
 import google.registry.model.eppinput.ResourceCommand;
@@ -62,7 +62,7 @@ public EppResponse run() throws EppException {
     ImmutableList<String> targetIds = ((Check) resourceCommand).getTargetIds();
     verifyTargetIdCount(targetIds, maxChecks);
     ImmutableSet<String> existingIds =
-        checkResourcesExist(Contact.class, targetIds, clock.nowUtc());
+        ForeignKeyUtils.loadKeys(Contact.class, targetIds, clock.nowUtc()).keySet();
     ImmutableList.Builder<ContactCheck> checks = new ImmutableList.Builder<>();
     for (String id : targetIds) {
       boolean unused = !existingIds.contains(id);

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

@@ -181,7 +181,7 @@ public EppResponse run() throws EppException {
             .setAsOfDate(now)
             .build());
     ImmutableMap<String, VKey<Domain>> existingDomains =
-        ForeignKeyUtils.load(Domain.class, domainNames, now);
+        ForeignKeyUtils.loadKeys(Domain.class, domainNames, now);
     // Check block labels only when there are unregistered domains, since "In use" goes before
     // "Blocked by BSA".
     ImmutableSet<InternetDomainName> bsaBlockedDomainNames =

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

@@ -15,7 +15,7 @@
 package google.registry.flows.domain;
 
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
-import static google.registry.flows.ResourceFlowUtils.verifyExistence;
+import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.domain.DomainFlowUtils.addSecDnsExtensionIfPresent;
 import static google.registry.flows.domain.DomainFlowUtils.handleFeeRequest;
@@ -37,7 +37,6 @@
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.AfterValidationParameters;
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.BeforeResponseParameters;
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.BeforeResponseReturnData;
-import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.domain.DomainCommand.Info;
 import google.registry.model.domain.DomainCommand.Info.HostsRequest;
@@ -108,9 +107,7 @@ public EppResponse run() throws EppException {
     validateRegistrarIsLoggedIn(registrarId);
     extensionManager.validate();
     DateTime now = clock.nowUtc();
-    Domain domain =
-        verifyExistence(
-            Domain.class, targetId, ForeignKeyUtils.loadResource(Domain.class, targetId, now));
+    Domain domain = loadAndVerifyExistence(Domain.class, targetId, now);
     verifyOptionalAuthInfo(authInfo, domain);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());

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

@@ -16,7 +16,6 @@
 
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.verifyTargetIdCount;
-import static google.registry.model.EppResourceUtils.checkResourcesExist;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -26,6 +25,7 @@
 import google.registry.flows.FlowModule.RegistrarId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.CheckData.HostCheck;
 import google.registry.model.eppoutput.CheckData.HostCheckData;
@@ -61,7 +61,8 @@ public EppResponse run() throws EppException {
     extensionManager.validate(); // There are no legal extensions for this flow.
     ImmutableList<String> hostnames = ((Check) resourceCommand).getTargetIds();
     verifyTargetIdCount(hostnames, maxChecks);
-    ImmutableSet<String> existingIds = checkResourcesExist(Host.class, hostnames, clock.nowUtc());
+    ImmutableSet<String> existingIds =
+        ForeignKeyUtils.loadKeys(Host.class, hostnames, clock.nowUtc()).keySet();
     ImmutableList.Builder<HostCheck> checks = new ImmutableList.Builder<>();
     for (String hostname : hostnames) {
       boolean unused = !existingIds.contains(hostname);

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

@@ -154,7 +154,7 @@ public EppResponse run() throws EppException {
     EppResource owningResource = firstNonNull(oldSuperordinateDomain, existingHost);
     verifyUpdateAllowed(
         command, existingHost, newSuperordinateDomain.orElse(null), owningResource, isHostRename);
-    if (isHostRename && ForeignKeyUtils.load(Host.class, newHostName, now) != null) {
+    if (isHostRename && ForeignKeyUtils.loadKey(Host.class, newHostName, now).isPresent()) {
       throw new HostAlreadyExistsException(newHostName);
     }
     AddRemove add = command.getInnerAdd();

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

@@ -16,19 +16,14 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
-import static google.registry.util.DateTimeUtils.latestOf;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryConfig;
 import google.registry.model.EppResource.BuilderWithTransferData;
-import google.registry.model.EppResource.ForeignKeyedEppResource;
 import google.registry.model.EppResource.ResourceWithTransferData;
 import google.registry.model.contact.Contact;
 import google.registry.model.domain.Domain;
@@ -41,13 +36,9 @@
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
-import google.registry.persistence.transaction.TransactionManager;
 import jakarta.persistence.Query;
-import java.util.Collection;
 import java.util.Comparator;
-import java.util.Optional;
 import java.util.function.Function;
-import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
@@ -90,98 +81,6 @@ private static <T extends EppResource> T cloneProjectedAtTime(T resource, DateTi
     return (T) resource.cloneProjectedAtTime(now);
   }
 
-  /**
-   * Loads the last created version of an {@link EppResource} from the database by foreign key,
-   * using a cache, if caching is enabled in config settings.
-   *
-   * <p>Returns null if no resource with this foreign key was ever created, or if the most recently
-   * created resource was deleted before time "now".
-   *
-   * <p>Loading an {@link EppResource} by itself is not sufficient to know its current state since
-   * it may have various expirable conditions and status values that might implicitly change its
-   * state as time progresses even if it has not been updated in the database. Rather, the resource
-   * must be combined with a timestamp to view its current state. We use a global last updated
-   * timestamp to guarantee monotonically increasing write times, and forward our projected time to
-   * the greater of this timestamp or "now". This guarantees that we're not projecting into the
-   * past.
-   *
-   * <p>Do not call this cached version for anything that needs transactional consistency. It should
-   * only be used when it's OK if the data is potentially being out of date, e.g. RDAP.
-   *
-   * @param foreignKey id to match
-   * @param now the current logical time to project resources at
-   */
-  public static <T extends EppResource> Optional<T> loadByForeignKeyByCacheIfEnabled(
-      Class<T> clazz, String foreignKey, DateTime now) {
-    return loadByForeignKeyHelper(
-        tm(), clazz, foreignKey, now, RegistryConfig.isEppResourceCachingEnabled());
-  }
-
-  /**
-   * Loads the last created version of an {@link EppResource} from the replica database by foreign
-   * key, using a cache.
-   *
-   * <p>This method ignores the config setting for caching, and is reserved for use cases that can
-   * tolerate slightly stale data.
-   */
-  public static <T extends EppResource> Optional<T> loadByForeignKeyByCache(
-      Class<T> clazz, String foreignKey, DateTime now) {
-    return loadByForeignKeyHelper(replicaTm(), clazz, foreignKey, now, true);
-  }
-
-  static <T extends EppResource> Optional<T> loadByForeignKeyHelper(
-      TransactionManager txnManager,
-      Class<T> clazz,
-      String foreignKey,
-      DateTime now,
-      boolean useCache) {
-    checkArgument(
-        ForeignKeyedEppResource.class.isAssignableFrom(clazz),
-        "loadByForeignKey may only be called for foreign keyed EPP resources");
-    VKey<T> key =
-        useCache
-            ? ForeignKeyUtils.loadByCache(clazz, ImmutableList.of(foreignKey), now).get(foreignKey)
-            : ForeignKeyUtils.load(clazz, foreignKey, now);
-    // The returned key is null if the resource is hard deleted or soft deleted by the given time.
-    if (key == null) {
-      return Optional.empty();
-    }
-    T resource =
-        useCache
-            ? EppResource.loadByCache(key)
-            // This transaction is buried very deeply inside many outer nested calls, hence merits
-            // the use of reTransact() for now pending a substantial refactoring.
-            : txnManager.reTransact(() -> txnManager.loadByKeyIfPresent(key).orElse(null));
-    if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
-      return Optional.empty();
-    }
-    // When setting status values based on a time, choose the greater of "now" and the resource's
-    // UpdateAutoTimestamp. For non-mutating uses (info, RDAP, etc.), this is equivalent to rolling
-    // "now" forward to at least the last update on the resource, so that a read right after a write
-    // doesn't appear stale. For mutating flows, if we had to roll now forward then the flow will
-    // fail when it tries to save anything, since "now" is needed to be > the last update time for
-    // writes.
-    return Optional.of(
-        cloneProjectedAtTime(
-            resource, latestOf(now, resource.getUpdateTimestamp().getTimestamp())));
-  }
-
-  /**
-   * Checks multiple {@link EppResource} objects from the database by unique ids.
-   *
-   * <p>There are currently no resources that support checks and do not use foreign keys. If we need
-   * to support that case in the future, we can loosen the type to allow any {@link EppResource} and
-   * add code to do the lookup by id directly.
-   *
-   * @param clazz the resource type to load
-   * @param uniqueIds a list of ids to match
-   * @param now the logical time of the check
-   */
-  public static <T extends EppResource> ImmutableSet<String> checkResourcesExist(
-      Class<T> clazz, Collection<String> uniqueIds, final DateTime now) {
-    return ForeignKeyUtils.load(clazz, uniqueIds, now).keySet();
-  }
-
   /**
    * Returns a Function that transforms an EppResource to the given DateTime, suitable for use with
    * Iterables.transform() over a collection of EppResources.
@@ -281,21 +180,6 @@ public static <T extends EppResource> T loadAtPointInTime(
             : null);
   }
 
-  /**
-   * Rewinds an {@link EppResource} object to a given point in time.
-   *
-   * <p>This method costs nothing if {@code resource} is already current. Otherwise, it returns an
-   * async operation that performs a single fetch operation.
-   *
-   * @return an asynchronous operation returning resource at {@code timestamp} or {@code null} if
-   *     resource is deleted or not yet created
-   * @see #loadAtPointInTime(EppResource, DateTime)
-   */
-  public static <T extends EppResource> Supplier<T> loadAtPointInTimeAsync(
-      final T resource, final DateTime timestamp) {
-    return () -> loadAtPointInTime(resource, timestamp);
-  }
-
   /**
    * Returns the most recent revision of a given EppResource before or at the provided timestamp,
    * falling back to using the resource as-is if there are no revisions.
@@ -370,13 +254,11 @@ public static ImmutableSet<VKey<Domain>> getLinkedDomainKeys(
   /**
    * Returns whether the given contact or host is linked to (that is, referenced by) a domain.
    *
-   * <p>This is an eventually consistent query.
-   *
    * @param key the referent key
    * @param now the logical time of the check
    */
   public static boolean isLinked(VKey<? extends EppResource> key, DateTime now) {
-    return getLinkedDomainKeys(key, now, 1).size() > 0;
+    return !getLinkedDomainKeys(key, now, 1).isEmpty();
   }
 
   private EppResourceUtils() {}