From 546eba68bd9caf23eb631af331b3e3d5b4959137 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Wed, 23 Jun 2021 15:39:22 -0400 Subject: [PATCH] Add SQL functionality to DeleteLoadTestDataAction (#1211) * Add SQL functionality to DeleteLoadTestDataAction This isn't directly meant to be run in production so some of the rough edges (doesn't delete domains, can't delete contacts that are referenced by an existing domain) are fine. We can handle those in DeleteProberTestAction when we do the more comprehensive deletions. --- .../batch/DeleteLoadTestDataAction.java | 122 +++++++++++++++--- .../ofy/DatastoreTransactionManager.java | 5 + .../JpaTransactionManagerImpl.java | 8 +- .../transaction/TransactionManager.java | 12 +- .../tools/AckPollMessagesCommand.java | 8 +- 5 files changed, 128 insertions(+), 27 deletions(-) diff --git a/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java b/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java index 26ddc5b77..86b4e1704 100644 --- a/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java +++ b/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java @@ -15,12 +15,14 @@ package google.registry.batch; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static google.registry.config.RegistryEnvironment.PRODUCTION; import static google.registry.mapreduce.MapreduceRunner.PARAM_DRY_RUN; import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput; import static google.registry.model.ofy.ObjectifyService.auditedOfy; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.request.Action.Method.POST; +import static google.registry.util.DateTimeUtils.END_OF_TIME; import com.google.appengine.tools.mapreduce.Mapper; import com.google.common.collect.ImmutableList; @@ -28,16 +30,24 @@ import com.google.common.collect.ImmutableSet; import com.google.common.flogger.FluentLogger; import com.googlecode.objectify.Key; import google.registry.config.RegistryEnvironment; +import google.registry.flows.poll.PollFlowUtils; import google.registry.mapreduce.MapreduceRunner; import google.registry.model.EppResource; +import google.registry.model.EppResourceUtils; import google.registry.model.contact.ContactResource; +import google.registry.model.domain.DomainBase; import google.registry.model.host.HostResource; import google.registry.model.index.EppResourceIndex; import google.registry.model.index.ForeignKeyIndex; +import google.registry.model.poll.PollMessage; +import google.registry.model.reporting.HistoryEntry; +import google.registry.model.reporting.HistoryEntryDao; +import google.registry.persistence.VKey; import google.registry.request.Action; import google.registry.request.Parameter; import google.registry.request.Response; import google.registry.request.auth.Auth; +import google.registry.util.Clock; import java.util.List; import javax.inject.Inject; @@ -46,8 +56,8 @@ import javax.inject.Inject; * the associated ForeignKey and EppResourceIndex entities. * *

This only deletes contacts and hosts, NOT domains. To delete domains, use {@link - * DeleteLoadTestDataAction} and pass it the TLD(s) that the load test domains were created on. Note - * that DeleteLoadTestDataAction is safe enough to run in production whereas this mapreduce is not, + * DeleteProberDataAction} and pass it the TLD(s) that the load test domains were created on. Note + * that DeleteProberDataAction is safe enough to run in production whereas this mapreduce is not, * but this one does not need to be runnable in production because load testing isn't run against * production. */ @@ -68,15 +78,22 @@ public class DeleteLoadTestDataAction implements Runnable { */ private static final ImmutableSet LOAD_TEST_REGISTRARS = ImmutableSet.of("proxy"); - @Inject - @Parameter(PARAM_DRY_RUN) - boolean isDryRun; - - @Inject MapreduceRunner mrRunner; - @Inject Response response; + private final boolean isDryRun; + private final MapreduceRunner mrRunner; + private final Response response; + private final Clock clock; @Inject - DeleteLoadTestDataAction() {} + DeleteLoadTestDataAction( + @Parameter(PARAM_DRY_RUN) boolean isDryRun, + MapreduceRunner mrRunner, + Response response, + Clock clock) { + this.isDryRun = isDryRun; + this.mrRunner = mrRunner; + this.response = response; + this.clock = clock; + } @Override public void run() { @@ -87,14 +104,85 @@ public class DeleteLoadTestDataAction implements Runnable { !RegistryEnvironment.get().equals(PRODUCTION), "This mapreduce is not safe to run on PRODUCTION."); - mrRunner - .setJobName("Delete load test data") - .setModuleName("backend") - .runMapOnly( - new DeleteLoadTestDataMapper(isDryRun), - ImmutableList.of( - createEntityInput(ContactResource.class), createEntityInput(HostResource.class))) - .sendLinkToMapreduceConsole(response); + if (tm().isOfy()) { + mrRunner + .setJobName("Delete load test data") + .setModuleName("backend") + .runMapOnly( + new DeleteLoadTestDataMapper(isDryRun), + ImmutableList.of( + createEntityInput(ContactResource.class), createEntityInput(HostResource.class))) + .sendLinkToMapreduceConsole(response); + } else { + tm().transact( + () -> { + LOAD_TEST_REGISTRARS.forEach(this::deletePollMessages); + tm().loadAllOfStream(ContactResource.class).forEach(this::deleteContact); + tm().loadAllOfStream(HostResource.class).forEach(this::deleteHost); + }); + } + } + + private void deletePollMessages(String registrarId) { + ImmutableList pollMessages = + PollFlowUtils.createPollMessageQuery(registrarId, END_OF_TIME).list(); + if (isDryRun) { + logger.atInfo().log( + "Would delete %d poll messages for registrar %s.", pollMessages.size(), registrarId); + } else { + pollMessages.forEach(tm()::delete); + } + } + + private void deleteContact(ContactResource contact) { + if (!LOAD_TEST_REGISTRARS.contains(contact.getPersistedCurrentSponsorClientId())) { + return; + } + // We cannot remove contacts from domains in the general case, so we cannot delete contacts + // that are linked to domains (since it would break the foreign keys) + if (EppResourceUtils.isLinked(contact.createVKey(), clock.nowUtc())) { + logger.atWarning().log( + "Cannot delete contact with repo ID %s since it is referenced from a domain", + contact.getRepoId()); + return; + } + deleteResource(contact); + } + + private void deleteHost(HostResource host) { + if (!LOAD_TEST_REGISTRARS.contains(host.getPersistedCurrentSponsorClientId())) { + return; + } + VKey hostVKey = host.createVKey(); + // We can remove hosts from linked domains, so we should do so then delete the hosts + ImmutableSet> linkedDomains = + EppResourceUtils.getLinkedDomainKeys(hostVKey, clock.nowUtc(), null); + tm().loadByKeys(linkedDomains) + .values() + .forEach( + domain -> { + ImmutableSet> remainingHosts = + domain.getNsHosts().stream() + .filter(vkey -> !vkey.equals(hostVKey)) + .collect(toImmutableSet()); + tm().put(domain.asBuilder().setNameservers(remainingHosts).build()); + }); + deleteResource(host); + } + + private void deleteResource(EppResource eppResource) { + // In SQL, the only objects parented on the resource are poll messages (deleted above) and + // history objects. + ImmutableList historyObjects = + HistoryEntryDao.loadHistoryObjectsForResource(eppResource.createVKey()); + if (isDryRun) { + logger.atInfo().log( + "Would delete repo ID %s along with %d history objects", + eppResource.getRepoId(), historyObjects.size()); + } else { + historyObjects.forEach(tm()::delete); + tm().delete(eppResource); + } } /** Provides the map method that runs for each existing contact and host entity. */ diff --git a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java index eb5703ded..c60eaecf2 100644 --- a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java +++ b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java @@ -279,6 +279,11 @@ public class DatastoreTransactionManager implements TransactionManager { return ImmutableList.copyOf(getPossibleAncestorQuery(clazz)); } + @Override + public Stream loadAllOfStream(Class clazz) { + return Streams.stream(getPossibleAncestorQuery(clazz)); + } + @Override public Optional loadSingleton(Class clazz) { List elements = getPossibleAncestorQuery(clazz).limit(2).list(); diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index 1af6935c4..fa4a9a285 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java @@ -489,13 +489,17 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { @Override public ImmutableList loadAllOf(Class clazz) { + return loadAllOfStream(clazz).collect(toImmutableList()); + } + + @Override + public Stream loadAllOfStream(Class clazz) { checkArgumentNotNull(clazz, "clazz must be specified"); assertInTransaction(); return getEntityManager() .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz) .getResultStream() - .map(this::detach) - .collect(toImmutableList()); + .map(this::detach); } @Override diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java index b5f9fb750..e92d1b381 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java +++ b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java @@ -22,6 +22,7 @@ import google.registry.persistence.VKey; import java.util.NoSuchElementException; import java.util.Optional; import java.util.function.Supplier; +import java.util.stream.Stream; import org.joda.time.DateTime; /** @@ -240,6 +241,15 @@ public interface TransactionManager { */ ImmutableList loadByEntities(Iterable entities); + /** + * Returns a list of all entities of the given type that exist in the database. + * + *

The resulting list is empty if there are no entities of this type. In Datastore mode, if the + * class is a member of the cross-TLD entity group (i.e. if it has the {@link InCrossTld} + * annotation, then the correct ancestor query will automatically be applied. + */ + ImmutableList loadAllOf(Class clazz); + /** * Returns a stream of all entities of the given type that exist in the database. * @@ -247,7 +257,7 @@ public interface TransactionManager { * the class is a member of the cross-TLD entity group (i.e. if it has the {@link InCrossTld} * annotation, then the correct ancestor query will automatically be applied. */ - ImmutableList loadAllOf(Class clazz); + Stream loadAllOfStream(Class clazz); /** * Loads the only instance of this particular class, or empty if none exists. diff --git a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java b/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java index c557e85a9..55ff08414 100644 --- a/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java +++ b/core/src/main/java/google/registry/tools/AckPollMessagesCommand.java @@ -125,13 +125,7 @@ final class AckPollMessagesCommand implements CommandWithRemoteApi { if (!isNullOrEmpty(message)) { query = query.where("msg", LIKE, "%" + message + "%"); } - - query.stream() - // Detach it so that we can print out the old, non-acked version - // (for autorenews, acking changes the next event time) - // TODO(mmuller): remove after PR 1116 is merged. - .peek(jpaTm().getEntityManager()::detach) - .forEach(this::actOnPollMessage); + query.stream().forEach(this::actOnPollMessage); }); }