From 1fb08c4c17a44eb7a5db201331c6c59e7ccbb1f0 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Fri, 2 Apr 2021 19:57:26 -0400 Subject: [PATCH] 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. --- .../registry/flows/poll/PollAckFlow.java | 24 +++-- .../registry/flows/poll/PollFlowUtils.java | 59 ++++++++-- .../registry/flows/poll/PollRequestFlow.java | 22 ++-- .../registry/model/poll/PollMessage.java | 6 +- .../poll/PollMessageExternalKeyConverter.java | 26 ++--- .../tools/AckPollMessagesCommand.java | 102 +++++++++++++----- .../registry/flows/poll/PollAckFlowTest.java | 38 +++---- .../flows/poll/PollRequestFlowTest.java | 28 ++--- .../PollMessageExternalKeyConverterTest.java | 79 ++++++++------ .../registry/testing/DatabaseHelper.java | 16 +++ .../tools/AckPollMessagesCommandTest.java | 82 ++++++++------ .../google/registry/model/schema.txt | 6 +- 12 files changed, 322 insertions(+), 166 deletions(-) diff --git a/core/src/main/java/google/registry/flows/poll/PollAckFlow.java b/core/src/main/java/google/registry/flows/poll/PollAckFlow.java index d027f166f..67a0f87b4 100644 --- a/core/src/main/java/google/registry/flows/poll/PollAckFlow.java +++ b/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 pollMessageKey; + VKey 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 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) { diff --git a/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java b/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java index 6befedeb7..454f4d2a5 100644 --- a/core/src/main/java/google/registry/flows/poll/PollFlowUtils.java +++ b/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 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 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 datastorePollMessageQuery(String registrarId, DateTime now) { + return ofy() + .load() + .type(PollMessage.class) + .filter("clientId", registrarId) + .filter("eventTime <=", now.toDate()) + .order("eventTime"); + } + + private PollFlowUtils() {} } diff --git a/core/src/main/java/google/registry/flows/poll/PollRequestFlow.java b/core/src/main/java/google/registry/flows/poll/PollRequestFlow.java index a5dc53250..bbe2f8849 100644 --- a/core/src/main/java/google/registry/flows/poll/PollRequestFlow.java +++ b/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 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(); } diff --git a/core/src/main/java/google/registry/model/poll/PollMessage.java b/core/src/main/java/google/registry/model/poll/PollMessage.java index 50211cac1..09cc6b01a 100644 --- a/core/src/main/java/google/registry/model/poll/PollMessage.java +++ b/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 parent; + @Parent @DoNotHydrate @Transient Key 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 getParentKey() { + public Key getParentKey() { return parent; } @@ -239,7 +239,7 @@ public abstract class PollMessage extends ImmutableObject return thisCastToDerived(); } - public B setParentKey(Key parentKey) { + public B setParentKey(Key parentKey) { getInstance().parent = parentKey; return thisCastToDerived(); } diff --git a/core/src/main/java/google/registry/model/poll/PollMessageExternalKeyConverter.java b/core/src/main/java/google/registry/model/poll/PollMessageExternalKeyConverter.java index ecba67bf1..0eccac613 100644 --- a/core/src/main/java/google/registry/model/poll/PollMessageExternalKeyConverter.java +++ b/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. * - *

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 + *

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 parsePollMessageExternalId(String externalKey) { + public static VKey parsePollMessageExternalId(String externalKey) { List 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(); diff --git a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java b/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java index b66208cc9..62f432d39 100644 --- a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java +++ b/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 query = getPollMessagesQuery(clientId, clock.nowUtc()).keys(); - // TODO(b/160325686): Remove the batch logic after db migration. - for (List> keys : Iterables.partition(query, BATCH_SIZE)) { - tm().transact( - () -> { - // Load poll messages and filter to just those of interest. - ImmutableList 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. + * + *

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 query = datastorePollMessageQuery(clientId, clock.nowUtc()).keys(); + for (List> 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 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. + * + *

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())); + } } diff --git a/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java b/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java index 5a6e1f399..6a3ceec60 100644 --- a/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java +++ b/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java @@ -32,14 +32,16 @@ import google.registry.flows.poll.PollAckFlow.NotAuthorizedToAckMessageException import google.registry.model.contact.ContactResource; import google.registry.model.domain.DomainBase; import google.registry.model.poll.PollMessage; +import google.registry.testing.DualDatabaseTest; import google.registry.testing.ReplayExtension; +import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Order; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Unit tests for {@link PollAckFlow}. */ +@DualDatabaseTest class PollAckFlowTest extends FlowTestCase { @Order(value = Order.DEFAULT - 2) @@ -89,13 +91,13 @@ class PollAckFlowTest extends FlowTestCase { .build()); } - @Test + @TestOfyAndSql void testDryRun() throws Exception { persistOneTimePollMessage(MESSAGE_ID); dryRunFlowAssertResponse(loadFile("poll_ack_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_contactPollMessage() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "2-2-ROID-4-3-2011")); persistResource( @@ -110,7 +112,7 @@ class PollAckFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_ack_response_empty.xml")); } - @Test + @TestOfyAndSql void testFailure_contactPollMessage_withIncorrectYearField() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "2-2-ROID-4-3-1999")); persistResource( @@ -125,14 +127,14 @@ class PollAckFlowTest extends FlowTestCase { assertThrows(MessageDoesNotExistException.class, this::runFlow); } - @Test + @TestOfyAndSql void testSuccess_messageOnContactResource() throws Exception { persistOneTimePollMessage(MESSAGE_ID); assertTransactionalFlow(true); runFlowAssertResponse(loadFile("poll_ack_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_recentActiveAutorenew() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "1-3-EXAMPLE-4-3-2010")); persistAutorenewPollMessage(clock.nowUtc().minusMonths(6), END_OF_TIME); @@ -140,7 +142,7 @@ class PollAckFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_ack_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_oldActiveAutorenew() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "1-3-EXAMPLE-4-3-2009")); persistAutorenewPollMessage(clock.nowUtc().minusYears(2), END_OF_TIME); @@ -156,7 +158,7 @@ class PollAckFlowTest extends FlowTestCase { ImmutableMap.of("MSGID", "1-3-EXAMPLE-4-3-2009", "COUNT", "4"))); } - @Test + @TestOfyAndSql void testSuccess_oldInactiveAutorenew() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "1-3-EXAMPLE-4-3-2010")); persistAutorenewPollMessage(clock.nowUtc().minusMonths(6), clock.nowUtc()); @@ -164,7 +166,7 @@ class PollAckFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_ack_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_moreMessages() throws Exception { // Create five messages to be queued for retrieval, one of which will be acked. for (int i = 0; i < 5; i++) { @@ -177,28 +179,28 @@ class PollAckFlowTest extends FlowTestCase { ImmutableMap.of("MSGID", "1-3-EXAMPLE-4-3-2011", "COUNT", "4"))); } - @Test + @TestOfyAndSql void testFailure_noSuchMessage() throws Exception { assertTransactionalFlow(true); Exception e = assertThrows(MessageDoesNotExistException.class, this::runFlow); assertThat(e).hasMessageThat().contains(String.format("(1-3-EXAMPLE-4-%d-2011)", MESSAGE_ID)); } - @Test + @TestOfyAndSql void testFailure_invalidId_tooFewComponents() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "1-2-3")); assertTransactionalFlow(true); assertThrows(InvalidMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_invalidId_tooManyComponents() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "2-2-ROID-4-3-1999-2007")); assertTransactionalFlow(true); assertThrows(InvalidMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_contactPollMessage_withMissingYearField() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "2-2-ROID-4-3")); persistResource( @@ -213,28 +215,28 @@ class PollAckFlowTest extends FlowTestCase { assertThrows(InvalidMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_invalidId_stringInsteadOfNumeric() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "ABC-12345")); assertTransactionalFlow(true); assertThrows(InvalidMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_invalidEppResourceClassId() throws Exception { setEppInput("poll_ack.xml", ImmutableMap.of("MSGID", "999-1-1-1")); assertTransactionalFlow(true); assertThrows(InvalidMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_missingId() throws Exception { setEppInput("poll_ack_missing_id.xml"); assertTransactionalFlow(true); assertThrows(MissingMessageIdException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_differentRegistrar() throws Exception { persistResource( new PollMessage.OneTime.Builder() @@ -248,7 +250,7 @@ class PollAckFlowTest extends FlowTestCase { assertThrows(NotAuthorizedToAckMessageException.class, this::runFlow); } - @Test + @TestOfyAndSql void testFailure_messageInFuture() throws Exception { persistResource( new PollMessage.OneTime.Builder() diff --git a/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java b/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java index cc9fad844..310258a38 100644 --- a/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java +++ b/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java @@ -38,14 +38,16 @@ import google.registry.model.reporting.HistoryEntry; import google.registry.model.transfer.TransferResponse.ContactTransferResponse; import google.registry.model.transfer.TransferResponse.DomainTransferResponse; import google.registry.model.transfer.TransferStatus; +import google.registry.testing.DualDatabaseTest; import google.registry.testing.ReplayExtension; import google.registry.testing.SetClockExtension; +import google.registry.testing.TestOfyAndSql; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Order; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Unit tests for {@link PollRequestFlow}. */ +@DualDatabaseTest class PollRequestFlowTest extends FlowTestCase { @Order(value = Order.DEFAULT - 3) @@ -92,14 +94,14 @@ class PollRequestFlowTest extends FlowTestCase { .build()); } - @Test + @TestOfyAndSql void testSuccess_domainTransferApproved() throws Exception { persistPendingTransferPollMessage(); assertTransactionalFlow(false); runFlowAssertResponse(loadFile("poll_response_domain_transfer.xml")); } - @Test + @TestOfyAndSql void testSuccess_clTridNotSpecified() throws Exception { setEppInput("poll_no_cltrid.xml"); persistPendingTransferPollMessage(); @@ -107,7 +109,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_domain_transfer_no_cltrid.xml")); } - @Test + @TestOfyAndSql void testSuccess_contactTransferPending() throws Exception { setClientIdForFlow("TheRegistrar"); persistResource( @@ -130,7 +132,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_contact_transfer.xml")); } - @Test + @TestOfyAndSql void testSuccess_domainPendingActionComplete() throws Exception { persistResource( new PollMessage.OneTime.Builder() @@ -145,7 +147,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_domain_pending_notification.xml")); } - @Test + @TestOfyAndSql void testSuccess_domainAutorenewMessage() throws Exception { persistResource( new PollMessage.Autorenew.Builder() @@ -159,12 +161,12 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_autorenew.xml")); } - @Test + @TestOfyAndSql void testSuccess_empty() throws Exception { runFlowAssertResponse(loadFile("poll_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_wrongRegistrar() throws Exception { persistResource( new PollMessage.OneTime.Builder() @@ -176,7 +178,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_futurePollMessage() throws Exception { persistResource( new PollMessage.OneTime.Builder() @@ -188,7 +190,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_futureAutorenew() throws Exception { persistResource( new PollMessage.Autorenew.Builder() @@ -202,7 +204,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_empty.xml")); } - @Test + @TestOfyAndSql void testSuccess_contactDelete() throws Exception { // Contact delete poll messages do not have any response data, so ensure that no // response data block is produced in the poll message. @@ -223,7 +225,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_contact_delete.xml")); } - @Test + @TestOfyAndSql void testSuccess_hostDelete() throws Exception { // Host delete poll messages do not have any response data, so ensure that no // response data block is produced in the poll message. @@ -245,7 +247,7 @@ class PollRequestFlowTest extends FlowTestCase { runFlowAssertResponse(loadFile("poll_response_host_delete.xml")); } - @Test + @TestOfyAndSql void testFailure_messageIdProvided() throws Exception { setEppInput("poll_with_id.xml"); assertTransactionalFlow(false); diff --git a/core/src/test/java/google/registry/model/poll/PollMessageExternalKeyConverterTest.java b/core/src/test/java/google/registry/model/poll/PollMessageExternalKeyConverterTest.java index a9dc6a60b..2fca5cab0 100644 --- a/core/src/test/java/google/registry/model/poll/PollMessageExternalKeyConverterTest.java +++ b/core/src/test/java/google/registry/model/poll/PollMessageExternalKeyConverterTest.java @@ -25,21 +25,25 @@ import static google.registry.testing.DatabaseHelper.persistResource; import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.assertThrows; -import com.googlecode.objectify.Key; +import google.registry.model.domain.DomainHistory; import google.registry.model.domain.Period; import google.registry.model.eppcommon.Trid; import google.registry.model.ofy.Ofy; import google.registry.model.poll.PollMessageExternalKeyConverter.PollMessageExternalKeyParseException; import google.registry.model.reporting.HistoryEntry; +import google.registry.persistence.VKey; import google.registry.testing.AppEngineExtension; +import google.registry.testing.DatabaseHelper; +import google.registry.testing.DualDatabaseTest; import google.registry.testing.FakeClock; import google.registry.testing.InjectExtension; +import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Unit tests for {@link PollMessageExternalKeyConverter}. */ +@DualDatabaseTest public class PollMessageExternalKeyConverterTest { @RegisterExtension @@ -55,21 +59,23 @@ public class PollMessageExternalKeyConverterTest { void beforeEach() { inject.setStaticField(Ofy.class, "clock", clock); createTld("foobar"); - historyEntry = persistResource(new HistoryEntry.Builder() - .setParent(persistActiveDomain("foo.foobar")) - .setType(HistoryEntry.Type.DOMAIN_CREATE) - .setPeriod(Period.create(1, Period.Unit.YEARS)) - .setXmlBytes("".getBytes(UTF_8)) - .setModificationTime(clock.nowUtc()) - .setClientId("foo") - .setTrid(Trid.create("ABC-123", "server-trid")) - .setBySuperuser(false) - .setReason("reason") - .setRequestedByRegistrar(false) - .build()); + historyEntry = + persistResource( + new DomainHistory.Builder() + .setParent(persistActiveDomain("foo.foobar")) + .setType(HistoryEntry.Type.DOMAIN_CREATE) + .setPeriod(Period.create(1, Period.Unit.YEARS)) + .setXmlBytes("".getBytes(UTF_8)) + .setModificationTime(clock.nowUtc()) + .setClientId("TheRegistrar") + .setTrid(Trid.create("ABC-123", "server-trid")) + .setBySuperuser(false) + .setReason("reason") + .setRequestedByRegistrar(false) + .build()); } - @Test + @TestOfyAndSql void testSuccess_domain() { PollMessage.OneTime pollMessage = persistResource( @@ -80,14 +86,14 @@ public class PollMessageExternalKeyConverterTest { .setParent(historyEntry) .build()); assertThat(makePollMessageExternalId(pollMessage)).isEqualTo("1-2-FOOBAR-4-5-2007"); - assertThat(parsePollMessageExternalId("1-2-FOOBAR-4-5-2007")) - .isEqualTo(Key.create(pollMessage)); + assertVKeysEqual(parsePollMessageExternalId("1-2-FOOBAR-4-5-2007"), pollMessage.createVKey()); } - @Test + @TestOfyAndSql void testSuccess_contact() { historyEntry = - persistResource(historyEntry.asBuilder().setParent(persistActiveContact("tim")).build()); + persistResource( + DatabaseHelper.createHistoryEntryForEppResource(persistActiveContact("tim"))); PollMessage.OneTime pollMessage = persistResource( new PollMessage.OneTime.Builder() @@ -96,14 +102,15 @@ public class PollMessageExternalKeyConverterTest { .setMsg("Test poll message") .setParent(historyEntry) .build()); - assertThat(makePollMessageExternalId(pollMessage)).isEqualTo("2-5-ROID-4-6-2007"); - assertThat(parsePollMessageExternalId("2-5-ROID-4-6-2007")).isEqualTo(Key.create(pollMessage)); + assertThat(makePollMessageExternalId(pollMessage)).isEqualTo("2-5-ROID-6-7-2007"); + assertVKeysEqual(parsePollMessageExternalId("2-5-ROID-6-7-2007"), pollMessage.createVKey()); } - @Test + @TestOfyAndSql void testSuccess_host() { historyEntry = - persistResource(historyEntry.asBuilder().setParent(persistActiveHost("time.zyx")).build()); + persistResource( + DatabaseHelper.createHistoryEntryForEppResource(persistActiveHost("time.xyz"))); PollMessage.OneTime pollMessage = persistResource( new PollMessage.OneTime.Builder() @@ -112,18 +119,18 @@ public class PollMessageExternalKeyConverterTest { .setMsg("Test poll message") .setParent(historyEntry) .build()); - assertThat(makePollMessageExternalId(pollMessage)).isEqualTo("3-5-ROID-4-6-2007"); - assertThat(parsePollMessageExternalId("3-5-ROID-4-6-2007")).isEqualTo(Key.create(pollMessage)); + assertThat(makePollMessageExternalId(pollMessage)).isEqualTo("3-5-ROID-6-7-2007"); + assertVKeysEqual(parsePollMessageExternalId("3-5-ROID-6-7-2007"), pollMessage.createVKey()); } - @Test + @TestOfyAndSql void testFailure_missingYearField() { assertThrows( PollMessageExternalKeyParseException.class, () -> parsePollMessageExternalId("1-2-FOOBAR-4-5")); } - @Test + @TestOfyAndSql void testFailure_invalidEppResourceTypeId() { // Populate the testdata correctly as for 1-2-FOOBAR-4-5 so we know that the only thing that // is wrong here is the EppResourceTypeId. @@ -133,24 +140,36 @@ public class PollMessageExternalKeyConverterTest { () -> parsePollMessageExternalId("4-2-FOOBAR-4-5-2007")); } - @Test + @TestOfyAndSql void testFailure_tooFewComponentParts() { assertThrows( PollMessageExternalKeyParseException.class, () -> parsePollMessageExternalId("1-3-EXAMPLE")); } - @Test + @TestOfyAndSql void testFailure_tooManyComponentParts() { assertThrows( PollMessageExternalKeyParseException.class, () -> parsePollMessageExternalId("1-3-EXAMPLE-4-5-2007-2009")); } - @Test + @TestOfyAndSql void testFailure_nonNumericIds() { assertThrows( PollMessageExternalKeyParseException.class, () -> parsePollMessageExternalId("A-B-FOOBAR-D-E-F")); } + + // We may have VKeys of slightly varying types, e.g. VKey (superclass) and + // VKey (subclass). We should treat these as equal since the DB does. + private static void assertVKeysEqual( + VKey one, VKey two) { + assertThat( + one.getKind().isAssignableFrom(two.getKind()) + || two.getKind().isAssignableFrom(one.getKind())) + .isTrue(); + assertThat(one.getSqlKey()).isEqualTo(two.getSqlKey()); + assertThat(one.getOfyKey()).isEqualTo(two.getOfyKey()); + } } diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java index 4c57cf8e3..867e2afbd 100644 --- a/core/src/test/java/google/registry/testing/DatabaseHelper.java +++ b/core/src/test/java/google/registry/testing/DatabaseHelper.java @@ -119,6 +119,7 @@ import google.registry.tmch.LordnTaskUtils; import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.NoSuchElementException; import java.util.Optional; import java.util.Set; import java.util.function.Function; @@ -1208,6 +1209,7 @@ public class DatabaseHelper { .setType(getHistoryEntryType(parentResource)) .setModificationTime(DateTime.now(DateTimeZone.UTC)) .setParent(parentResource) + .setClientId(parentResource.getPersistedCurrentSponsorClientId()) .build()); } @@ -1329,5 +1331,19 @@ public class DatabaseHelper { return transactIfJpaTm(() -> tm().loadAllOf(clazz)); } + /** + * Loads the set of entities by their keys from the DB. + * + *

If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for + * convenience, so you don't need to wrap it in a transaction at the callsite. + * + *

Nonexistent keys / entities are absent from the resulting map, but no {@link + * NoSuchElementException} will be thrown. + */ + public static ImmutableMap, T> loadByKeysIfPresent( + Iterable> keys) { + return transactIfJpaTm(() -> tm().loadByKeysIfPresent(keys)); + } + private DatabaseHelper() {} } diff --git a/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java b/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java index 796e9b288..c61f17be7 100644 --- a/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java +++ b/core/src/test/java/google/registry/tools/AckPollMessagesCommandTest.java @@ -16,39 +16,59 @@ package google.registry.tools; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; -import static google.registry.persistence.transaction.TransactionManagerFactory.tm; +import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.loadByKeys; +import static google.registry.testing.DatabaseHelper.loadByKeysIfPresent; +import static google.registry.testing.DatabaseHelper.newDomainBase; import static google.registry.testing.DatabaseHelper.persistResource; import com.google.common.collect.ImmutableList; import com.googlecode.objectify.Key; import google.registry.model.domain.DomainBase; +import google.registry.model.domain.DomainHistory; 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.persistence.VKey; +import google.registry.testing.DualDatabaseTest; import google.registry.testing.FakeClock; import google.registry.testing.InjectExtension; +import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; /** Unit tests for {@link AckPollMessagesCommand}. */ +@DualDatabaseTest public class AckPollMessagesCommandTest extends CommandTestCase { private FakeClock clock = new FakeClock(DateTime.parse("2015-02-04T08:16:32.064Z")); @RegisterExtension public final InjectExtension inject = new InjectExtension(); + private DomainHistory domainHistory; + @BeforeEach final void beforeEach() { inject.setStaticField(Ofy.class, "clock", clock); command.clock = clock; + createTld("tld"); + DomainBase domain = newDomainBase("example.tld").asBuilder().setRepoId("FSDGS-TLD").build(); + persistResource(domain); + domainHistory = + persistResource( + new DomainHistory.Builder() + .setModificationTime(clock.nowUtc()) + .setDomainRepoId(domain.getRepoId()) + .setType(HistoryEntry.Type.DOMAIN_CREATE) + .setId(2406L) + .build()); + clock.advanceOneMilli(); } - @Test + @TestOfyAndSql void testSuccess_doesntDeletePollMessagesInFuture() throws Exception { VKey pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); @@ -60,7 +80,7 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm4 = futurePollMessage.createVKey(); runCommand("-c", "TheRegistrar"); - assertThat(tm().loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3, pm4)).values()) + assertThat(loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3, pm4)).values()) .containsExactly(futurePollMessage); assertInStdout( "1-FSDGS-TLD-2406-624-2013,2013-05-01T22:33:44.000Z,ninelives", @@ -69,7 +89,7 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); @@ -78,10 +98,8 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm3 = autorenew.createVKey(); runCommand("-c", "TheRegistrar"); - assertThat(tm().loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3)).values()) + assertThat(loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3)).values()) .containsExactly(resaved); assertInStdout( - "1-AAFSGS-TLD-99406-624-2011,2011-04-15T22:33:44.000Z,autorenew", + "1-FSDGS-TLD-2406-625-2011,2011-04-15T22:33:44.000Z,autorenew", "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"); } - @Test + @TestOfyAndSql void testSuccess_deletesExpiredAutorenewPollMessages() throws Exception { VKey pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "foobar").createVKey(); @@ -107,10 +125,8 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm3 = autorenew.createVKey(); runCommand("-c", "TheRegistrar"); - assertThat(tm().loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3))).isEmpty(); + assertThat(loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3))).isEmpty(); assertInStdout( - "1-AAFSGS-TLD-99406-624-2011,2011-04-15T22:33:44.000Z,autorenew", + "1-FSDGS-TLD-2406-625-2011,2011-04-15T22:33:44.000Z,autorenew", "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"); } - @Test + @TestOfyAndSql void testSuccess_onlyDeletesPollMessagesMatchingMessage() throws Exception { VKey pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good") @@ -139,11 +155,11 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm4 = notMatched2.createVKey(); runCommand("-c", "TheRegistrar", "-m", "food"); - assertThat(tm().loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3, pm4)).values()) + assertThat(loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3, pm4)).values()) .containsExactly(notMatched1, notMatched2); } - @Test + @TestOfyAndSql void testSuccess_onlyDeletesPollMessagesMatchingClientId() throws Exception { VKey pm1 = persistPollMessage(316L, DateTime.parse("2014-01-01T22:33:44Z"), "food is good") @@ -155,20 +171,18 @@ public class AckPollMessagesCommandTest extends CommandTestCase pm3 = notMatched.createVKey(); runCommand("-c", "TheRegistrar"); - assertThat(tm().loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3)).values()) + assertThat(loadByKeysIfPresent(ImmutableList.of(pm1, pm2, pm3)).values()) .containsExactly(notMatched); } - @Test + @TestOfyAndSql 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"); @@ -176,20 +190,22 @@ public class AckPollMessagesCommandTest extends CommandTestCase parent; + @Parent com.googlecode.objectify.Key parent; java.lang.String clientId; java.lang.String msg; org.joda.time.DateTime eventTime; } class google.registry.model.poll.PollMessage$Autorenew { @Id java.lang.Long id; - @Parent com.googlecode.objectify.Key parent; + @Parent com.googlecode.objectify.Key parent; java.lang.String clientId; java.lang.String msg; java.lang.String targetId; @@ -526,7 +526,7 @@ class google.registry.model.poll.PollMessage$Autorenew { } class google.registry.model.poll.PollMessage$OneTime { @Id java.lang.Long id; - @Parent com.googlecode.objectify.Key parent; + @Parent com.googlecode.objectify.Key parent; java.lang.String clientId; java.lang.String msg; java.util.List contactPendingActionNotificationResponses;