Convert poll-message-related classes to use SQL as well (#1050)

* Convert poll-message-related classes to use SQL as well

Two relatively complex parts. The first is that we needed a small
refactor on the AckPollMessagesCommand because we could theoretically be
acking more poll messages than the Datastore transaction size boundary.
This means that the normal flow of "gather the poll messages from the DB
into one collection, then act on it" needs to be changed to a more
functional flow.

The second is that acking the poll message (deleting it in most cases)
reduces the number of remaining poll messages in SQL but not in
Datastore, since in Datastore the deletion does not take effect until
after the transaction is over.

core/src/main/java/google/registry/flows/poll/PollAckFlow.java

@@ -16,15 +16,13 @@ package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.poll.PollFlowUtils.ackPollMessage;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.parsePollMessageExternalId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.ObjectDoesNotExistException;
@@ -39,6 +37,8 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.model.poll.PollMessageExternalKeyConverter.PollMessageExternalKeyParseException;
+import google.registry.persistence.VKey;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -71,7 +71,7 @@ public class PollAckFlow implements TransactionalFlow {
       throw new MissingMessageIdException();
     }
 
-    Key<PollMessage> pollMessageKey;
+    VKey<PollMessage> pollMessageKey;
     // Try parsing the messageId, and throw an exception if it's invalid.
     try {
       pollMessageKey = parsePollMessageExternalId(messageId);
@@ -84,12 +84,13 @@ public class PollAckFlow implements TransactionalFlow {
     // 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. Same for if the message ID year isn't the same as the actual
     // poll message's event time (that means they're passing in an old already-acked ID).
-    PollMessage pollMessage = ofy().load().key(pollMessageKey).now();
-    if (pollMessage == null
-        || !isBeforeOrAt(pollMessage.getEventTime(), now)
-        || !makePollMessageExternalId(pollMessage).equals(messageId)) {
+    Optional<PollMessage> maybePollMessage = tm().loadByKeyIfPresent(pollMessageKey);
+    if (!maybePollMessage.isPresent()
+        || !isBeforeOrAt(maybePollMessage.get().getEventTime(), now)
+        || !makePollMessageExternalId(maybePollMessage.get()).equals(messageId)) {
       throw new MessageDoesNotExistException(messageId);
     }
+    PollMessage pollMessage = maybePollMessage.get();
 
     // Make sure this client is authorized to ack this message. It could be that the message is
     // supposed to go to a different registrar.
@@ -106,8 +107,11 @@ public class PollAckFlow implements TransactionalFlow {
     // acked, then we return a special status code indicating that. Note that the query will
     // include the message being acked.
 
-    int messageCount = tm().doTransactionless(() -> getPollMessagesQuery(clientId, now).count());
-    if (!includeAckedMessageInCount) {
+    int messageCount = tm().doTransactionless(() -> getPollMessageCount(clientId, now));
+    // Within the same transaction, Datastore will not reflect the updated count (potentially
+    // reduced by one thanks to the acked poll message). SQL will, however, so we shouldn't reduce
+    // the count in the SQL case.
+    if (!includeAckedMessageInCount && tm().isOfy()) {
       messageCount--;
     }
     if (messageCount <= 0) {

core/src/main/java/google/registry/flows/poll/PollFlowUtils.java

@@ -16,25 +16,56 @@ package google.registry.flows.poll;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.poll.PollMessage;
+import java.util.Optional;
 import org.joda.time.DateTime;
 
 /** Static utility functions for poll flows. */
 public final class PollFlowUtils {
 
-  private PollFlowUtils() {}
+  public static final String SQL_POLL_MESSAGE_QUERY =
+      "FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now ORDER BY eventTime ASC";
+  private static final String SQL_POLL_MESSAGE_COUNT_QUERY =
+      "SELECT COUNT(*) FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now";
 
-  /** Returns a query for poll messages for the logged in registrar which are not in the future. */
-  public static Query<PollMessage> getPollMessagesQuery(String clientId, DateTime now) {
-    return ofy().load()
-        .type(PollMessage.class)
-        .filter("clientId", clientId)
-        .filter("eventTime <=", now.toDate())
-        .order("eventTime");
+  /** Returns the number of poll messages for the given registrar that are not in the future. */
+  public static int getPollMessageCount(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return datastorePollMessageQuery(registrarId, now).count();
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_COUNT_QUERY, Long.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .getSingleResult()
+                      .intValue());
+    }
+  }
+
+  /** Returns the first (by event time) poll message not in the future for this registrar. */
+  public static Optional<PollMessage> getFirstPollMessage(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return Optional.ofNullable(datastorePollMessageQuery(registrarId, now).first().now());
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_QUERY, PollMessage.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .setMaxResults(1)
+                      .getResultStream()
+                      .findFirst());
+    }
   }
 
   /**
@@ -74,4 +105,16 @@ public final class PollFlowUtils {
     }
     return includeAckedMessageInCount;
   }
+
+  /** A Datastore query for poll messages from the given registrar that are not in the future. */
+  public static Query<PollMessage> datastorePollMessageQuery(String registrarId, DateTime now) {
+    return ofy()
+        .load()
+        .type(PollMessage.class)
+        .filter("clientId", registrarId)
+        .filter("eventTime <=", now.toDate())
+        .order("eventTime");
+  }
+
+  private PollFlowUtils() {}
 }

core/src/main/java/google/registry/flows/poll/PollRequestFlow.java

@@ -15,7 +15,8 @@
 package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getFirstPollMessage;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACK_MESSAGE;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
@@ -31,6 +32,7 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.util.Clock;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -63,18 +65,20 @@ public class PollRequestFlow implements Flow {
     }
     // Return the oldest message from the queue.
     DateTime now = clock.nowUtc();
-    PollMessage pollMessage = getPollMessagesQuery(clientId, now).first().now();
-    if (pollMessage == null) {
+    Optional<PollMessage> maybePollMessage = getFirstPollMessage(clientId, now);
+    if (!maybePollMessage.isPresent()) {
       return responseBuilder.setResultFromCode(SUCCESS_WITH_NO_MESSAGES).build();
     }
+    PollMessage pollMessage = maybePollMessage.get();
     return responseBuilder
         .setResultFromCode(SUCCESS_WITH_ACK_MESSAGE)
-        .setMessageQueueInfo(new MessageQueueInfo.Builder()
-            .setQueueDate(pollMessage.getEventTime())
-            .setMsg(pollMessage.getMsg())
-            .setQueueLength(getPollMessagesQuery(clientId, now).count())
-            .setMessageId(makePollMessageExternalId(pollMessage))
-            .build())
+        .setMessageQueueInfo(
+            new MessageQueueInfo.Builder()
+                .setQueueDate(pollMessage.getEventTime())
+                .setMsg(pollMessage.getMsg())
+                .setQueueLength(getPollMessageCount(clientId, now))
+                .setMessageId(makePollMessageExternalId(pollMessage))
+                .build())
         .setMultipleResData(pollMessage.getResponseData())
         .build();
   }

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -106,7 +106,7 @@ public abstract class PollMessage extends ImmutableObject
   @Column(name = "poll_message_id")
   Long id;
 
-  @Parent @DoNotHydrate @Transient Key<HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<? extends HistoryEntry> parent;
 
   /** The registrar that this poll message will be delivered to. */
   @Index
@@ -134,7 +134,7 @@ public abstract class PollMessage extends ImmutableObject
 
   @Ignore Long hostHistoryRevisionId;
 
-  public Key<HistoryEntry> getParentKey() {
+  public Key<? extends HistoryEntry> getParentKey() {
     return parent;
   }
 
@@ -239,7 +239,7 @@ public abstract class PollMessage extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setParentKey(Key<HistoryEntry> parentKey) {
+    public B setParentKey(Key<? extends HistoryEntry> parentKey) {
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }

core/src/main/java/google/registry/model/poll/PollMessageExternalKeyConverter.java

@@ -24,6 +24,7 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import java.util.List;
 
 /**
@@ -78,14 +79,14 @@ public class PollMessageExternalKeyConverter {
   /**
    * Returns an Objectify Key to a PollMessage corresponding with the external ID.
    *
-   * <p>Note that the year field that is included at the end of the poll message isn't actually
-   * used for anything; it exists solely to create unique externally visible IDs for autorenews. We
-   * thus ignore it (for now) for backwards compatibility reasons, so that registrars can still ACK
+   * <p>Note that the year field that is included at the end of the poll message isn't actually used
+   * for anything; it exists solely to create unique externally visible IDs for autorenews. We thus
+   * ignore it (for now) for backwards compatibility reasons, so that registrars can still ACK
    * existing poll message IDs they may have lying around.
    *
    * @throws PollMessageExternalKeyParseException if the external key has an invalid format.
    */
-  public static Key<PollMessage> parsePollMessageExternalId(String externalKey) {
+  public static VKey<PollMessage> parsePollMessageExternalId(String externalKey) {
     List<String> idComponents = Splitter.on('-').splitToList(externalKey);
     if (idComponents.size() != 6) {
       throw new PollMessageExternalKeyParseException();
@@ -96,16 +97,17 @@ public class PollMessageExternalKeyConverter {
       if (resourceClazz == null) {
         throw new PollMessageExternalKeyParseException();
       }
-      return Key.create(
+      return VKey.from(
           Key.create(
               Key.create(
-                  null,
-                  resourceClazz,
-                  String.format("%s-%s", idComponents.get(1), idComponents.get(2))),
-              HistoryEntry.class,
-              Long.parseLong(idComponents.get(3))),
-          PollMessage.class,
-          Long.parseLong(idComponents.get(4)));
+                  Key.create(
+                      null,
+                      resourceClazz,
+                      String.format("%s-%s", idComponents.get(1), idComponents.get(2))),
+                  HistoryEntry.class,
+                  Long.parseLong(idComponents.get(3))),
+              PollMessage.class,
+              Long.parseLong(idComponents.get(4))));
       // Note that idComponents.get(5) is entirely ignored; we never use the year field internally.
     } catch (NumberFormatException e) {
       throw new PollMessageExternalKeyParseException();

core/src/main/java/google/registry/tools/AckPollMessagesCommand.java

@@ -15,16 +15,16 @@
 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.flows.poll.PollFlowUtils.SQL_POLL_MESSAGE_QUERY;
+import static google.registry.flows.poll.PollFlowUtils.datastorePollMessageQuery;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 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;
@@ -35,6 +35,7 @@ import google.registry.model.poll.PollMessage.OneTime;
 import google.registry.util.Clock;
 import java.util.List;
 import javax.inject.Inject;
+import javax.persistence.TypedQuery;
 
 /**
  * Command to acknowledge one-time poll messages for a registrar.
@@ -60,11 +61,14 @@ import javax.inject.Inject;
 @Parameters(separators = " =", commandDescription = "Acknowledge one-time poll messages.")
 final class AckPollMessagesCommand implements CommandWithRemoteApi {
 
+  private static final String SQL_POLL_MESSAGE_QUERY_BY_MESSAGE =
+      "FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now AND msg LIKE :msg"
+          + " ORDER BY eventTime ASC";
+
   @Parameter(
       names = {"-c", "--client"},
       description = "Client identifier of the registrar whose poll messages should be ACKed",
-      required = true
-  )
+      required = true)
   private String clientId;
 
   @Parameter(
@@ -84,28 +88,72 @@ final class AckPollMessagesCommand implements CommandWithRemoteApi {
 
   @Override
   public void run() {
-    QueryKeys<PollMessage> query = getPollMessagesQuery(clientId, clock.nowUtc()).keys();
-    // TODO(b/160325686): Remove the batch logic after db migration.
-    for (List<Key<PollMessage>> keys : Iterables.partition(query, BATCH_SIZE)) {
-      tm().transact(
-              () -> {
-                // Load poll messages and filter to just those of interest.
-                ImmutableList<PollMessage> pollMessages =
-                    ofy().load().keys(keys).values().stream()
-                        .filter(pm -> isNullOrEmpty(message) || pm.getMsg().contains(message))
-                        .collect(toImmutableList());
-                if (!dryRun) {
-                  pollMessages.forEach(PollFlowUtils::ackPollMessage);
-                }
-                pollMessages.forEach(
-                    pm ->
-                        System.out.println(
-                            Joiner.on(',')
-                                .join(
-                                    makePollMessageExternalId(pm),
-                                    pm.getEventTime(),
-                                    pm.getMsg())));
-              });
+    if (tm().isOfy()) {
+      ackPollMessagesDatastore();
+    } else {
+      ackPollMessagesSql();
     }
   }
+
+  /**
+   * Loads and acks the matching poll messages from Datastore.
+   *
+   * <p>We have to first load the poll message keys then batch-load the objects themselves due to
+   * the Datastore size limits.
+   */
+  private void ackPollMessagesDatastore() {
+    QueryKeys<PollMessage> query = datastorePollMessageQuery(clientId, clock.nowUtc()).keys();
+    for (List<Key<PollMessage>> keys : Iterables.partition(query, BATCH_SIZE)) {
+      tm().transact(
+              () ->
+                  // Load poll messages and filter to just those of interest.
+                  ofy().load().keys(keys).values().stream()
+                      .filter(pm -> isNullOrEmpty(message) || pm.getMsg().contains(message))
+                      .forEach(this::actOnPollMessage));
+    }
+  }
+
+  /** Loads and acks all matching poll messages from SQL in one transaction. */
+  private void ackPollMessagesSql() {
+    jpaTm()
+        .transact(
+            () -> {
+              TypedQuery<PollMessage> typedQuery;
+              if (isNullOrEmpty(message)) {
+                typedQuery = jpaTm().query(SQL_POLL_MESSAGE_QUERY, PollMessage.class);
+              } else {
+                typedQuery =
+                    jpaTm()
+                        .query(SQL_POLL_MESSAGE_QUERY_BY_MESSAGE, PollMessage.class)
+                        .setParameter("msg", "%" + message + "%");
+              }
+              typedQuery
+                  .setParameter("registrarId", clientId)
+                  .setParameter("now", clock.nowUtc())
+                  .getResultStream()
+                  // Detach it so that we can print out the old, non-acked version
+                  // (for autorenews, acking changes the next event time)
+                  .peek(jpaTm().getEntityManager()::detach)
+                  .forEach(this::actOnPollMessage);
+            });
+  }
+
+  /**
+   * Acks the poll message if not running in dry-run mode, prints regardless.
+   *
+   * <p>This is a separate function because the processing of poll messages is transactionally
+   * different between the Datastore and SQL implementations. Datastore must process the messages in
+   * batches, whereas we can load all messages from SQL in one transaction.
+   */
+  private void actOnPollMessage(PollMessage pollMessage) {
+    if (!dryRun) {
+      PollFlowUtils.ackPollMessage(pollMessage);
+    }
+    System.out.println(
+        Joiner.on(',')
+            .join(
+                makePollMessageExternalId(pollMessage),
+                pollMessage.getEventTime(),
+                pollMessage.getMsg()));
+  }
 }