Consolidate expand billing events action into 'batch' package

Now it lives alongside the delete prober data action, as well as any future
batch/maintenance tasks that should run periodically.

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

java/google/registry/batch/BUILD

@@ -6,21 +6,27 @@ licenses(["notice"])  # Apache 2.0
 
 
 java_library(
-    name = "billing",
+    name = "batch",
     srcs = glob(["*.java"]),
     deps = [
+        "//java/com/google/common/annotations",
         "//java/com/google/common/base",
         "//java/com/google/common/collect",
+        "//java/com/google/common/html",
+        "//java/com/google/common/io",
         "//java/com/google/common/net",
         "//third_party/java/appengine:appengine-api",
+        "//third_party/java/appengine_gcs_client",
         "//third_party/java/appengine_mapreduce2:appengine_mapreduce",
         "//third_party/java/dagger",
         "//third_party/java/joda_money",
         "//third_party/java/joda_time",
+        "//third_party/java/json_simple",
         "//third_party/java/jsr305_annotations",
         "//third_party/java/jsr330_inject",
         "//third_party/java/objectify:objectify-v4_1",
         "//third_party/java/servlet/servlet_api",
+        "//java/google/registry/config",
         "//java/google/registry/mapreduce",
         "//java/google/registry/mapreduce/inputs",
         "//java/google/registry/model",

java/google/registry/batch/DeleteProberDataAction.java

@@ -0,0 +1,168 @@
+// Copyright 2016 The Domain Registry Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.batch;
+
+import static com.google.common.base.Verify.verifyNotNull;
+import static google.registry.mapreduce.MapreduceRunner.PARAM_DRY_RUN;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.registry.Registries.getTldsOfType;
+import static google.registry.request.Action.Method.POST;
+
+import com.google.appengine.tools.mapreduce.Mapper;
+import com.google.common.base.Function;
+import com.google.common.base.Predicate;
+import com.google.common.base.Splitter;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.Work;
+import google.registry.mapreduce.MapreduceRunner;
+import google.registry.mapreduce.inputs.EppResourceInputs;
+import google.registry.model.domain.DomainApplication;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
+import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.TldType;
+import google.registry.request.Action;
+import google.registry.request.Parameter;
+import google.registry.request.Response;
+import google.registry.util.FormattingLogger;
+import google.registry.util.PipelineUtils;
+import java.util.List;
+import javax.inject.Inject;
+
+/**
+ * Deletes all prober DomainResources and their subordinate history entries, poll messages, and
+ * billing events, along with their ForeignKeyDomainIndex and EppResourceIndex entities.
+ *
+ * <p>See: https://www.youtube.com/watch?v=xuuv0syoHnM
+ */
+@Action(path = "/_dr/task/deleteProberData", method = POST)
+public class DeleteProberDataAction implements Runnable {
+
+  private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
+
+  @Inject @Parameter(PARAM_DRY_RUN) boolean isDryRun;
+  @Inject MapreduceRunner mrRunner;
+  @Inject Response response;
+  @Inject DeleteProberDataAction() {}
+
+  @Override
+  public void run() {
+    response.sendJavaScriptRedirect(PipelineUtils.createJobPath(mrRunner
+        .setJobName("Delete prober data")
+        .setModuleName("backend")
+        .runMapOnly(
+            new DeleteProberDataMapper(getProberRoidSuffixes(), isDryRun),
+            ImmutableList.of(EppResourceInputs.createKeyInput(DomainBase.class)))));
+  }
+
+  private static ImmutableSet<String> getProberRoidSuffixes() {
+    return FluentIterable.from(getTldsOfType(TldType.TEST))
+        .filter(new Predicate<String>() {
+          @Override
+          public boolean apply(String tld) {
+            // Extra sanity check to prevent us from nuking prod data if a real TLD accidentally
+            // gets set to type TEST.
+            return tld.endsWith(".test");
+          }})
+        .transform(
+            new Function<String, String>() {
+              @Override
+              public String apply(String tld) {
+                return Registry.get(tld).getRoidSuffix();
+              }})
+        .toSet();
+  }
+
+  /** Provides the map method that runs for each existing DomainBase entity. */
+  public static class DeleteProberDataMapper extends Mapper<Key<DomainBase>, Void, Void> {
+
+    private static final long serialVersionUID = 1737761271804180412L;
+
+    private final ImmutableSet<String> proberRoidSuffixes;
+    private final Boolean isDryRun;
+
+    public DeleteProberDataMapper(ImmutableSet<String> proberRoidSuffixes, Boolean isDryRun) {
+      this.proberRoidSuffixes = proberRoidSuffixes;
+      this.isDryRun = isDryRun;
+    }
+
+    @Override
+    public final void map(Key<DomainBase> key) {
+      try {
+        String roidSuffix = Iterables.getLast(Splitter.on('-').split(key.getName()));
+        if (proberRoidSuffixes.contains(roidSuffix)) {
+          deleteDomain(key);
+        } else {
+          getContext().incrementCounter(String.format("skipped, non-prober data"));
+        }
+      } catch (Throwable t) {
+        logger.severefmt(t, "Error while deleting prober data for key %s", key);
+        getContext().incrementCounter(String.format("error, kind %s", key.getKind()));
+      }
+    }
+
+    private void deleteDomain(final Key<DomainBase> domainKey) {
+      final DomainBase domain = ofy().load().key(domainKey).now();
+      if (domain == null) {
+        // Depending on how stale Datastore indexes are, we can get keys to resources that are
+        // already deleted (e.g. by a recent previous invocation of this mapreduce). So ignore them.
+        getContext().incrementCounter("already deleted");
+        return;
+      }
+      if (domain instanceof DomainApplication) {
+        // Cover the case where we somehow have a domain application with a prober ROID suffix.
+        getContext().incrementCounter("skipped, domain application");
+        return;
+      }
+      if (domain.getFullyQualifiedDomainName().equals("nic." + domain.getTld())) {
+        getContext().incrementCounter("skipped, NIC domain");
+        return;
+      }
+      int dependentsDeleted = ofy().transact(new Work<Integer>() {
+        @Override
+        public Integer run() {
+          EppResourceIndex eppIndex = ofy().load().entity(EppResourceIndex.create(domainKey)).now();
+          verifyNotNull(eppIndex, "Missing EppResourceIndex for domain %s", domain);
+          ForeignKeyIndex<?> fki = ofy().load().key(ForeignKeyDomainIndex.createKey(domain)).now();
+          verifyNotNull(fki, "Missing ForeignKeyDomainIndex for domain %s", domain);
+          // This ancestor query selects all descendant HistoryEntries, BillingEvents, and
+          // PollMessages, as well as the domain itself.
+          List<Key<Object>> domainAndDependentKeys = ofy().load().ancestor(domainKey).keys().list();
+          if (isDryRun) {
+            logger.infofmt(
+                "Would delete the following entities: %s",
+                new ImmutableList.Builder<Object>()
+                    .add(fki)
+                    .add(eppIndex)
+                    .addAll(domainAndDependentKeys)
+                    .build());
+          } else {
+            ofy().deleteWithoutBackup().keys(domainAndDependentKeys);
+            ofy().deleteWithoutBackup().entities(eppIndex, fki);
+          }
+          return domainAndDependentKeys.size() - 1;
+        }
+      });
+      getContext().incrementCounter(String.format("deleted, kind %s", domainKey.getKind()));
+      getContext().incrementCounter("deleted, dependent keys", dependentsDeleted);
+    }
+  }
+}

java/google/registry/batch/ExpandRecurringBillingEventsAction.java

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.billing;
+package google.registry.batch;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Sets.difference;

java/google/registry/module/backend/BUILD

@@ -22,7 +22,6 @@ java_library(
         "//java/google/registry/backup",
         "//java/google/registry/batch",
         "//java/google/registry/bigquery",
-        "//java/google/registry/billing",
         "//java/google/registry/config",
         "//java/google/registry/cron",
         "//java/google/registry/dns",

java/google/registry/module/backend/BackendModule.java

@@ -21,7 +21,7 @@ import static google.registry.request.RequestParameters.extractRequiredParameter
 import com.google.common.base.Optional;
 import dagger.Module;
 import dagger.Provides;
-import google.registry.billing.ExpandRecurringBillingEventsAction;
+import google.registry.batch.ExpandRecurringBillingEventsAction;
 import google.registry.request.Parameter;
 import google.registry.request.RequestParameters;
 import javax.servlet.http.HttpServletRequest;

java/google/registry/module/backend/BackendRequestComponent.java

@@ -21,7 +21,7 @@ import google.registry.backup.DeleteOldCommitLogsAction;
 import google.registry.backup.ExportCommitLogDiffAction;
 import google.registry.backup.RestoreCommitLogsAction;
 import google.registry.batch.DeleteProberDataAction;
-import google.registry.billing.ExpandRecurringBillingEventsAction;
+import google.registry.batch.ExpandRecurringBillingEventsAction;
 import google.registry.cron.CommitLogFanoutAction;
 import google.registry.cron.CronModule;
 import google.registry.cron.TldFanoutAction;

javatests/google/registry/batch/BUILD

@@ -9,23 +9,34 @@ load("//java/com/google/testing/builddefs:GenTestRules.bzl", "GenTestRules")
 
 
 java_library(
-    name = "billing",
+    name = "batch",
     srcs = glob(["*.java"]),
     deps = [
         "//java/com/google/common/base",
         "//java/com/google/common/collect",
+        "//java/com/google/common/io",
+        "//java/com/google/common/net",
         "//third_party/java/appengine:appengine-api-testonly",
-        "//third_party/java/appengine_mapreduce2:appengine_mapreduce",
+        "//third_party/java/appengine:appengine-stubs",
+        "//third_party/java/appengine_gcs_client",
+        "//third_party/java/dagger",
         "//third_party/java/joda_money",
         "//third_party/java/joda_time",
+        "//third_party/java/jsr305_annotations",
         "//third_party/java/junit",
+        "//third_party/java/mockito",
         "//third_party/java/objectify:objectify-v4_1",
+        "//third_party/java/re2j",
+        "//third_party/java/servlet/servlet_api",
         "//third_party/java/truth",
-        "//java/google/registry/billing",
+        "//java/google/registry/batch",
+        "//java/google/registry/config",
+        "//java/google/registry/gcs",
+        "//java/google/registry/groups",
         "//java/google/registry/mapreduce",
         "//java/google/registry/model",
+        "//java/google/registry/request",
         "//java/google/registry/util",
-        "//javatests/google/registry/model",
         "//javatests/google/registry/testing",
         "//javatests/google/registry/testing/mapreduce",
     ],
@@ -36,5 +47,5 @@ GenTestRules(
     default_test_size = "medium",
     shard_count = 3,
     test_files = glob(["*Test.java"]),
-    deps = [":billing"],
+    deps = [":batch"],
 )

javatests/google/registry/batch/DeleteProberDataActionTest.java

@@ -0,0 +1,176 @@
+// Copyright 2016 The Domain Registry Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.batch;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.persistActiveDomain;
+import static google.registry.testing.DatastoreHelper.persistDeletedDomain;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static google.registry.testing.DatastoreHelper.persistSimpleResource;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.mapreduce.MapreduceRunner;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Reason;
+import google.registry.model.domain.DomainResource;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.TldType;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.ExceptionRule;
+import google.registry.testing.FakeResponse;
+import google.registry.testing.mapreduce.MapreduceTestCase;
+import java.util.Set;
+import org.joda.money.Money;
+import org.joda.time.DateTime;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link DeleteProberDataAction}. */
+@RunWith(JUnit4.class)
+public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDataAction> {
+
+  private static final DateTime DELETION_TIME = DateTime.parse("2010-01-01T00:00:00.000Z");
+
+  @Rule
+  public final ExceptionRule thrown = new ExceptionRule();
+
+  @Before
+  public void init() {
+    // Entities in these two should not be touched.
+    createTld("tld", "TLD");
+    // Since "example" doesn't end with .test, its entities won't be deleted even though it is of
+    // TEST type.
+    createTld("example", "EXAMPLE");
+    persistResource(Registry.get("example").asBuilder().setTldType(TldType.TEST).build());
+
+    // Entities in these two should be deleted.
+    createTld("ib-any.test", "IBANYT");
+    persistResource(Registry.get("ib-any.test").asBuilder().setTldType(TldType.TEST).build());
+    createTld("oa-canary.test", "OACANT");
+    persistResource(Registry.get("oa-canary.test").asBuilder().setTldType(TldType.TEST).build());
+
+    action = new DeleteProberDataAction();
+    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.response = new FakeResponse();
+    action.isDryRun = false;
+  }
+
+  private void runMapreduce() throws Exception {
+    action.run();
+    executeTasksUntilEmpty("mapreduce");
+  }
+
+  @Test
+  public void test_deletesAllAndOnlyProberData() throws Exception {
+    Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
+    Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
+    Set<ImmutableObject> ibEntities = persistLotsOfDomains("ib-any.test");
+    Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
+    runMapreduce();
+    assertNotDeleted(tldEntities);
+    assertNotDeleted(exampleEntities);
+    assertDeleted(ibEntities);
+    assertDeleted(oaEntities);
+  }
+
+  @Test
+  public void testSuccess_doesntDeleteNicDomainForProbers() throws Exception {
+    DomainResource nic = persistActiveDomain("nic.ib-any.test");
+    ForeignKeyIndex<DomainResource> fkiNic =
+        ForeignKeyIndex.load(DomainResource.class, "nic.ib-any.test", START_OF_TIME);
+    Set<ImmutableObject> ibEntities = persistLotsOfDomains("ib-any.test");
+    runMapreduce();
+    assertDeleted(ibEntities);
+    assertNotDeleted(ImmutableSet.<ImmutableObject>of(nic, fkiNic));
+  }
+
+  @Test
+  public void testSuccess_dryRunDoesntDeleteData() throws Exception {
+    Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
+    Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
+    action.isDryRun = true;
+    assertNotDeleted(tldEntities);
+    assertNotDeleted(oaEntities);
+  }
+
+  /**
+   * Persists and returns a domain and a descendant history entry, billing event, and poll message,
+   * along with the ForeignKeyIndex and EppResourceIndex.
+   */
+  private static Set<ImmutableObject> persistDomainAndDescendants(String fqdn) {
+    DomainResource domain = persistDeletedDomain(fqdn, DELETION_TIME);
+    HistoryEntry historyEntry = persistSimpleResource(
+        new HistoryEntry.Builder()
+            .setParent(domain)
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .build());
+    BillingEvent.OneTime billingEvent = persistSimpleResource(
+        new BillingEvent.OneTime.Builder()
+            .setParent(historyEntry)
+            .setBillingTime(DELETION_TIME.plusYears(1))
+            .setCost(Money.parse("USD 10"))
+            .setPeriodYears(1)
+            .setReason(Reason.CREATE)
+            .setClientId("TheRegistrar")
+            .setEventTime(DELETION_TIME)
+            .setTargetId(fqdn)
+            .build());
+    PollMessage.OneTime pollMessage = persistSimpleResource(
+        new PollMessage.OneTime.Builder()
+            .setParent(historyEntry)
+            .setEventTime(DELETION_TIME)
+            .setClientId("TheRegistrar")
+            .setMsg("Domain registered")
+            .build());
+    ForeignKeyIndex<DomainResource> fki =
+        ForeignKeyIndex.load(DomainResource.class, fqdn, START_OF_TIME);
+    EppResourceIndex eppIndex =
+        ofy().load().entity(EppResourceIndex.create(Key.create(domain))).now();
+    return ImmutableSet.<ImmutableObject>of(
+        domain, historyEntry, billingEvent, pollMessage, fki, eppIndex);
+  }
+
+  private static Set<ImmutableObject> persistLotsOfDomains(String tld) {
+    ImmutableSet.Builder<ImmutableObject> persistedObjects = new ImmutableSet.Builder<>();
+    for (int i = 0; i < 20; i++) {
+      persistedObjects.addAll(persistDomainAndDescendants(String.format("domain%d.%s", i, tld)));
+    }
+    return persistedObjects.build();
+  }
+
+  private static void assertNotDeleted(Iterable<ImmutableObject> entities) {
+    for (ImmutableObject entity : entities) {
+      assertThat(ofy().load().entity(entity).now()).isNotNull();
+    }
+  }
+
+  private static void assertDeleted(Iterable<ImmutableObject> entities) {
+    for (ImmutableObject entity : entities) {
+      assertThat(ofy().load().entity(entity).now()).isNull();
+    }
+  }
+}

javatests/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.billing;
+package google.registry.batch;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.common.Cursor.CursorType.RECURRING_BILLING;
@@ -96,7 +96,7 @@ public class ExpandRecurringBillingEventsActionTest
         ofy().save().entity(Cursor.createGlobal(RECURRING_BILLING, cursorTime));
       }});
   }
-  
+
   void runMapreduce() throws Exception {
     action.response = new FakeResponse();
     action.run();
@@ -369,7 +369,7 @@ public class ExpandRecurringBillingEventsActionTest
         .setCancellationMatchingBillingEvent(Key.create(recurring))
         .setTargetId(domain.getFullyQualifiedDomainName())
         .build();
-    assertBillingEventsForResource(domain, expected, recurring); 
+    assertBillingEventsForResource(domain, expected, recurring);
     assertCursorAt(clock.nowUtc());
   }
 
@@ -379,7 +379,7 @@ public class ExpandRecurringBillingEventsActionTest
     // Simulate a quick second run of the mapreduce (this should be a no-op).
     saveCursor(clock.nowUtc().minusSeconds(1));
     runMapreduce();
-    assertBillingEventsForResource(domain, recurring); 
+    assertBillingEventsForResource(domain, recurring);
     assertCursorAt(clock.nowUtc());
   }
 
@@ -419,7 +419,7 @@ public class ExpandRecurringBillingEventsActionTest
           .setTargetId(domain.getFullyQualifiedDomainName())
           .build());
     }
-    assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class)); 
+    assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class));
     assertCursorAt(clock.nowUtc());
   }
 
@@ -449,7 +449,7 @@ public class ExpandRecurringBillingEventsActionTest
           .setTargetId(domain.getFullyQualifiedDomainName())
           .build());
     }
-    assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class)); 
+    assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class));
     assertCursorAt(clock.nowUtc());
   }
 
@@ -606,7 +606,7 @@ public class ExpandRecurringBillingEventsActionTest
         .setCancellationMatchingBillingEvent(Key.create(recurring))
         .setTargetId(domain.getFullyQualifiedDomainName())
         .build();
-    assertBillingEventsForResource(domain, expected, recurring); 
+    assertBillingEventsForResource(domain, expected, recurring);
     assertCursorAt(clock.nowUtc());
   }
 
@@ -645,7 +645,7 @@ public class ExpandRecurringBillingEventsActionTest
         .setBillingTime(billingDate.plusYears(1))
         .setEventTime(eventDate.plusYears(1))
         .build();
-    assertBillingEventsForResource(domain, recurring, cheaper, expensive); 
+    assertBillingEventsForResource(domain, recurring, cheaper, expensive);
     assertCursorAt(clock.nowUtc());
   }