Convert GenerateZoneFilesAction to SQL (#1668)

I'm not 100% sure that this is strictly necessary, but for now we can
replicate the ability to generate zonefiles for any point in time in the
recent past.

core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java

@@ -15,27 +15,21 @@
 package google.registry.tools.server;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static com.google.common.collect.Iterators.filter;
 import static com.google.common.io.BaseEncoding.base16;
-import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.loadAtPointInTime;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.joda.time.DateTimeZone.UTC;
 
-import com.google.appengine.tools.mapreduce.Mapper;
-import com.google.appengine.tools.mapreduce.Reducer;
-import com.google.appengine.tools.mapreduce.ReducerInput;
 import com.google.cloud.storage.BlobId;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.gcs.GcsUtils;
 import google.registry.mapreduce.MapreduceRunner;
-import google.registry.mapreduce.inputs.NullInput;
-import google.registry.model.EppResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.host.HostResource;
@@ -51,11 +45,13 @@ import java.io.PrintWriter;
 import java.io.Writer;
 import java.net.Inet4Address;
 import java.net.InetAddress;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import javax.inject.Inject;
+import org.hibernate.CacheMode;
+import org.hibernate.ScrollMode;
+import org.hibernate.ScrollableResults;
+import org.hibernate.query.Query;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
@@ -73,8 +69,13 @@ import org.joda.time.Duration;
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
 public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonAction {
 
+  private static final FluentLogger log = FluentLogger.forEnclosingClass();
+
   public static final String PATH = "/_dr/task/generateZoneFiles";
 
+  /** Number of domains to process in one batch. */
+  private static final int BATCH_SIZE = 1000;
+
   /** Format for the zone file name. */
   private static final String FILENAME_FORMAT = "%s-%s.zone";
 
@@ -128,20 +129,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
           "Invalid export time: must be < %d days ago",
           datastoreRetention.getStandardDays()));
     }
-    if (!exportTime.equals(exportTime.toDateTime(UTC).withTimeAtStartOfDay())) {
+    tlds.forEach(tld -> generateForTld(tld, exportTime));
-      throw new BadRequestException("Invalid export time: must be midnight UTC");
-    }
-    String mapreduceConsoleLink =
-        mrRunner
-            .setJobName("Generate bind file stanzas")
-            .setModuleName("tools")
-            .setDefaultReduceShards(tlds.size())
-            .runMapreduce(
-                new GenerateBindFileMapper(
-                    tlds, exportTime, dnsDefaultATtl, dnsDefaultNsTtl, dnsDefaultDsTtl),
-                new GenerateBindFileReducer(bucket, exportTime, gcsUtils),
-                ImmutableList.of(new NullInput<>(), createEntityInput(DomainBase.class)))
-            .getLinkToMapreduceConsole();
     ImmutableList<String> filenames =
         tlds.stream()
             .map(
@@ -150,115 +138,79 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
                         GCS_PATH_FORMAT, bucket, String.format(FILENAME_FORMAT, tld, exportTime)))
             .collect(toImmutableList());
     return ImmutableMap.of(
-        "mapreduceConsoleLink", mapreduceConsoleLink,
         "filenames", filenames);
   }
 
-  /** Mapper to find domains that were active at a given time. */
+  private void generateForTld(String tld, DateTime exportTime) {
-  static class GenerateBindFileMapper extends Mapper<EppResource, String, String> {
+    ImmutableList<String> stanzas = jpaTm().transact(() -> getStanzasForTld(tld, exportTime));
-
+    BlobId outputBlobId = BlobId.of(bucket, String.format(FILENAME_FORMAT, tld, exportTime));
-    private static final long serialVersionUID = 4647941823789859913L;
+    try (OutputStream gcsOutput = gcsUtils.openOutputStream(outputBlobId);
-
+        Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8);
-    private final ImmutableSet<String> tlds;
+        PrintWriter writer = new PrintWriter(osWriter)) {
-    private final DateTime exportTime;
+      writer.printf(HEADER_FORMAT, tld);
-    private final Duration dnsDefaultATtl;
+      stanzas.forEach(writer::println);
-    private final Duration dnsDefaultNsTtl;
+      writer.flush();
-    private final Duration dnsDefaultDsTtl;
+    } catch (IOException e) {
-
+      throw new RuntimeException(e);
-    GenerateBindFileMapper(
-        ImmutableSet<String> tlds,
-        DateTime exportTime,
-        Duration dnsDefaultATtl,
-        Duration dnsDefaultNsTtl,
-        Duration dnsDefaultDsTtl) {
-      this.tlds = tlds;
-      this.exportTime = exportTime;
-      this.dnsDefaultATtl = dnsDefaultATtl;
-      this.dnsDefaultNsTtl = dnsDefaultNsTtl;
-      this.dnsDefaultDsTtl = dnsDefaultDsTtl;
-    }
-
-    @Override
-    public void map(EppResource resource) {
-      if (resource == null) {  // Force the reducer to always generate a bind header for each tld.
-        for (String tld : tlds) {
-          emit(tld, null);
-        }
-      } else {
-        mapDomain((DomainBase) resource);
-      }
-    }
-
-    // Originally, we mapped over domains and hosts separately, emitting the necessary information
-    // for each. But that doesn't work. All subordinate hosts in the specified TLD(s) would always
-    // be emitted in the final file, which is incorrect. Rather, to match the actual DNS glue
-    // records, we only want to emit host information for in-bailiwick hosts in the specified
-    // TLD(s), meaning those that act as nameservers for their respective superordinate domains.
-    private void mapDomain(DomainBase domain) {
-      // Domains never change their tld, so we can check if it's from the wrong tld right away.
-      if (tlds.contains(domain.getTld())) {
-        domain = loadAtPointInTime(domain, exportTime);
-        // A null means the domain was deleted (or not created) at this time.
-        if (domain != null && domain.shouldPublishToDns()) {
-          String stanza = domainStanza(domain, exportTime, dnsDefaultNsTtl, dnsDefaultDsTtl);
-          if (!stanza.isEmpty()) {
-            emit(domain.getTld(), stanza);
-            getContext().incrementCounter(domain.getTld() + " domains");
-          }
-          emitForSubordinateHosts(domain);
-        }
-      }
-    }
-
-    private void emitForSubordinateHosts(DomainBase domain) {
-      ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
-      if (!subordinateHosts.isEmpty()) {
-        for (HostResource unprojectedHost : tm().loadByKeys(domain.getNameservers()).values()) {
-          HostResource host = loadAtPointInTime(unprojectedHost, exportTime);
-          // A null means the host was deleted (or not created) at this time.
-          if ((host != null) && subordinateHosts.contains(host.getHostName())) {
-            String stanza = hostStanza(host, dnsDefaultATtl, domain.getTld());
-            if (!stanza.isEmpty()) {
-              emit(domain.getTld(), stanza);
-              getContext().incrementCounter(domain.getTld() + " hosts");
-            }
-          }
-        }
-      }
     }
   }
 
-  /** Reducer to write zone files to GCS. */
+  private ImmutableList<String> getStanzasForTld(String tld, DateTime exportTime) {
-  static class GenerateBindFileReducer extends Reducer<String, String, Void> {
+    ImmutableList.Builder<String> result = new ImmutableList.Builder<>();
-
+    ScrollableResults scrollableResults =
-    private static final long serialVersionUID = -8489050680083119352L;
+        jpaTm()
-
+            .query("FROM Domain WHERE tld = :tld AND deletionTime > :exportTime")
-    private final String bucket;
+            .setParameter("tld", tld)
-    private final DateTime exportTime;
+            .setParameter("exportTime", exportTime)
-    private final GcsUtils gcsUtils;
+            .unwrap(Query.class)
-
+            .setCacheMode(CacheMode.IGNORE)
-    GenerateBindFileReducer(String bucket, DateTime exportTime, GcsUtils gcsUtils) {
+            .scroll(ScrollMode.FORWARD_ONLY);
-      this.bucket = bucket;
+    for (int i = 1; scrollableResults.next(); i = (i + 1) % BATCH_SIZE) {
-      this.exportTime = exportTime;
+      DomainBase domain = (DomainBase) scrollableResults.get(0);
-      this.gcsUtils = gcsUtils;
+      populateStanzasForDomain(domain, exportTime, result);
+      if (i == 0) {
+        jpaTm().getEntityManager().flush();
+        jpaTm().getEntityManager().clear();
+      }
     }
+    return result.build();
+  }
 
-    @Override
+  private void populateStanzasForDomain(
-    public void reduce(String tld, ReducerInput<String> stanzas) {
+      DomainBase domain, DateTime exportTime, ImmutableList.Builder<String> result) {
-      String stanzaCounter = tld + " stanzas";
+    domain = loadAtPointInTime(domain, exportTime);
-      BlobId filename = BlobId.of(bucket, String.format(FILENAME_FORMAT, tld, exportTime));
+    // A null means the domain was deleted (or not created) at this time.
-      try (OutputStream gcsOutput = gcsUtils.openOutputStream(filename);
+    if (domain == null || !domain.shouldPublishToDns()) {
-          Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8);
+      return;
-          PrintWriter writer = new PrintWriter(osWriter)) {
+    }
-        writer.printf(HEADER_FORMAT, tld);
+    String stanza = domainStanza(domain, exportTime);
-        for (Iterator<String> stanzaIter = filter(stanzas, Objects::nonNull);
+    if (!stanza.isEmpty()) {
-            stanzaIter.hasNext(); ) {
+      result.add(stanza);
-          writer.println(stanzaIter.next());
+    }
-          getContext().incrementCounter(stanzaCounter);
+    populateStanzasForSubordinateHosts(domain, exportTime, result);
+  }
+
+  private void populateStanzasForSubordinateHosts(
+      DomainBase domain, DateTime exportTime, ImmutableList.Builder<String> result) {
+    ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
+    if (!subordinateHosts.isEmpty()) {
+      for (HostResource unprojectedHost : tm().loadByKeys(domain.getNameservers()).values()) {
+        HostResource host = loadAtPointInTime(unprojectedHost, exportTime);
+        // A null means the host was deleted (or not created) at this time.
+        if (host != null && subordinateHosts.contains(host.getHostName())) {
+          String stanza = hostStanza(host, domain.getTld());
+          if (!stanza.isEmpty()) {
+            result.add(stanza);
+          }
+        } else if (host == null) {
+          log.atSevere().log(
+              "Domain %s contained nameserver %s that didn't exist at time %s",
+              domain.getRepoId(), unprojectedHost.getRepoId(), exportTime);
+        } else {
+          log.atSevere().log(
+              "Domain %s contained nameserver %s not in subordinate hosts at time %s",
+              domain.getRepoId(), unprojectedHost.getRepoId(), exportTime);
         }
-        writer.flush();
-      } catch (IOException e) {
-        throw new RuntimeException(e);
       }
     }
   }
@@ -266,17 +218,16 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
   /**
    * Generates DNS records for a domain (NS and DS).
    *
-   * For domain foo.tld, these look like this:
+   * <p>For domain foo.tld, these look like this:
-   * {@code
+   *
+   * <pre>
+   * {
    *   foo 180 IN NS ns.example.com.
    *   foo 86400 IN DS 1 2 3 000102
    * }
+   * </pre>
    */
-  private static String domainStanza(
+  private String domainStanza(DomainBase domain, DateTime exportTime) {
-      DomainBase domain,
-      DateTime exportTime,
-      Duration dnsDefaultNsTtl,
-      Duration dnsDefaultDsTtl) {
     StringBuilder result = new StringBuilder();
     String domainLabel = stripTld(domain.getDomainName(), domain.getTld());
     for (HostResource nameserver : tm().loadByKeys(domain.getNameservers()).values()) {
@@ -306,12 +257,15 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
    * Generates DNS records for a domain (A and AAAA).
    *
    * <p>These look like this:
-   * {@code
+   *
+   * <pre>
+   * {
    *   ns.foo.tld 3600 IN A 127.0.0.1
    *   ns.foo.tld 3600 IN AAAA 0:0:0:0:0:0:0:1
    * }
+   * </pre>
    */
-  private static String hostStanza(HostResource host, Duration dnsDefaultATtl, String tld) {
+  private String hostStanza(HostResource host, String tld) {
     StringBuilder result = new StringBuilder();
     for (InetAddress addr : host.getInetAddresses()) {
       // must be either IPv4 or IPv6

core/src/test/java/google/registry/tools/server/GenerateZoneFilesActionTest.java

@@ -37,28 +37,32 @@ import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
-import google.registry.testing.TmOverrideExtension;
+import google.registry.testing.TestSqlOnly;
-import google.registry.testing.mapreduce.MapreduceTestCase;
 import java.net.InetAddress;
 import java.util.Map;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.Duration;
-import org.junit.jupiter.api.Order;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Tests for {@link GenerateZoneFilesAction}. */
-class GenerateZoneFilesActionTest extends MapreduceTestCase<GenerateZoneFilesAction> {
+@DualDatabaseTest
+class GenerateZoneFilesActionTest {
 
   @RegisterExtension
-  @Order(Order.DEFAULT - 1)
+  public final AppEngineExtension appEngine =
-  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withOfy();
+      AppEngineExtension.builder()
+          .withDatastoreAndCloudSql()
+          .withLocalModules()
+          .withTaskQueue()
+          .build();
 
   private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
 
-  @Test
+  @TestSqlOnly
   void testGenerate() throws Exception {
     DateTime now = DateTime.now(DateTimeZone.UTC).withTimeAtStartOfDay();
     createTlds("tld", "com");
@@ -118,7 +122,6 @@ class GenerateZoneFilesActionTest extends MapreduceTestCase<GenerateZoneFilesAct
         .build());
 
     GenerateZoneFilesAction action = new GenerateZoneFilesAction();
-    action.mrRunner = makeDefaultRunner();
     action.bucket = "zonefiles-bucket";
     action.gcsUtils = gcsUtils;
     action.datastoreRetention = standardDays(29);
@@ -132,13 +135,6 @@ class GenerateZoneFilesActionTest extends MapreduceTestCase<GenerateZoneFilesAct
             ImmutableMap.<String, Object>of("tlds", ImmutableList.of("tld"), "exportTime", now));
     assertThat(response)
         .containsEntry("filenames", ImmutableList.of("gs://zonefiles-bucket/tld-" + now + ".zone"));
-    assertThat(response).containsKey("mapreduceConsoleLink");
-    assertThat(response.get("mapreduceConsoleLink").toString())
-        .startsWith(
-            "Mapreduce console: https://backend-dot-projectid.appspot.com"
-                + "/_ah/pipeline/status.html?root=");
-
-    executeTasksUntilEmpty("mapreduce");
 
     BlobId gcsFilename = BlobId.of("zonefiles-bucket", String.format("tld-%s.zone", now));
     String generatedFile = new String(gcsUtils.readBytesFrom(gcsFilename), UTF_8);