Accept a list of TLDs on which to delete "prober" (test) data.

Currently, DeleteProberDataAction goes over all the TLDs of type "TEST" that
end with .test, and deletes all their DomainResources and their subordinate
history entries, poll messages, billing events, ForeignKeyDomainIndex and
EppResourceIndex entities.

After this change, we can optionally supply TLDs to work on for the request using one or more "tld=" parameter. The default (if none are supplied) will still be "all TEST TLDs that end in .test".

All given TLDs must exist, and must all be of type TEST.

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

java/google/registry/batch/DeleteProberDataAction.java

@@ -14,6 +14,7 @@
 
 package google.registry.batch;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.flows.ResourceFlowUtils.updateForeignKeyIndexDeletionTime;
@@ -22,6 +23,7 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registries.getTldsOfType;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.request.Action.Method.POST;
+import static google.registry.request.RequestParameters.PARAM_TLD;
 import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.appengine.tools.mapreduce.Mapper;
@@ -30,8 +32,10 @@ import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Sets;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryEnvironment;
 import google.registry.dns.DnsQueue;
 import google.registry.mapreduce.MapreduceRunner;
 import google.registry.mapreduce.inputs.EppResourceInputs;
@@ -71,9 +75,12 @@ public class DeleteProberDataAction implements Runnable {
   private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
 
   @Inject @Parameter(PARAM_DRY_RUN) boolean isDryRun;
+  /** List of TLDs to work on. If empty - will work on all TLDs that end with .test. */
+  @Inject @Parameter(PARAM_TLD) ImmutableSet<String> tlds;
   @Inject @Config("registryAdminClientId") String registryAdminClientId;
   @Inject MapreduceRunner mrRunner;
   @Inject Response response;
+  @Inject RegistryEnvironment registryEnvironment;
   @Inject DeleteProberDataAction() {}
 
   @Override
@@ -89,10 +96,23 @@ public class DeleteProberDataAction implements Runnable {
             ImmutableList.of(EppResourceInputs.createKeyInput(DomainBase.class)))));
   }
 
-  private static ImmutableSet<String> getProberRoidSuffixes() {
+  private ImmutableSet<String> getProberRoidSuffixes() {
-    return getTldsOfType(TldType.TEST)
+    checkArgument(
+        !RegistryEnvironment.PRODUCTION.equals(registryEnvironment)
+            || tlds.stream().allMatch(tld -> tld.endsWith(".test")),
+        "On production, can only work on TLDs that end with .test");
+    ImmutableSet<String> deletableTlds =
+        getTldsOfType(TldType.TEST)
+            .stream()
+            .filter(tld -> tlds.isEmpty() ? tld.endsWith(".test") : tlds.contains(tld))
+            .collect(toImmutableSet());
+    checkArgument(
+        tlds.isEmpty() || deletableTlds.equals(tlds),
+        "If tlds are given, they must all exist and be TEST tlds. Given: %s, not found: %s",
+        tlds,
+        Sets.difference(tlds, deletableTlds));
+    return deletableTlds
         .stream()
-        .filter(tld -> tld.endsWith(".test"))
         .map(tld -> Registry.get(tld).getRoidSuffix())
         .collect(toImmutableSet());
   }
@@ -164,7 +184,8 @@ public class DeleteProberDataAction implements Runnable {
       }
 
       DomainResource domain = (DomainResource) domainBase;
-      if (domain.getFullyQualifiedDomainName().equals("nic." + domain.getTld())) {
+      String domainName = domain.getFullyQualifiedDomainName();
+      if (domainName.equals("nic." + domain.getTld())) {
         getContext().incrementCounter("skipped, NIC domain");
         return;
       }
@@ -173,7 +194,9 @@ public class DeleteProberDataAction implements Runnable {
         return;
       }
       if (!domain.getSubordinateHosts().isEmpty()) {
-        logger.warningfmt("Cannot delete domain %s because it has subordinate hosts.", domainKey);
+        logger.warningfmt(
+            "Cannot delete domain %s (%s) because it has subordinate hosts.",
+            domainName, domainKey);
         getContext().incrementCounter("skipped, had subordinate host(s)");
         return;
       }
@@ -184,7 +207,7 @@ public class DeleteProberDataAction implements Runnable {
       // time the mapreduce is run.
       if (EppResourceUtils.isActive(domain, now)) {
         if (isDryRun) {
-          logger.infofmt("Would soft-delete the active domain: %s", domainKey);
+          logger.infofmt("Would soft-delete the active domain: %s (%s)", domainName, domainKey);
         } else {
           softDeleteDomain(domain);
         }

javatests/google/registry/batch/BUILD

@@ -13,6 +13,7 @@ java_library(
     deps = [
         "//java/google/registry/batch",
         "//java/google/registry/bigquery",
+        "//java/google/registry/config",
         "//java/google/registry/flows",
         "//java/google/registry/mapreduce",
         "//java/google/registry/model",

javatests/google/registry/batch/DeleteProberDataActionTest.java

@@ -33,6 +33,7 @@ import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
+import google.registry.config.RegistryEnvironment;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -68,6 +69,10 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
     createTld("example", "EXAMPLE");
     persistResource(Registry.get("example").asBuilder().setTldType(TldType.TEST).build());
 
+    // Since "not-test.test" isn't of TEST type, its entities won't be deleted even though it ends
+    // with .test.
+    createTld("not-test.test", "EXTEST");
+
     // Entities in these two should be deleted.
     createTld("ib-any.test", "IBANYT");
     persistResource(Registry.get("ib-any.test").asBuilder().setTldType(TldType.TEST).build());
@@ -82,6 +87,8 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
     action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.isDryRun = false;
+    action.tlds = ImmutableSet.of();
+    action.registryEnvironment = RegistryEnvironment.SANDBOX;
     action.registryAdminClientId = "TheRegistrar";
   }
 
@@ -94,15 +101,64 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   public void test_deletesAllAndOnlyProberData() throws Exception {
     Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
     Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
+    Set<ImmutableObject> notTestEntities = persistLotsOfDomains("not-test.test");
     Set<ImmutableObject> ibEntities = persistLotsOfDomains("ib-any.test");
     Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
     runMapreduce();
     assertNotDeleted(tldEntities);
     assertNotDeleted(exampleEntities);
+    assertNotDeleted(notTestEntities);
     assertDeleted(ibEntities);
     assertDeleted(oaEntities);
   }
 
+  @Test
+  public void testSuccess_deletesAllAndOnlyGivenTlds() throws Exception {
+    Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
+    Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
+    Set<ImmutableObject> notTestEntities = persistLotsOfDomains("not-test.test");
+    Set<ImmutableObject> ibEntities = persistLotsOfDomains("ib-any.test");
+    Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
+    action.tlds = ImmutableSet.of("example", "ib-any.test");
+    runMapreduce();
+    assertNotDeleted(tldEntities);
+    assertNotDeleted(notTestEntities);
+    assertNotDeleted(oaEntities);
+    assertDeleted(exampleEntities);
+    assertDeleted(ibEntities);
+  }
+
+  @Test
+  public void testFail_givenNonTestTld() throws Exception {
+    action.tlds = ImmutableSet.of("not-test.test");
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, this::runMapreduce);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("If tlds are given, they must all exist and be TEST tlds");
+  }
+
+  @Test
+  public void testFail_givenNonExistentTld() throws Exception {
+    action.tlds = ImmutableSet.of("non-existent.test");
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, this::runMapreduce);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("If tlds are given, they must all exist and be TEST tlds");
+  }
+
+  @Test
+  public void testFail_givenNonDotTestTldOnProd() throws Exception {
+    action.tlds = ImmutableSet.of("example");
+    action.registryEnvironment = RegistryEnvironment.PRODUCTION;
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, this::runMapreduce);
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("On production, can only work on TLDs that end with .test");
+  }
+
   @Test
   public void testSuccess_doesntDeleteNicDomainForProbers() throws Exception {
     DomainResource nic = persistActiveDomain("nic.ib-any.test");