diff --git a/java/google/registry/export/BUILD b/java/google/registry/export/BUILD index 33644fb92..9e4aa2c43 100644 --- a/java/google/registry/export/BUILD +++ b/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", diff --git a/java/google/registry/export/DriveModule.java b/java/google/registry/export/DriveModule.java index 8ee8bf8a3..72889417a 100644 --- a/java/google/registry/export/DriveModule.java +++ b/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(); + } } diff --git a/java/google/registry/export/ExportDomainListsAction.java b/java/google/registry/export/ExportDomainListsAction.java index 74fa257f9..d0eb808f1 100644 --- a/java/google/registry/export/ExportDomainListsAction.java +++ b/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 - * TLD.txt into the domain-lists bucket. Note that this overwrites the files in place. + *

Each TLD's active domain names are exported as a newline-delimited flat text file with the + * name TLD.txt into the domain-lists bucket. Note that this overwrites the files in place. */ -@Action( - path = "/_dr/task/exportDomainLists", - method = POST, - auth = Auth.AUTH_INTERNAL_ONLY -) +@Action(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 - public void reduce(String tld, ReducerInput fqdns) { + private void exportToDrive(String tld, String domains) { + 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); - PrintWriter writer = new PrintWriter(osWriter)) { - 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); + Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8)) { + osWriter.write(domains); } 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 fqdns) { + ImmutableList 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); } } } diff --git a/javatests/google/registry/export/ExportDomainListsActionTest.java b/javatests/google/registry/export/ExportDomainListsActionTest.java index 9169e1122..e08ba14cd 100644 --- a/javatests/google/registry/export/ExportDomainListsActionTest.java +++ b/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 { private GcsService gcsService; + private DriveConnection driveConnection = mock(DriveConnection.class); + private ArgumentCaptor 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 readGcsFile(gcsService, nonexistentFile)); @@ -95,25 +120,37 @@ public class ExportDomainListsActionTest extends MapreduceTestCase