From d113b718d7ebd3781d97355ed087dbf096ccc4e8 Mon Sep 17 00:00:00 2001 From: Lai Jiang Date: Thu, 10 Jun 2021 12:03:17 -0400 Subject: [PATCH] Make ExportDomainListsAction SQL-aware (#1195) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This change is [Reviewable](https://reviewable.io/reviews/google/nomulus/1195) --- .../export/ExportDomainListsAction.java | 166 ++++++++++++------ .../export/ExportDomainListsActionTest.java | 42 +++-- 2 files changed, 140 insertions(+), 68 deletions(-) diff --git a/core/src/main/java/google/registry/export/ExportDomainListsAction.java b/core/src/main/java/google/registry/export/ExportDomainListsAction.java index 3eb5fcd37..6b69fdb83 100644 --- a/core/src/main/java/google/registry/export/ExportDomainListsAction.java +++ b/core/src/main/java/google/registry/export/ExportDomainListsAction.java @@ -19,9 +19,10 @@ import static com.google.common.base.Verify.verifyNotNull; import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput; import static google.registry.model.EppResourceUtils.isActive; import static google.registry.model.registry.Registries.getTldsOfType; +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.cloudstorage.GcsFilename; import com.google.appengine.tools.cloudstorage.RetryParams; @@ -45,12 +46,14 @@ 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.Clock; import google.registry.util.NonFinalForTesting; import java.io.IOException; import java.io.ObjectInputStream; import java.io.OutputStream; import java.io.OutputStreamWriter; import java.io.Writer; +import java.util.List; import java.util.function.Supplier; import javax.inject.Inject; import org.joda.time.DateTime; @@ -70,9 +73,13 @@ public class ExportDomainListsAction implements Runnable { private static final FluentLogger logger = FluentLogger.forEnclosingClass(); private static final int MAX_NUM_REDUCE_SHARDS = 100; + public static final String REGISTERED_DOMAINS_FILENAME = "registered_domains.txt"; @Inject MapreduceRunner mrRunner; @Inject Response response; + @Inject Clock clock; + @Inject DriveConnection driveConnection; + @Inject @Config("domainListsGcsBucket") String gcsBucket; @Inject @Config("gcsBufferSize") int gcsBufferSize; @Inject ExportDomainListsAction() {} @@ -81,15 +88,99 @@ public class ExportDomainListsAction implements Runnable { public void run() { ImmutableSet realTlds = getTldsOfType(TldType.REAL); logger.atInfo().log("Exporting domain lists for tlds %s", realTlds); - mrRunner - .setJobName("Export domain lists") - .setModuleName("backend") - .setDefaultReduceShards(Math.min(realTlds.size(), MAX_NUM_REDUCE_SHARDS)) - .runMapreduce( - new ExportDomainListsMapper(DateTime.now(UTC), realTlds), - new ExportDomainListsReducer(gcsBucket, gcsBufferSize), - ImmutableList.of(createEntityInput(DomainBase.class))) - .sendLinkToMapreduceConsole(response); + if (tm().isOfy()) { + mrRunner + .setJobName("Export domain lists") + .setModuleName("backend") + .setDefaultReduceShards(Math.min(realTlds.size(), MAX_NUM_REDUCE_SHARDS)) + .runMapreduce( + new ExportDomainListsMapper(clock.nowUtc(), realTlds), + new ExportDomainListsReducer(gcsBucket, gcsBufferSize), + ImmutableList.of(createEntityInput(DomainBase.class))) + .sendLinkToMapreduceConsole(response); + } else { + realTlds.forEach( + tld -> { + List domains = + tm().transact( + () -> + // Note that if we had "creationTime <= :now" in the condition (not + // necessary as there is no pending creation, the order of deletionTime + // and creationTime in the query would have been significant and it + // should come after deletionTime. When Hibernate substitutes "now" it + // will first validate that the **first** field that is to be compared + // with it (deletionTime) is assignable from the substituted Java object + // (click.nowUtc()). Since creationTime is a CreateAutoTimestamp, if it + // comes first, we will need to substitute "now" with + // CreateAutoTimestamp.create(clock.nowUtc()). This might look a bit + // strange as the Java object type is clearly incompatible between the + // two fields deletionTime (DateTime) and creationTime, yet they are + // compared with the same "now". It is actually OK because in the end + // Hibernate converts everything to SQL types (and Java field names to + // SQL column names) to run the query. Both CreateAutoTimestamp and + // DateTime are persisted as timestamp_z in SQL. It is only the + // validation that compares the Java types, and only with the first + // field that compares with the substituted value. + jpaTm() + .query( + "SELECT fullyQualifiedDomainName FROM Domain " + + "WHERE tld = :tld " + + "AND deletionTime > :now " + + "ORDER by fullyQualifiedDomainName ASC", + String.class) + .setParameter("tld", tld) + .setParameter("now", clock.nowUtc()) + .getResultList()); + String domainsList = Joiner.on("\n").join(domains); + logger.atInfo().log( + "Exporting %d domains for TLD %s to GCS and Drive.", domains.size(), tld); + exportToGcs(tld, domainsList, gcsBucket, gcsBufferSize); + exportToDrive(tld, domainsList, driveConnection); + }); + } + } + + protected static boolean exportToDrive( + String tld, String domains, DriveConnection driveConnection) { + verifyNotNull(driveConnection, "Expecting non-null driveConnection"); + try { + Registry registry = Registry.get(tld); + if (registry.getDriveFolderId() == null) { + logger.atInfo().log( + "Skipping registered domains export for TLD %s because Drive folder isn't specified", + tld); + } else { + String resultMsg = + driveConnection.createOrUpdateFile( + REGISTERED_DOMAINS_FILENAME, + MediaType.PLAIN_TEXT_UTF_8, + registry.getDriveFolderId(), + domains.getBytes(UTF_8)); + logger.atInfo().log( + "Exporting registered domains succeeded for TLD %s, response was: %s", tld, resultMsg); + } + } catch (Throwable e) { + logger.atSevere().withCause(e).log( + "Error exporting registered domains for TLD %s to Drive, skipping...", tld); + return false; + } + return true; + } + + protected static boolean exportToGcs( + String tld, String domains, String gcsBucket, int gcsBufferSize) { + 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)) { + osWriter.write(domains); + } catch (Throwable e) { + logger.atSevere().withCause(e).log( + "Error exporting registered domains for TLD %s to GCS, skipping...", tld); + return false; + } + return true; } static class ExportDomainListsMapper extends Mapper { @@ -122,9 +213,6 @@ public class ExportDomainListsAction implements Runnable { private static Supplier driveConnectionSupplier = Suppliers.memoize(() -> 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; @@ -147,53 +235,21 @@ public class ExportDomainListsAction implements Runnable { driveConnection = driveConnectionSupplier.get(); } - private void exportToDrive(String tld, String domains) { - verifyNotNull(driveConnection, "expecting non-null driveConnection"); - try { - Registry registry = Registry.get(tld); - if (registry.getDriveFolderId() == null) { - logger.atInfo().log( - "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.atInfo().log( - "Exporting registered domains succeeded for TLD %s, response was: %s", - tld, resultMsg); - } - } catch (Throwable e) { - logger.atSevere().withCause(e).log( - "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)) { - osWriter.write(domains); - } catch (IOException e) { - logger.atSevere().withCause(e).log( - "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.atInfo().log("Exporting %d domains for TLD %s to GCS and Drive.", domains.size(), tld); - exportToGcs(tld, domainsList); - exportToDrive(tld, domainsList); + if (exportToGcs(tld, domainsList, gcsBucket, gcsBufferSize)) { + getContext().incrementCounter("domain lists successful written out to GCS"); + } else { + getContext().incrementCounter("domain lists failed to write out to GCS"); + } + if (exportToDrive(tld, domainsList, driveConnection)) { + getContext().incrementCounter("domain lists successfully written out to Drive"); + } else { + getContext().incrementCounter("domain lists failed to write out to Drive"); + } } @VisibleForTesting diff --git a/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java b/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java index 6cc21d265..431291ee1 100644 --- a/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java +++ b/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java @@ -16,8 +16,7 @@ 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.export.ExportDomainListsAction.REGISTERED_DOMAINS_FILENAME; import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistDeletedDomain; @@ -34,25 +33,40 @@ 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.net.MediaType; import google.registry.export.ExportDomainListsAction.ExportDomainListsReducer; +import google.registry.model.ofy.Ofy; import google.registry.model.registry.Registry; import google.registry.model.registry.Registry.TldType; import google.registry.storage.drive.DriveConnection; +import google.registry.testing.DualDatabaseTest; +import google.registry.testing.FakeClock; import google.registry.testing.FakeResponse; +import google.registry.testing.InjectExtension; +import google.registry.testing.TestOfyAndSql; +import google.registry.testing.TestOfyOnly; import google.registry.testing.mapreduce.MapreduceTestCase; import java.io.FileNotFoundException; import org.joda.time.DateTime; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.ArgumentCaptor; /** Unit tests for {@link ExportDomainListsAction}. */ +@DualDatabaseTest class ExportDomainListsActionTest extends MapreduceTestCase { private GcsService gcsService; private DriveConnection driveConnection = mock(DriveConnection.class); private ArgumentCaptor bytesExportedToDrive = ArgumentCaptor.forClass(byte[].class); private final FakeResponse response = new FakeResponse(); + private final FakeClock clock = new FakeClock(DateTime.parse("2020-02-02T02:02:02Z")); + + @Order(Order.DEFAULT - 1) + @RegisterExtension + public final InjectExtension inject = + new InjectExtension().withStaticFieldOverride(Ofy.class, "clock", clock); @BeforeEach void beforeEach() { @@ -68,10 +82,12 @@ class ExportDomainListsActionTest extends MapreduceTestCase