Flatten and inject the poll flows

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

java/google/registry/flows/Flow.java

@@ -24,6 +24,7 @@ import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.eppoutput.EppResponse.ResponseData;
 import google.registry.model.eppoutput.EppResponse.ResponseExtension;
 import google.registry.model.eppoutput.Result;
+import google.registry.model.poll.MessageQueueInfo;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -82,13 +83,22 @@ public abstract class Flow {
       Result.Code code,
       @Nullable ResponseData responseData,
       @Nullable ImmutableList<? extends ResponseExtension> extensions) {
+    return createOutput(
+        code, responseData == null ? null : ImmutableList.of(responseData), extensions, null);
+  }
+
+  protected EppOutput createOutput(
+      Result.Code code,
+      @Nullable ImmutableList<ResponseData> responseData,
+      @Nullable ImmutableList<? extends ResponseExtension> responseExtensions,
+      @Nullable MessageQueueInfo messageQueueInfo) {
     return EppOutput.create(new EppResponse.Builder()
         .setTrid(trid)
         .setResult(Result.create(code))
+        .setMessageQueueInfo(messageQueueInfo)
+        .setResData(responseData)
+        .setExtensions(responseExtensions)
         .setExecutionTime(now)
-        .setCreatedRepoId(getCreatedRepoId())
-        .setResData(responseData == null ? null : ImmutableList.of(responseData))
-        .setExtensions(extensions)
         .build());
   }
 

java/google/registry/flows/FlowModule.java

@@ -26,6 +26,7 @@ import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.eppinput.EppInput;
+import google.registry.model.eppinput.EppInput.Poll;
 import google.registry.model.eppinput.EppInput.ResourceCommandWrapper;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppinput.ResourceCommand.SingleResourceCommand;
@@ -193,6 +194,13 @@ public class FlowModule {
     return ((SingleResourceCommand) resourceCommand).getTargetId();
   }
 
+  @Provides
+  @FlowScope
+  @PollMessageId
+  static String providePollMessageId(EppInput eppInput) {
+    return Strings.nullToEmpty(((Poll) eppInput.getCommandWrapper().getCommand()).getMessageId());
+  }
+
   /**
    * Provides a partially filled in {@link HistoryEntry} builder.
    *
@@ -246,6 +254,11 @@ public class FlowModule {
   @Documented
   public @interface TargetId {}
 
+  /** Dagger qualifier for the message id for poll flows. */
+  @Qualifier
+  @Documented
+  public @interface PollMessageId {}
+
   /** Dagger qualifier for whether a flow is in dry run mode. */
   @Qualifier
   @Documented

java/google/registry/flows/poll/PollAckFlow.java

@@ -15,17 +15,22 @@
 package google.registry.flows.poll;
 
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
 import static google.registry.model.eppoutput.Result.Code.Success;
 import static google.registry.model.eppoutput.Result.Code.SuccessWithNoMessages;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.googlecode.objectify.Key;
+import com.googlecode.objectify.Work;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.ObjectDoesNotExistException;
 import google.registry.flows.EppException.ParameterValueSyntaxErrorException;
 import google.registry.flows.EppException.RequiredParameterMissingException;
+import google.registry.flows.FlowModule.ClientId;
+import google.registry.flows.FlowModule.PollMessageId;
+import google.registry.flows.LoggedInFlow;
 import google.registry.flows.TransactionalFlow;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.poll.MessageQueueInfo;
@@ -42,34 +47,36 @@ import org.joda.time.DateTime;
  * @error {@link PollAckFlow.MissingMessageIdException}
  * @error {@link PollAckFlow.NotAuthorizedToAckMessageException}
  */
-public class PollAckFlow extends PollFlow implements TransactionalFlow {
+public class PollAckFlow extends LoggedInFlow implements TransactionalFlow {
 
+  @Inject @ClientId String clientId;
+  @Inject @PollMessageId String messageId;
   @Inject PollAckFlow() {}
 
   @Override
   public final EppOutput run() throws EppException {
-    if (command.getMessageId() == null) {
+    if (messageId.isEmpty()) {
       throw new MissingMessageIdException();
     }
 
     Key<PollMessage> pollMessageKey;
     // Try parsing the messageId, and throw an exception if it's invalid.
     try {
-      pollMessageKey = PollMessage.EXTERNAL_KEY_CONVERTER.reverse().convert(command.getMessageId());
+      pollMessageKey = PollMessage.EXTERNAL_KEY_CONVERTER.reverse().convert(messageId);
     } catch (PollMessageExternalKeyParseException e) {
-      throw new InvalidMessageIdException(command.getMessageId());
+      throw new InvalidMessageIdException(messageId);
     }
 
     // Load the message to be acked. If a message is queued to be delivered in the future, we treat
     // it as if it doesn't exist yet.
     PollMessage pollMessage = ofy().load().key(pollMessageKey).now();
     if (pollMessage == null || !isBeforeOrAt(pollMessage.getEventTime(), now)) {
-      throw new MessageDoesNotExistException(command.getMessageId());
+      throw new MessageDoesNotExistException(messageId);
     }
 
     // Make sure this client is authorized to ack this message. It could be that the message is
     // supposed to go to a different registrar.
-    if (!getClientId().equals(pollMessage.getClientId())) {
+    if (!clientId.equals(pollMessage.getClientId())) {
       throw new NotAuthorizedToAckMessageException();
     }
 
@@ -100,23 +107,23 @@ public class PollAckFlow extends PollFlow implements TransactionalFlow {
     // We need to return the new queue length. If this was the last message in the queue being
     // acked, then we return a special status code indicating that. Note that the query will
     // include the message being acked.
-    int messageCount = getMessageQueueLength();
+
+    int messageCount = ofy().doTransactionless(new Work<Integer>() {
+      @Override
+      public Integer run() {
+        return getPollMessagesQuery(clientId, now).count();
+      }});
     if (!includeAckedMessageInCount) {
       messageCount--;
     }
     if (messageCount <= 0) {
       return createOutput(SuccessWithNoMessages);
     }
-
+    return createOutput(Success, null, null, MessageQueueInfo.create(
-    return createOutput(
+        null,  // eventTime
-        Success,
+        null,  // msg
-        MessageQueueInfo.create(
+        messageCount,
-            null,              // eventTime
+        messageId));
-            null,              // msg
-            messageCount,
-            command.getMessageId()),
-        null,   // responseData
-        null);  // extensions
   }
 
   /** Registrar is not authorized to ack this message. */

java/google/registry/flows/poll/PollFlow.java

@@ -1,89 +0,0 @@
-// Copyright 2016 The Domain Registry 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.flows.poll;
-
-import static google.registry.model.ofy.ObjectifyService.ofy;
-
-import com.google.common.collect.ImmutableList;
-import com.googlecode.objectify.Key;
-import com.googlecode.objectify.Work;
-import com.googlecode.objectify.cmd.Query;
-import google.registry.flows.EppException;
-import google.registry.flows.LoggedInFlow;
-import google.registry.model.eppinput.EppInput.Poll;
-import google.registry.model.eppoutput.EppOutput;
-import google.registry.model.eppoutput.EppResponse;
-import google.registry.model.eppoutput.EppResponse.ResponseData;
-import google.registry.model.eppoutput.EppResponse.ResponseExtension;
-import google.registry.model.eppoutput.Result;
-import google.registry.model.poll.MessageQueueInfo;
-import google.registry.model.poll.PollMessage;
-import java.util.List;
-import javax.annotation.Nullable;
-
-/** Base class of EPP Poll command flows. Mostly provides datastore helper methods. */
-public abstract class PollFlow extends LoggedInFlow {
-
-  protected Poll command;
-
-  @Override
-  @SuppressWarnings("unchecked")
-  protected final void initLoggedInFlow() throws EppException {
-    command = (Poll) eppInput.getCommandWrapper().getCommand();
-  }
-
-  /**
-   * Returns a query for all poll messages for the logged in registrar in the current TLD which are
-   * not in the future.
-   */
-  private Query<PollMessage> getQuery() {
-    return ofy().doTransactionless(new Work<Query<PollMessage>>() {
-      @Override
-      public Query<PollMessage> run() {
-        return ofy().load()
-            .type(PollMessage.class)
-            .filter("clientId", getClientId())
-            .filter("eventTime <=", now.toDate());
-      }});
-  }
-
-  /** Return the length of the message queue for the logged in registrar. */
-  protected int getMessageQueueLength() {
-    return getQuery().keys().list().size();
-  }
-
-  /**
-   * Retrieves the Keys of all active PollMessage entities for the current client ordered by
-   * eventTime.
-   */
-  protected List<Key<PollMessage>> getMessageQueueKeysInOrder() {
-    return getQuery().order("eventTime").keys().list();
-  }
-
-  protected EppOutput createOutput(
-      Result.Code code,
-      MessageQueueInfo messageQueueInfo,
-      @Nullable ImmutableList<ResponseData> responseData,
-      @Nullable ImmutableList<ResponseExtension> responseExtensions) {
-    return EppOutput.create(new EppResponse.Builder()
-        .setTrid(trid)
-        .setResult(Result.create(code))
-        .setMessageQueueInfo(messageQueueInfo)
-        .setResData(responseData)
-        .setExtensions(responseExtensions)
-        .setExecutionTime(now)
-        .build());
-  }
-}

java/google/registry/flows/poll/PollFlowUtils.java

@@ -0,0 +1,36 @@
+// Copyright 2016 The Domain Registry 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.flows.poll;
+
+import static google.registry.model.ofy.ObjectifyService.ofy;
+
+import com.googlecode.objectify.cmd.Query;
+import google.registry.model.poll.PollMessage;
+import org.joda.time.DateTime;
+
+/** Static utility functions for poll flows. */
+public final class PollFlowUtils {
+
+  private PollFlowUtils() {}
+
+  /** Returns a query for poll messages for the logged in registrar which are not in the future. */
+  static Query<PollMessage> getPollMessagesQuery(String clientId, DateTime now) {
+    return ofy().load()
+        .type(PollMessage.class)
+        .filter("clientId", clientId)
+        .filter("eventTime <=", now.toDate())
+        .order("eventTime");
+  }
+}

java/google/registry/flows/poll/PollRequestFlow.java

@@ -14,18 +14,20 @@
 
 package google.registry.flows.poll;
 
+import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
 import static google.registry.model.eppoutput.Result.Code.SuccessWithAckMessage;
 import static google.registry.model.eppoutput.Result.Code.SuccessWithNoMessages;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValueSyntaxErrorException;
+import google.registry.flows.FlowModule.ClientId;
+import google.registry.flows.FlowModule.PollMessageId;
+import google.registry.flows.LoggedInFlow;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
-import java.util.List;
 import javax.inject.Inject;
 
 /**
@@ -33,34 +35,31 @@ import javax.inject.Inject;
  *
  * @error {@link PollRequestFlow.UnexpectedMessageIdException}
  */
-public class PollRequestFlow extends PollFlow {
+public class PollRequestFlow extends LoggedInFlow {
 
+  @Inject @ClientId String clientId;
+  @Inject @PollMessageId String messageId;
   @Inject PollRequestFlow() {}
 
   @Override
   public final EppOutput run() throws EppException {
-    if (command.getMessageId() != null) {
+    if (!messageId.isEmpty()) {
       throw new UnexpectedMessageIdException();
     }
-
+    // Return the oldest message from the queue.
-    List<Key<PollMessage>> pollMessageKeys = getMessageQueueKeysInOrder();
+    PollMessage pollMessage = getPollMessagesQuery(clientId, now).first().now();
-    // Retrieve the oldest message from the queue that still exists -- since the query is eventually
+    if (pollMessage == null) {
-    // consistent, it may return keys to some entities that no longer exist.
+      return createOutput(SuccessWithNoMessages);
-    for (Key<PollMessage> key : pollMessageKeys) {
-      PollMessage pollMessage = ofy().load().key(key).now();
-      if (pollMessage != null) {
-        return createOutput(
-            SuccessWithAckMessage,
-            MessageQueueInfo.create(
-                pollMessage.getEventTime(),
-                pollMessage.getMsg(),
-                pollMessageKeys.size(),
-                PollMessage.EXTERNAL_KEY_CONVERTER.convert(key)),
-            forceEmptyToNull(pollMessage.getResponseData()),
-            forceEmptyToNull(pollMessage.getResponseExtensions()));
-      }
     }
-    return createOutput(SuccessWithNoMessages);
+    return createOutput(
+        SuccessWithAckMessage,
+        forceEmptyToNull(pollMessage.getResponseData()),
+        forceEmptyToNull(pollMessage.getResponseExtensions()),
+        MessageQueueInfo.create(
+            pollMessage.getEventTime(),
+            pollMessage.getMsg(),
+            getPollMessagesQuery(clientId, now).count(),
+            PollMessage.EXTERNAL_KEY_CONVERTER.convert(Key.create(pollMessage))));
   }
 
   /** Unexpected message id. */