Export registered domains to drive folder

The export happens in a reducer, whois instance fields all need to be serializable. The DriveConnection is therefore installed as a class variable that can be replaced with a mock during test. Class variables are not serialized. Only fields related to a particular instance are.

Note that DriveConnection is a misnomer. It is not a connection at all. It is just a thin wrapper class around the Drive service class, which provides convenient methods to write into Drive. Regardless, it cannot be serialized.

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

java/google/registry/export/BUILD

@@ -16,6 +16,7 @@ java_library(
         "//java/google/registry/mapreduce/inputs",
         "//java/google/registry/model",
         "//java/google/registry/request",
+        "//java/google/registry/request:modules",
         "//java/google/registry/request/auth",
         "//java/google/registry/storage/drive",
         "//java/google/registry/util",

java/google/registry/export/DriveModule.java

@@ -19,11 +19,19 @@ import com.google.api.client.http.HttpTransport;
 import com.google.api.client.json.JsonFactory;
 import com.google.api.services.drive.Drive;
 import com.google.api.services.drive.DriveScopes;
+import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryConfig.ConfigModule;
+import google.registry.request.Modules.AppIdentityCredentialModule;
+import google.registry.request.Modules.Jackson2Module;
+import google.registry.request.Modules.UrlFetchTransportModule;
+import google.registry.request.Modules.UseAppIdentityCredentialForGoogleApisModule;
+import google.registry.storage.drive.DriveConnection;
 import java.util.Set;
 import java.util.function.Function;
+import javax.inject.Singleton;
 
 /** Dagger module for Google {@link Drive} service connection objects. */
 @Module
@@ -39,4 +47,18 @@ public final class DriveModule {
         .setApplicationName(projectId)
         .build();
   }
+
+  @Singleton
+  @Component(
+      modules = {
+        DriveModule.class,
+        UrlFetchTransportModule.class,
+        Jackson2Module.class,
+        AppIdentityCredentialModule.class,
+        UseAppIdentityCredentialForGoogleApisModule.class,
+        ConfigModule.class
+      })
+  interface DriveComponent {
+    DriveConnection driveConnection();
+  }
 }

java/google/registry/export/ExportDomainListsAction.java

@@ -28,21 +28,25 @@ import com.google.appengine.tools.cloudstorage.RetryParams;
 import com.google.appengine.tools.mapreduce.Mapper;
 import com.google.appengine.tools.mapreduce.Reducer;
 import com.google.appengine.tools.mapreduce.ReducerInput;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.net.MediaType;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.gcs.GcsUtils;
 import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.domain.DomainResource;
+import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.request.Action;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.storage.drive.DriveConnection;
 import google.registry.util.FormattingLogger;
+import google.registry.util.NonFinalForTesting;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.io.OutputStreamWriter;
-import java.io.PrintWriter;
 import java.io.Writer;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
@@ -50,14 +54,10 @@ import org.joda.time.DateTime;
 /**
  * A mapreduce that exports the list of active domains on all real TLDs to Google Cloud Storage.
  *
- * Each TLD's active domain names are exported as a newline-delimited flat text file with the name
+ * <p>Each TLD's active domain names are exported as a newline-delimited flat text file with the
- * TLD.txt into the domain-lists bucket.  Note that this overwrites the files in place.
+ * name TLD.txt into the domain-lists bucket. Note that this overwrites the files in place.
  */
-@Action(
+@Action(path = "/_dr/task/exportDomainLists", method = POST, auth = Auth.AUTH_INTERNAL_ONLY)
-  path = "/_dr/task/exportDomainLists",
-  method = POST,
-  auth = Auth.AUTH_INTERNAL_ONLY
-)
 public class ExportDomainListsAction implements Runnable {
 
   private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
@@ -108,32 +108,69 @@ public class ExportDomainListsAction implements Runnable {
 
     private static final long serialVersionUID = 7035260977259119087L;
 
+    @NonFinalForTesting
+    private static DriveConnection driveConnection =
+        DaggerDriveModule_DriveComponent.create().driveConnection();
+
+    static final String REGISTERED_DOMAINS_FILENAME = "registered_domains.txt";
+    static final MediaType EXPORT_MIME_TYPE = MediaType.PLAIN_TEXT_UTF_8;
+
     private final String gcsBucket;
     private final int gcsBufferSize;
 
+    static void setDriveConnectionForTesting(DriveConnection driveConnection) {
+      ExportDomainListsReducer.driveConnection = driveConnection;
+    }
+
     public ExportDomainListsReducer(String gcsBucket, int gcsBufferSize) {
       this.gcsBucket = gcsBucket;
       this.gcsBufferSize = gcsBufferSize;
     }
 
-    @Override
+    private void exportToDrive(String tld, String domains) {
-    public void reduce(String tld, ReducerInput<String> fqdns) {
+      try {
+        Registry registry = Registry.get(tld);
+        if (registry.getDriveFolderId() == null) {
+          logger.infofmt(
+              "Skipping registered domains export for TLD %s because Drive folder isn't specified",
+              tld);
+        } else {
+          String resultMsg =
+              driveConnection.createOrUpdateFile(
+                  REGISTERED_DOMAINS_FILENAME,
+                  EXPORT_MIME_TYPE,
+                  registry.getDriveFolderId(),
+                  domains.getBytes(UTF_8));
+          logger.infofmt(
+              "Exporting registered domains succeeded for TLD %s, response was: %s",
+              tld, resultMsg);
+        }
+      } catch (Throwable e) {
+        logger.severefmt(e, "Error exporting registered domains for TLD %s to Drive", tld);
+      }
+      getContext().incrementCounter("domain lists written out to Drive");
+    }
+
+    private void exportToGcs(String tld, String domains) {
       GcsFilename filename = new GcsFilename(gcsBucket, tld + ".txt");
       GcsUtils cloudStorage =
           new GcsUtils(createGcsService(RetryParams.getDefaultInstance()), gcsBufferSize);
       try (OutputStream gcsOutput = cloudStorage.openOutputStream(filename);
-          Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8);
+          Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8)) {
-          PrintWriter writer = new PrintWriter(osWriter)) {
+        osWriter.write(domains);
-        long count;
-        for (count = 0; fqdns.hasNext(); count++) {
-          writer.println(fqdns.next());
-        }
-        writer.flush();
-        getContext().incrementCounter("tld domain lists written out");
-        logger.infofmt("Wrote out %d domains for tld %s.", count, tld);
       } catch (IOException e) {
-        throw new RuntimeException(e);
+        logger.severefmt(e, "Error exporting registered domains for TLD %s to GCS.", tld);
       }
+      getContext().incrementCounter("domain lists written out to GCS");
+    }
+
+    @Override
+    public void reduce(String tld, ReducerInput<String> fqdns) {
+      ImmutableList<String> domains = ImmutableList.sortedCopyOf(() -> fqdns);
+      String domainsList = Joiner.on('\n').join(domains);
+      logger.infofmt("Exporting %d domains for TLD %s to GCS and Drive.", domains.size(), tld);
+      exportToGcs(tld, domainsList);
+      exportToDrive(tld, domainsList);
     }
   }
 }

javatests/google/registry/export/ExportDomainListsActionTest.java

@@ -16,6 +16,8 @@ package google.registry.export;
 
 import static com.google.appengine.tools.cloudstorage.GcsServiceFactory.createGcsService;
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.export.ExportDomainListsAction.ExportDomainListsReducer.EXPORT_MIME_TYPE;
+import static google.registry.export.ExportDomainListsAction.ExportDomainListsReducer.REGISTERED_DOMAINS_FILENAME;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistActiveDomain;
 import static google.registry.testing.DatastoreHelper.persistActiveDomainApplication;
@@ -24,14 +26,19 @@ import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.GcsTestingUtils.readGcsFile;
 import static google.registry.testing.JUnitBackports.assertThrows;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.appengine.tools.cloudstorage.ListOptions;
 import com.google.appengine.tools.cloudstorage.ListResult;
-import com.google.common.base.Splitter;
+import google.registry.export.ExportDomainListsAction.ExportDomainListsReducer;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
+import google.registry.storage.drive.DriveConnection;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import java.io.FileNotFoundException;
@@ -40,19 +47,25 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
 
 /** Unit tests for {@link ExportDomainListsAction}. */
 @RunWith(JUnit4.class)
 public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAction> {
 
   private GcsService gcsService;
+  private DriveConnection driveConnection = mock(DriveConnection.class);
+  private ArgumentCaptor<byte[]> bytesExportedToDrive = ArgumentCaptor.forClass(byte[].class);
 
   @Before
   public void init() {
     createTld("tld");
     createTld("testtld");
+    persistResource(Registry.get("tld").asBuilder().setDriveFolderId("brouhaha").build());
     persistResource(Registry.get("testtld").asBuilder().setTldType(TldType.TEST).build());
 
+    ExportDomainListsReducer.setDriveConnectionForTesting(driveConnection);
+
     action = new ExportDomainListsAction();
     action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
@@ -66,6 +79,16 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     executeTasksUntilEmpty("mapreduce");
   }
 
+  private void verifyExportedToDrive(String folderId, String domains) throws Exception {
+    verify(driveConnection)
+        .createOrUpdateFile(
+            eq(REGISTERED_DOMAINS_FILENAME),
+            eq(EXPORT_MIME_TYPE),
+            eq(folderId),
+            bytesExportedToDrive.capture());
+    assertThat(new String(bytesExportedToDrive.getValue(), "UTF-8")).isEqualTo(domains);
+  }
+
   @Test
   public void test_outputsOnlyActiveDomains() throws Exception {
     persistActiveDomain("onetwo.tld");
@@ -73,9 +96,11 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     persistDeletedDomain("mortuary.tld", DateTime.parse("2001-03-14T10:11:12Z"));
     runMapreduce();
     GcsFilename existingFile = new GcsFilename("outputbucket", "tld.txt");
-    String tlds = new String(readGcsFile(gcsService, existingFile), UTF_8).trim();
+    String tlds = new String(readGcsFile(gcsService, existingFile), UTF_8);
     // Check that it only contains the active domains, not the dead one.
-    assertThat(Splitter.on('\n').splitToList(tlds)).containsExactly("onetwo.tld", "rudnitzky.tld");
+    assertThat(tlds).isEqualTo("onetwo.tld\nrudnitzky.tld");
+    verifyExportedToDrive("brouhaha", "onetwo.tld\nrudnitzky.tld");
+    verifyNoMoreInteractions(driveConnection);
   }
 
   @Test
@@ -87,7 +112,7 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     GcsFilename existingFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, existingFile), UTF_8).trim();
     // Check that it only contains the domains on the real TLD, and not the test one.
-    assertThat(Splitter.on('\n').splitToList(tlds)).containsExactly("onetwo.tld", "rudnitzky.tld");
+    assertThat(tlds).isEqualTo("onetwo.tld\nrudnitzky.tld");
     // Make sure that the test TLD file wasn't written out.
     GcsFilename nonexistentFile = new GcsFilename("outputbucket", "testtld.txt");
     assertThrows(FileNotFoundException.class, () -> readGcsFile(gcsService, nonexistentFile));
@@ -95,25 +120,37 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     assertThat(ls.next().getName()).isEqualTo("tld.txt");
     // Make sure that no other files were written out.
     assertThat(ls.hasNext()).isFalse();
+    verifyExportedToDrive("brouhaha", "onetwo.tld\nrudnitzky.tld");
+    verifyNoMoreInteractions(driveConnection);
   }
 
   @Test
   public void test_outputsDomainsFromDifferentTldsToMultipleFiles() throws Exception {
     createTld("tldtwo");
+    persistResource(Registry.get("tldtwo").asBuilder().setDriveFolderId("hooray").build());
+
+    createTld("tldthree");
     // You'd think this test was written around Christmas, but it wasn't.
     persistActiveDomain("dasher.tld");
     persistActiveDomain("prancer.tld");
     persistActiveDomain("rudolph.tldtwo");
     persistActiveDomain("santa.tldtwo");
     persistActiveDomain("buddy.tldtwo");
+    persistActiveDomain("cupid.tldthree");
     runMapreduce();
     GcsFilename firstTldFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, firstTldFile), UTF_8).trim();
-    assertThat(Splitter.on('\n').splitToList(tlds)).containsExactly("dasher.tld", "prancer.tld");
+    assertThat(tlds).isEqualTo("dasher.tld\nprancer.tld");
     GcsFilename secondTldFile = new GcsFilename("outputbucket", "tldtwo.txt");
     String moreTlds = new String(readGcsFile(gcsService, secondTldFile), UTF_8).trim();
-    assertThat(Splitter.on('\n').splitToList(moreTlds))
+    assertThat(moreTlds).isEqualTo("buddy.tldtwo\nrudolph.tldtwo\nsanta.tldtwo");
-        .containsExactly("rudolph.tldtwo", "santa.tldtwo", "buddy.tldtwo");
+    GcsFilename thirdTldFile = new GcsFilename("outputbucket", "tldthree.txt");
+    String evenMoreTlds = new String(readGcsFile(gcsService, thirdTldFile), UTF_8).trim();
+    assertThat(evenMoreTlds).isEqualTo("cupid.tldthree");
+    verifyExportedToDrive("brouhaha", "dasher.tld\nprancer.tld");
+    verifyExportedToDrive("hooray", "buddy.tldtwo\nrudolph.tldtwo\nsanta.tldtwo");
+    // tldthree does not have a drive id, so no export to drive is performed.
+    verifyNoMoreInteractions(driveConnection);
   }
 
   @Test
@@ -124,6 +161,8 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     GcsFilename firstTldFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, firstTldFile), UTF_8).trim();
     // Check that it didn't output nagajolokia.tld.
-    assertThat(Splitter.on('\n').splitToList(tlds)).containsExactly("chilipepper.tld");
+    assertThat(tlds).isEqualTo("chilipepper.tld");
+    verifyExportedToDrive("brouhaha", "chilipepper.tld");
+    verifyNoMoreInteractions(driveConnection);
   }
 }