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;