Load referenced contact/hosts from EPP more efficiently

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=122459862

java/google/registry/flows/EppException.java

@@ -17,7 +17,9 @@ package google.registry.flows;
 import static java.lang.annotation.ElementType.TYPE;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 
+import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSet;
 
 import google.registry.model.annotations.ExternalMessagingName;
 import google.registry.model.eppinput.EppInput.InnerCommand;
@@ -117,11 +119,21 @@ public abstract class EppException extends Exception {
       super(
           String.format(
               "The %s with given ID (%s) doesn't exist.",
-              !type.isAnnotationPresent(ExternalMessagingName.class)
+              type.isAnnotationPresent(ExternalMessagingName.class)
-                  ? "object"
+                  ? type.getAnnotation(ExternalMessagingName.class).value()
-                  : type.getAnnotation(ExternalMessagingName.class).value(),
+                  : "object",
               id));
     }
+
+    public ObjectDoesNotExistException(Class<?> type, ImmutableSet<String> ids) {
+      super(
+          String.format(
+              "The %s with given IDs (%s) don't exist.",
+              type.isAnnotationPresent(ExternalMessagingName.class)
+                  ? type.getAnnotation(ExternalMessagingName.class).value() + " objects"
+                  : "objects",
+              Joiner.on(',').join(ids)));
+    }
   }
 
   /** Abstract exception class. Do not throw this directly or catch in tests. */

java/google/registry/flows/domain/DomainApplicationCreateFlow.java

@@ -89,7 +89,7 @@ import java.util.List;
  * @error {@link DomainFlowUtils.InvalidPunycodeException}
  * @error {@link DomainFlowUtils.LaunchPhaseMismatchException}
  * @error {@link DomainFlowUtils.LeadingDashException}
- * @error {@link DomainFlowUtils.LinkedResourceDoesNotExistException}
+ * @error {@link DomainFlowUtils.LinkedResourcesDoNotExistException}
  * @error {@link DomainFlowUtils.MissingContactTypeException}
  * @error {@link DomainFlowUtils.NameserversNotAllowedException}
  * @error {@link DomainFlowUtils.NoMarksFoundMatchingDomainException}

java/google/registry/flows/domain/DomainApplicationUpdateFlow.java

@@ -46,7 +46,7 @@ import google.registry.model.reporting.HistoryEntry;
  * @error {@link BaseDomainUpdateFlow.SecDnsAllUsageException}
  * @error {@link BaseDomainUpdateFlow.UrgentAttributeNotSupportedException}
  * @error {@link DomainFlowUtils.DuplicateContactForRoleException}
- * @error {@link DomainFlowUtils.LinkedResourceDoesNotExistException}
+ * @error {@link DomainFlowUtils.LinkedResourcesDoNotExistException}
  * @error {@link DomainFlowUtils.MissingAdminContactException}
  * @error {@link DomainFlowUtils.MissingContactTypeException}
  * @error {@link DomainFlowUtils.MissingTechnicalContactException}

java/google/registry/flows/domain/DomainCreateFlow.java

@@ -75,7 +75,7 @@ import java.util.Set;
  * @error {@link DomainFlowUtils.InvalidIdnDomainLabelException}
  * @error {@link DomainFlowUtils.InvalidPunycodeException}
  * @error {@link DomainFlowUtils.LeadingDashException}
- * @error {@link DomainFlowUtils.LinkedResourceDoesNotExistException}
+ * @error {@link DomainFlowUtils.LinkedResourcesDoNotExistException}
  * @error {@link DomainFlowUtils.LinkedResourceInPendingDeleteProhibitsOperationException}
  * @error {@link DomainFlowUtils.MissingAdminContactException}
  * @error {@link DomainFlowUtils.MissingContactTypeException}

java/google/registry/flows/domain/DomainFlowUtils.java

@@ -64,7 +64,7 @@ import google.registry.model.domain.DesignatedContact;
 import google.registry.model.domain.DesignatedContact.Type;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.CreateOrUpdate;
-import google.registry.model.domain.DomainCommand.InvalidReferenceException;
+import google.registry.model.domain.DomainCommand.InvalidReferencesException;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.BaseFeeCommand;
@@ -405,8 +405,8 @@ public class DomainFlowUtils {
       throws EppException {
     try {
       return command.cloneAndLinkReferences(now);
-    } catch (InvalidReferenceException e) {
+    } catch (InvalidReferencesException e) {
-      throw new LinkedResourceDoesNotExistException(e.getType(), e.getForeignKey());
+      throw new LinkedResourcesDoNotExistException(e.getType(), e.getForeignKeys());
     }
   }
 
@@ -681,9 +681,9 @@ public class DomainFlowUtils {
   }
 
   /** Resource linked to this domain does not exist. */
-  static class LinkedResourceDoesNotExistException extends ObjectDoesNotExistException {
+  static class LinkedResourcesDoNotExistException extends ObjectDoesNotExistException {
-    public LinkedResourceDoesNotExistException(Class<?> type, String resourceId) {
+    public LinkedResourcesDoNotExistException(Class<?> type, ImmutableSet<String> resourceIds) {
-      super(type, resourceId);
+      super(type, resourceIds);
     }
   }
 

java/google/registry/flows/domain/DomainUpdateFlow.java

@@ -56,7 +56,7 @@ import java.util.Set;
  * @error {@link BaseDomainUpdateFlow.SecDnsAllUsageException}
  * @error {@link BaseDomainUpdateFlow.UrgentAttributeNotSupportedException}
  * @error {@link DomainFlowUtils.DuplicateContactForRoleException}
- * @error {@link DomainFlowUtils.LinkedResourceDoesNotExistException}
+ * @error {@link DomainFlowUtils.LinkedResourcesDoNotExistException}
  * @error {@link DomainFlowUtils.LinkedResourceInPendingDeleteProhibitsOperationException}
  * @error {@link DomainFlowUtils.MissingAdminContactException}
  * @error {@link DomainFlowUtils.MissingContactTypeException}

java/google/registry/model/domain/DomainCommand.java

@@ -16,13 +16,19 @@ package google.registry.model.domain;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.common.collect.Maps.transformValues;
+import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
-import static google.registry.model.index.ForeignKeyIndex.loadAndGetReference;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.util.CollectionUtils.difference;
 import static google.registry.util.CollectionUtils.nullSafeImmutableCopy;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
+import static google.registry.util.CollectionUtils.union;
 
+import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 
 import com.googlecode.objectify.Ref;
@@ -38,10 +44,12 @@ import google.registry.model.eppinput.ResourceCommand.ResourceCreateOrChange;
 import google.registry.model.eppinput.ResourceCommand.ResourceUpdate;
 import google.registry.model.eppinput.ResourceCommand.SingleResourceCommand;
 import google.registry.model.host.HostResource;
+import google.registry.model.index.ForeignKeyIndex;
 
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 
+import java.util.Map;
 import java.util.Set;
 
 import javax.xml.bind.annotation.XmlAttribute;
@@ -66,7 +74,7 @@ public class DomainCommand {
    */
   public interface CreateOrUpdate<T extends CreateOrUpdate<T>> extends SingleResourceCommand {
     /** Creates a copy of this command with hard links to hosts and contacts. */
-    public T cloneAndLinkReferences(DateTime now) throws InvalidReferenceException;
+    public T cloneAndLinkReferences(DateTime now) throws InvalidReferencesException;
   }
 
   /** The fields on "chgType" from {@link "http://tools.ietf.org/html/rfc5731"}. */
@@ -190,12 +198,27 @@ public class DomainCommand {
 
     /** Creates a copy of this {@link Create} with hard links to hosts and contacts. */
     @Override
-    public Create cloneAndLinkReferences(DateTime now) throws InvalidReferenceException {
+    public Create cloneAndLinkReferences(DateTime now) throws InvalidReferencesException {
       Create clone = clone(this);
       clone.nameservers = linkHosts(clone.nameserverFullyQualifiedHostNames, now);
+      if (registrantContactId == null) {
         clone.contacts = linkContacts(clone.foreignKeyedDesignatedContacts, now);
-      clone.registrant = clone.registrantContactId == null
+      } else {
-          ? null : loadReference(clone.registrantContactId, ContactResource.class, now);
+        // Load the registrant and contacts in one shot.
+        ForeignKeyedDesignatedContact registrantPlaceholder = new ForeignKeyedDesignatedContact();
+        registrantPlaceholder.contactId = clone.registrantContactId;
+        registrantPlaceholder.type = DesignatedContact.Type.REGISTRANT;
+        Set<DesignatedContact> contacts = linkContacts(
+            union(clone.foreignKeyedDesignatedContacts, registrantPlaceholder),
+            now);
+        for (DesignatedContact contact : contacts) {
+          if (DesignatedContact.Type.REGISTRANT.equals(contact.getType())) {
+            clone.registrant = contact.getContactRef();
+            clone.contacts = difference(contacts, contact);
+            break;
+          }
+        }
+      }
       return clone;
     }
   }
@@ -378,7 +401,7 @@ public class DomainCommand {
       }
 
       /** Creates a copy of this {@link AddRemove} with hard links to hosts and contacts. */
-      private AddRemove cloneAndLinkReferences(DateTime now) throws InvalidReferenceException {
+      private AddRemove cloneAndLinkReferences(DateTime now) throws InvalidReferencesException {
         AddRemove clone = clone(this);
         clone.nameservers = linkHosts(clone.nameserverFullyQualifiedHostNames, now);
         clone.contacts = linkContacts(clone.foreignKeyedDesignatedContacts, now);
@@ -390,10 +413,14 @@ public class DomainCommand {
     @XmlType(propOrder = {"registrantContactId", "authInfo"})
     public static class Change extends DomainCreateOrChange<DomainBase.Builder<?, ?>> {
       /** Creates a copy of this {@link Change} with hard links to hosts and contacts. */
-      Change cloneAndLinkReferences(DateTime now) throws InvalidReferenceException {
+      Change cloneAndLinkReferences(DateTime now) throws InvalidReferencesException {
         Change clone = clone(this);
         clone.registrant = clone.registrantContactId == null
-            ? null : loadReference(clone.registrantContactId, ContactResource.class, now);
+            ? null
+            : getOnlyElement(
+                loadReferences(
+                    ImmutableSet.of(clone.registrantContactId), ContactResource.class, now)
+                        .values());
         return clone;
       }
     }
@@ -423,7 +450,7 @@ public class DomainCommand {
      * those classes, which is harmless because the getters do that anyways.
      */
     @Override
-    public Update cloneAndLinkReferences(DateTime now) throws InvalidReferenceException {
+    public Update cloneAndLinkReferences(DateTime now) throws InvalidReferencesException {
       Update clone = clone(this);
       clone.innerAdd = clone.getInnerAdd().cloneAndLinkReferences(now);
       clone.innerRemove = clone.getInnerRemove().cloneAndLinkReferences(now);
@@ -433,44 +460,54 @@ public class DomainCommand {
   }
 
   private static Set<Ref<HostResource>> linkHosts(
-      Set<String> fullyQualifiedHostNames, DateTime now) throws InvalidReferenceException {
+      Set<String> fullyQualifiedHostNames, DateTime now) throws InvalidReferencesException {
     if (fullyQualifiedHostNames == null) {
       return null;
     }
-    ImmutableSet.Builder<Ref<HostResource>> linked = new ImmutableSet.Builder<>();
+    return ImmutableSet.copyOf(
-    for (String fullyQualifiedHostName : fullyQualifiedHostNames) {
+        loadReferences(fullyQualifiedHostNames, HostResource.class, now).values());
-      linked.add(loadReference(fullyQualifiedHostName, HostResource.class, now));
-    }
-    return linked.build();
   }
 
   private static Set<DesignatedContact> linkContacts(
-      Set<ForeignKeyedDesignatedContact> contacts, DateTime now) throws InvalidReferenceException {
+      Set<ForeignKeyedDesignatedContact> contacts, DateTime now) throws InvalidReferencesException {
     if (contacts == null) {
       return null;
     }
+    ImmutableSet.Builder<String> foreignKeys = new ImmutableSet.Builder<>();
+    for (ForeignKeyedDesignatedContact contact : contacts) {
+      foreignKeys.add(contact.contactId);
+    }
+    ImmutableMap<String, Ref<ContactResource>> loadedContacts =
+        loadReferences(foreignKeys.build(), ContactResource.class, now);
     ImmutableSet.Builder<DesignatedContact> linkedContacts = new ImmutableSet.Builder<>();
     for (ForeignKeyedDesignatedContact contact : contacts) {
       linkedContacts.add(DesignatedContact.create(
-          contact.type, loadReference(contact.contactId, ContactResource.class, now)));
+          contact.type, loadedContacts.get(contact.contactId)));
     }
     return linkedContacts.build();
   }
 
-  /** Load a reference to a resource by its foreign key. */
+  /** Load references to resources by their foreign keys. */
-  private static <T extends EppResource> Ref<T> loadReference(
+  private static <T extends EppResource> ImmutableMap<String, Ref<T>> loadReferences(
-      final String foreignKey, final Class<T> clazz, final DateTime now)
+      final Set<String> foreignKeys, final Class<T> clazz, final DateTime now)
-      throws InvalidReferenceException {
+      throws InvalidReferencesException {
-    Ref<T> ref = ofy().doTransactionless(new Work<Ref<T>>() {
+    Map<String, ForeignKeyIndex<T>> fkis = ofy().doTransactionless(
+        new Work<Map<String, ForeignKeyIndex<T>>>() {
           @Override
-      public Ref<T> run() {
+          public Map<String, ForeignKeyIndex<T>> run() {
-        return loadAndGetReference(clazz, foreignKey, now);
+            return ForeignKeyIndex.load(clazz, foreignKeys, now);
+          }});
+    if (!fkis.keySet().equals(foreignKeys)) {
+      throw new InvalidReferencesException(
+          clazz, ImmutableSet.copyOf(difference(foreignKeys, fkis.keySet())));
     }
-    });
+    return ImmutableMap.copyOf(transformValues(
-    if (ref == null) {
+        fkis,
-      throw new InvalidReferenceException(clazz, foreignKey);
+        new Function<ForeignKeyIndex<T>, Ref<T>>() {
-    }
+          @Override
-    return ref;
+          public Ref<T> apply(ForeignKeyIndex<T> fki) {
+            return fki.getReference();
+          }}));
   }
 
   /**
@@ -487,18 +524,18 @@ public class DomainCommand {
     String contactId;
   }
 
-  /** Exception to throw when a referenced object does not exist. */
+  /** Exception to throw when referenced objects don't exist. */
-  public static class InvalidReferenceException extends Exception {
+  public static class InvalidReferencesException extends Exception {
-    private String foreignKey;
+    private final ImmutableSet<String> foreignKeys;
-    private Class<?> type;
+    private final Class<?> type;
 
-    InvalidReferenceException(Class<?> type, String foreignKey) {
+    InvalidReferencesException(Class<?> type, ImmutableSet<String> foreignKeys) {
       this.type = checkNotNull(type);
-      this.foreignKey = foreignKey;
+      this.foreignKeys = foreignKeys;
     }
 
-    public String getForeignKey() {
+    public ImmutableSet<String> getForeignKeys() {
-      return foreignKey;
+      return foreignKeys;
     }
 
     public Class<?> getType() {

javatests/google/registry/flows/domain/DomainApplicationCreateFlowTest.java

@@ -77,7 +77,7 @@ import google.registry.flows.domain.DomainFlowUtils.InvalidIdnDomainLabelExcepti
 import google.registry.flows.domain.DomainFlowUtils.InvalidPunycodeException;
 import google.registry.flows.domain.DomainFlowUtils.LaunchPhaseMismatchException;
 import google.registry.flows.domain.DomainFlowUtils.LeadingDashException;
-import google.registry.flows.domain.DomainFlowUtils.LinkedResourceDoesNotExistException;
+import google.registry.flows.domain.DomainFlowUtils.LinkedResourcesDoNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.MissingContactTypeException;
 import google.registry.flows.domain.DomainFlowUtils.NameserversNotAllowedException;
 import google.registry.flows.domain.DomainFlowUtils.NoMarksFoundMatchingDomainException;
@@ -760,7 +760,7 @@ public class DomainApplicationCreateFlowTest
 
   @Test
   public void testFailure_missingHost() throws Exception {
-    thrown.expect(LinkedResourceDoesNotExistException.class, "(ns2.example.net)");
+    thrown.expect(LinkedResourcesDoNotExistException.class, "(ns2.example.net)");
     persistActiveHost("ns1.example.net");
     persistActiveContact("jd1234");
     persistActiveContact("sh8013");
@@ -769,7 +769,7 @@ public class DomainApplicationCreateFlowTest
 
   @Test
   public void testFailure_missingContact() throws Exception {
-    thrown.expect(LinkedResourceDoesNotExistException.class, "(sh8013)");
+    thrown.expect(LinkedResourcesDoNotExistException.class, "(sh8013)");
     persistActiveHost("ns1.example.net");
     persistActiveHost("ns2.example.net");
     persistActiveContact("jd1234");

javatests/google/registry/flows/domain/DomainApplicationUpdateFlowTest.java

@@ -48,7 +48,7 @@ import google.registry.flows.domain.BaseDomainUpdateFlow.SecDnsAllUsageException
 import google.registry.flows.domain.BaseDomainUpdateFlow.UrgentAttributeNotSupportedException;
 import google.registry.flows.domain.DomainApplicationUpdateFlow.ApplicationStatusProhibitsUpdateException;
 import google.registry.flows.domain.DomainFlowUtils.DuplicateContactForRoleException;
-import google.registry.flows.domain.DomainFlowUtils.LinkedResourceDoesNotExistException;
+import google.registry.flows.domain.DomainFlowUtils.LinkedResourcesDoNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.MissingAdminContactException;
 import google.registry.flows.domain.DomainFlowUtils.MissingContactTypeException;
 import google.registry.flows.domain.DomainFlowUtils.MissingTechnicalContactException;
@@ -467,7 +467,7 @@ public class DomainApplicationUpdateFlowTest
   @Test
   public void testFailure_missingHost() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(ns2.example.tld)");
     persistActiveHost("ns1.example.tld");
     persistActiveContact("sh8013");
@@ -479,7 +479,7 @@ public class DomainApplicationUpdateFlowTest
   @Test
   public void testFailure_missingContact() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(sh8013)");
     persistActiveHost("ns1.example.tld");
     persistActiveHost("ns2.example.tld");

javatests/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -83,8 +83,8 @@ import google.registry.flows.domain.DomainFlowUtils.FeesRequiredForPremiumNameEx
 import google.registry.flows.domain.DomainFlowUtils.InvalidIdnDomainLabelException;
 import google.registry.flows.domain.DomainFlowUtils.InvalidPunycodeException;
 import google.registry.flows.domain.DomainFlowUtils.LeadingDashException;
-import google.registry.flows.domain.DomainFlowUtils.LinkedResourceDoesNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.LinkedResourceInPendingDeleteProhibitsOperationException;
+import google.registry.flows.domain.DomainFlowUtils.LinkedResourcesDoNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.MissingAdminContactException;
 import google.registry.flows.domain.DomainFlowUtils.MissingContactTypeException;
 import google.registry.flows.domain.DomainFlowUtils.MissingRegistrantException;
@@ -668,7 +668,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
   @Test
   public void testFailure_missingHost() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(ns2.example.net)");
     persistActiveHost("ns1.example.net");
     persistActiveContact("jd1234");
@@ -720,7 +720,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
   @Test
   public void testFailure_missingContact() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(sh8013)");
     persistActiveHost("ns1.example.net");
     persistActiveHost("ns2.example.net");

javatests/google/registry/flows/domain/DomainUpdateFlowTest.java

@@ -57,8 +57,8 @@ import google.registry.flows.domain.BaseDomainUpdateFlow.MaxSigLifeChangeNotSupp
 import google.registry.flows.domain.BaseDomainUpdateFlow.SecDnsAllUsageException;
 import google.registry.flows.domain.BaseDomainUpdateFlow.UrgentAttributeNotSupportedException;
 import google.registry.flows.domain.DomainFlowUtils.DuplicateContactForRoleException;
-import google.registry.flows.domain.DomainFlowUtils.LinkedResourceDoesNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.LinkedResourceInPendingDeleteProhibitsOperationException;
+import google.registry.flows.domain.DomainFlowUtils.LinkedResourcesDoNotExistException;
 import google.registry.flows.domain.DomainFlowUtils.MissingAdminContactException;
 import google.registry.flows.domain.DomainFlowUtils.MissingContactTypeException;
 import google.registry.flows.domain.DomainFlowUtils.MissingTechnicalContactException;
@@ -858,7 +858,7 @@ public class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow,
   @Test
   public void testFailure_missingHost() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(ns2.example.foo)");
     persistActiveHost("ns1.example.foo");
     persistActiveContact("sh8013");
@@ -870,7 +870,7 @@ public class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow,
   @Test
   public void testFailure_missingContact() throws Exception {
     thrown.expect(
-        LinkedResourceDoesNotExistException.class,
+        LinkedResourcesDoNotExistException.class,
         "(sh8013)");
     persistActiveHost("ns1.example.foo");
     persistActiveHost("ns2.example.foo");