Add a nomulus tool command to bulk ACK poll messages

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

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

@@ -26,7 +26,7 @@ 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) {
+  public static Query<PollMessage> getPollMessagesQuery(String clientId, DateTime now) {
     return ofy().load()
         .type(PollMessage.class)
         .filter("clientId", clientId)

java/google/registry/tools/AckPollMessagesCommand.java

@@ -0,0 +1,111 @@
+// Copyright 2019 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.tools;
+
+import static com.google.common.base.Strings.isNullOrEmpty;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.cmd.QueryKeys;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.poll.PollMessage.Autorenew;
+import google.registry.model.poll.PollMessage.OneTime;
+import google.registry.util.Clock;
+import java.util.List;
+import javax.inject.Inject;
+
+/**
+ * Command to acknowledge one-time poll messages for a registrar.
+ *
+ * <p>This is useful to bulk ACK a large number of {@link PollMessage}s for a given registrar that
+ * are gumming up that registrar's queue. ACKed poll messages are printed to stdout, so they can be
+ * piped to a file and delivered to the registrar out of band if necessary. Note that the poll
+ * messages are printed in an abbreviated CSV format (i.e. not the full EPP XML output) for
+ * brevity's sake when dealing with many poll messages.
+ *
+ * <p>You may specify a string that poll messages to be ACKed should contain, which is useful if the
+ * overwhelming majority of a backlog is caused by a single type of poll message (e.g. contact
+ * delete confirmations) and it is desired that the registrar be able to ACK the rest of the poll
+ * messages in-band.
+ *
+ * <p>This command only ACKs {@link OneTime} poll messages because that's our only use case so far,
+ * but it could be extended to ACK {@link Autorenew} poll messages as well if needed. The main
+ * difference is that one-time poll messages are deleted when ACKed whereas Autorenews are sometimes
+ * modified and re-saved instead, if the corresponding domain is still active.
+ *
+ * <p>In all cases it is not permissible to ACK a poll message until it has been delivered (i.e. its
+ * event time is in the past), same as through EPP.
+ */
+@Parameters(separators = " =", commandDescription = "Acknowledge one-time poll messages.")
+final class AckPollMessagesCommand implements CommandWithRemoteApi {
+
+  @Parameter(
+      names = {"-c", "--client"},
+      description = "Client identifier of the registrar whose poll messages should be ACKed",
+      required = true
+  )
+  private String clientId;
+
+  @Parameter(
+      names = {"-m", "--message"},
+      description = "A string that poll messages to be ACKed must contain (else all will be ACKed)"
+  )
+  private String message;
+
+  @Parameter(
+      names = {"-d", "--dry_run"},
+      description = "Do not actually commit any mutations")
+  private boolean dryRun;
+
+  @Inject Clock clock;
+
+  private static final int BATCH_SIZE = 20;
+
+  @Override
+  public void run() {
+    QueryKeys<PollMessage> query = getPollMessagesQuery(clientId, clock.nowUtc()).keys();
+    for (List<Key<PollMessage>> keys : Iterables.partition(query, BATCH_SIZE)) {
+      ofy()
+          .transact(
+              () -> {
+                // Load poll messages and filter to just those of interest.
+                ImmutableList<PollMessage> pollMessages =
+                    ofy().load().keys(keys).values().stream()
+                        .filter(pm -> pm instanceof OneTime)
+                        .filter(pm -> isNullOrEmpty(message) || pm.getMsg().contains(message))
+                        .collect(toImmutableList());
+                if (!dryRun) {
+                  ofy().delete().entities(pollMessages).now();
+                }
+                pollMessages.forEach(
+                    pm ->
+                        System.out.println(
+                            Joiner.on(',')
+                                .join(
+                                    makePollMessageExternalId(pm),
+                                    pm.getEventTime(),
+                                    pm.getMsg())));
+              });
+    }
+  }
+}

java/google/registry/tools/RegistryTool.java

@@ -29,6 +29,7 @@ public final class RegistryTool {
    */
   public static final ImmutableMap<String, Class<? extends Command>> COMMAND_MAP =
       new ImmutableMap.Builder<String, Class<? extends Command>>()
+          .put("ack_poll_messages", AckPollMessagesCommand.class)
           .put("canonicalize_labels", CanonicalizeLabelsCommand.class)
           .put("check_domain", CheckDomainCommand.class)
           .put("check_domain_claims", CheckDomainClaimsCommand.class)

java/google/registry/tools/RegistryToolComponent.java

@@ -73,6 +73,7 @@ import javax.inject.Singleton;
       WhoisModule.class,
     })
 interface RegistryToolComponent {
+  void inject(AckPollMessagesCommand command);
   void inject(CheckDomainClaimsCommand command);
   void inject(CheckDomainCommand command);
   void inject(CountDomainsCommand command);

javatests/google/registry/tools/AckPollMessagesCommandTest.java

@@ -0,0 +1,132 @@
+// Copyright 2019 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatastoreHelper.persistResource;
+
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.poll.PollMessage.Autorenew;
+import google.registry.model.poll.PollMessage.OneTime;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectRule;
+import org.joda.time.DateTime;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+/** Unit tests for {@link AckPollMessagesCommand}. */
+public class AckPollMessagesCommandTest extends CommandTestCase<AckPollMessagesCommand> {
+
+  private FakeClock clock = new FakeClock(DateTime.parse("2015-02-04T08:16:32.064Z"));
+
+  @Rule public final InjectRule inject = new InjectRule();
+
+  @Before
+  public final void before() {
+    inject.setStaticField(Ofy.class, "clock", clock);
+    command.clock = clock;
+  }
+
+  @Test
+  public void testSuccess_doesntDeletePollMessagesInFuture() throws Exception {
+    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
+    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
+    OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger");
+    OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
+    runCommand("-c", "TheRegistrar");
+    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values()).containsExactly(pm4);
+    assertInStdout(
+        "1-FSDGS-TLD-2406-624-2013,2013-05-01T22:33:44.000Z,ninelives",
+        "1-FSDGS-TLD-2406-316-2014,2014-01-01T22:33:44.000Z,foobar",
+        "1-FSDGS-TLD-2406-791-2015,2015-01-08T22:33:44.000Z,ginger");
+    assertNotInStdout("1-FSDGS-TLD-2406-123-2015,2015-09-01T22:33:44.000Z,notme");
+  }
+
+  @Test
+  public void testSuccess_doesntDeleteAutorenewPollMessages() throws Exception {
+    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
+    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
+    Autorenew pm3 =
+        persistResource(
+            new PollMessage.Autorenew.Builder()
+                .setId(624L)
+                .setParentKey(
+                    Key.create(
+                        Key.create(DomainBase.class, "AAFSGS-TLD"), HistoryEntry.class, 99406L))
+                .setEventTime(DateTime.parse("2011-04-15T22:33:44Z"))
+                .setClientId("TheRegistrar")
+                .build());
+    runCommand("-c", "TheRegistrar");
+    assertThat(ofy().load().entities(pm1, pm2, pm3).values()).containsExactly(pm3);
+  }
+
+  @Test
+  public void testSuccess_onlyDeletesPollMessagesMatchingMessage() throws Exception {
+    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good");
+    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
+    OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "mmmmmfood");
+    OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "time flies");
+    runCommand("-c", "TheRegistrar", "-m", "food");
+    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values()).containsExactly(pm2, pm4);
+  }
+
+  @Test
+  public void testSuccess_onlyDeletesPollMessagesMatchingClientId() throws Exception {
+    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good");
+    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "theft is bad");
+    OneTime pm3 =
+        persistResource(
+            new PollMessage.OneTime.Builder()
+                .setId(2474L)
+                .setParentKey(
+                    Key.create(
+                        Key.create(DomainBase.class, "FSDGS-TLD"), HistoryEntry.class, 2406L))
+                .setClientId("NewRegistrar")
+                .setEventTime(DateTime.parse("2013-06-01T22:33:44Z"))
+                .setMsg("baaaahh")
+                .build());
+    runCommand("-c", "TheRegistrar");
+    assertThat(ofy().load().entities(pm1, pm2, pm3).values()).containsExactly(pm3);
+  }
+
+  @Test
+  public void testSuccess_dryRunDoesntDeleteAnything() throws Exception {
+    OneTime pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar");
+    OneTime pm2 = persistPollMessage(624L, DateTime.parse("2013-05-01T22:33:44Z"), "ninelives");
+    OneTime pm3 = persistPollMessage(791L, DateTime.parse("2015-01-08T22:33:44Z"), "ginger");
+    OneTime pm4 = persistPollMessage(123L, DateTime.parse("2015-09-01T22:33:44Z"), "notme");
+    runCommand("-c", "TheRegistrar", "-d");
+    assertThat(ofy().load().entities(pm1, pm2, pm3, pm4).values())
+        .containsExactly(pm1, pm2, pm3, pm4);
+  }
+
+  private static OneTime persistPollMessage(long id, DateTime eventTime, String message) {
+    return persistResource(
+        new PollMessage.OneTime.Builder()
+            .setId(id)
+            .setParentKey(
+                Key.create(Key.create(DomainBase.class, "FSDGS-TLD"), HistoryEntry.class, 2406L))
+            .setClientId("TheRegistrar")
+            .setEventTime(eventTime)
+            .setMsg(message)
+            .build());
+  }
+}