Remove aggressive check in RegistryJpaIO.Write (#1889)

* Prevent saving duplicate rows in spec11 pipeline

* Chain applies together

core/src/main/java/google/registry/beam/common/RegistryJpaIO.java

@@ -14,7 +14,6 @@
 
 package google.registry.beam.common;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static org.apache.beam.sdk.values.TypeDescriptors.integers;
 
@@ -345,8 +344,7 @@ public final class RegistryJpaIO {
       try {
         tm().transact(
                 () -> {
-                  // Don't modify existing objects as it could lead to race conditions
-                  entities.forEach(this::verifyObjectNonexistence);
+                  // TODO(b/263502442): properly handle creations and blind-writes.
                   tm().putAll(entities);
                 });
         counter.inc(entities.size());
@@ -364,8 +362,7 @@ public final class RegistryJpaIO {
         try {
           tm().transact(
                   () -> {
-                    // Don't modify existing objects as it could lead to race conditions
-                    verifyObjectNonexistence(entity);
+                    // TODO(b/263502442): properly handle creations and blind-writes.
                     tm().put(entity);
                   });
           counter.inc();
@@ -391,15 +388,5 @@ public final class RegistryJpaIO {
         return "Non-SqlEntity: " + entity;
       }
     }
-
-    /** SqlBatchWriter should not re-write existing entities due to potential race conditions. */
-    private void verifyObjectNonexistence(Object obj) {
-      // We cannot rely on calling "insert" on the objects because the underlying JPA persist call
-      // adds the input object to the persistence context, meaning that any modifications (e.g.
-      // updateTimestamp) are reflected in the input object. Beam doesn't allow modification of
-      // input objects, so this throws an exception.
-      // TODO(go/non-datastore-allocateid): also check that all the objects have IDs
-      checkArgument(!tm().exists(obj), "Entities created in SqlBatchWriter must not already exist");
-    }
   }
 }

core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java

@@ -26,6 +26,7 @@ import google.registry.beam.common.RegistryJpaIO;
 import google.registry.beam.common.RegistryJpaIO.Read;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.RegistryConfig.ConfigModule;
+import google.registry.model.IdService;
 import google.registry.model.domain.Domain;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
@@ -45,6 +46,7 @@ import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Reshuffle;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.TypeDescriptor;
@@ -154,25 +156,36 @@ public class Spec11Pipeline implements Serializable {
 
   static void saveToSql(
       PCollection<KV<DomainNameInfo, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
-    String transformId = "Spec11 Threat Matches";
     LocalDate date = LocalDate.parse(options.getDate(), ISODateTimeFormat.date());
-    threatMatches.apply(
-        "Write to Sql: " + transformId,
-        RegistryJpaIO.<KV<DomainNameInfo, ThreatMatch>>write()
-            .withName(transformId)
-            .withBatchSize(options.getSqlWriteBatchSize())
-            .withJpaConverter(
-                (kv) -> {
-                  DomainNameInfo domainNameInfo = kv.getKey();
-                  return new Spec11ThreatMatch.Builder()
-                      .setThreatTypes(
-                          ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
-                      .setCheckDate(date)
-                      .setDomainName(domainNameInfo.domainName())
-                      .setDomainRepoId(domainNameInfo.domainRepoId())
-                      .setRegistrarId(domainNameInfo.registrarId())
-                      .build();
-                }));
+    String transformId = "Spec11 Threat Matches";
+    threatMatches
+        .apply(
+            "Construct objects",
+            ParDo.of(
+                new DoFn<KV<DomainNameInfo, ThreatMatch>, Spec11ThreatMatch>() {
+                  @ProcessElement
+                  public void processElement(
+                      @Element KV<DomainNameInfo, ThreatMatch> input,
+                      OutputReceiver<Spec11ThreatMatch> output) {
+                    Spec11ThreatMatch spec11ThreatMatch =
+                        new Spec11ThreatMatch.Builder()
+                            .setThreatTypes(
+                                ImmutableSet.of(ThreatType.valueOf(input.getValue().threatType())))
+                            .setCheckDate(date)
+                            .setDomainName(input.getKey().domainName())
+                            .setDomainRepoId(input.getKey().domainRepoId())
+                            .setRegistrarId(input.getKey().registrarId())
+                            .setId(IdService.allocateId())
+                            .build();
+                    output.output(spec11ThreatMatch);
+                  }
+                }))
+        .apply("Prevent Fusing", Reshuffle.viaRandomKey())
+        .apply(
+            "Write to Sql: " + transformId,
+            RegistryJpaIO.<Spec11ThreatMatch>write()
+                .withName(transformId)
+                .withBatchSize(options.getSqlWriteBatchSize()));
   }
 
   static void saveToGcs(

core/src/main/java/google/registry/model/EppResourceUtils.java

@@ -40,8 +40,8 @@ import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
+import java.util.Collection;
 import java.util.Comparator;
-import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -183,7 +183,7 @@ public final class EppResourceUtils {
    * @param now the logical time of the check
    */
   public static <T extends EppResource> ImmutableSet<String> checkResourcesExist(
-      Class<T> clazz, List<String> uniqueIds, final DateTime now) {
+      Class<T> clazz, Collection<String> uniqueIds, final DateTime now) {
     return ForeignKeyUtils.load(clazz, uniqueIds, now).keySet();
   }
 

core/src/main/java/google/registry/model/reporting/Spec11ThreatMatch.java

@@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableSet;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
 import google.registry.util.DomainNameUtils;
+import java.io.Serializable;
 import java.util.Set;
 import javax.persistence.Column;
 import javax.persistence.Entity;
@@ -39,7 +40,7 @@ import org.joda.time.LocalDate;
       @Index(name = "spec11threatmatch_tld_idx", columnList = "tld"),
       @Index(name = "spec11threatmatch_check_date_idx", columnList = "checkDate")
     })
-public class Spec11ThreatMatch extends ImmutableObject implements Buildable {
+public class Spec11ThreatMatch extends ImmutableObject implements Buildable, Serializable {
 
   /** The type of threat detected. */
   public enum ThreatType {

core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java

@@ -35,11 +35,12 @@ import google.registry.model.domain.Domain;
 import google.registry.model.domain.secdns.DomainDsData;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.Host;
+import google.registry.tools.params.NameserversParameter;
 import google.registry.tools.soy.DomainRenewSoyInfo;
 import google.registry.tools.soy.UniformRapidSuspensionSoyInfo;
 import google.registry.util.Clock;
-import google.registry.util.DomainNameUtils;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -69,9 +70,12 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
 
   @Parameter(
       names = {"-h", "--hosts"},
-      description = "Comma-delimited set of fully qualified host names to replace the current hosts"
-          + " on the domain.")
-  private List<String> newHosts = new ArrayList<>();
+      description =
+          "Comma-delimited set of fully qualified host names to replace the current hosts"
+              + " on the domain.",
+      converter = NameserversParameter.class,
+      validateWith = NameserversParameter.class)
+  private Set<String> newHosts = new HashSet<>();
 
   @Parameter(
       names = {"-s", "--dsdata"},
@@ -126,14 +130,10 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
   protected void initMutatingEppToolCommand() {
     superuser = true;
     DateTime now = clock.nowUtc();
-    ImmutableList<String> newCanonicalHosts =
-        newHosts.stream().map(DomainNameUtils::canonicalizeHostname).collect(toImmutableList());
-    ImmutableSet<String> newHostsSet = ImmutableSet.copyOf(newCanonicalHosts);
     Optional<Domain> domainOpt = loadByForeignKey(Domain.class, domainName, now);
     checkArgumentPresent(domainOpt, "Domain '%s' does not exist or is deleted", domainName);
     Domain domain = domainOpt.get();
-    Set<String> missingHosts =
-        difference(newHostsSet, checkResourcesExist(Host.class, newCanonicalHosts, now));
+    Set<String> missingHosts = difference(newHosts, checkResourcesExist(Host.class, newHosts, now));
     checkArgument(missingHosts.isEmpty(), "Hosts do not exist: %s", missingHosts);
     checkArgument(
         locksToPreserve.isEmpty() || undo,
@@ -187,9 +187,9 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
             "domainName",
             domainName,
             "hostsToAdd",
-            difference(newHostsSet, existingNameservers),
+            difference(newHosts, existingNameservers),
             "hostsToRemove",
-            difference(existingNameservers, newHostsSet),
+            difference(existingNameservers, newHosts),
             "statusesToApply",
             statusesToApply,
             "statusesToRemove",

core/src/main/java/google/registry/tools/params/NameserversParameter.java

@@ -25,6 +25,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.re2j.Matcher;
 import com.google.re2j.Pattern;
+import google.registry.util.DomainNameUtils;
 import java.util.Set;
 import java.util.stream.Stream;
 
@@ -50,12 +51,9 @@ public final class NameserversParameter extends ParameterConverterValidator<Set<
     if (Strings.isNullOrEmpty(value)) {
       return ImmutableSet.of();
     }
-    return Splitter.on(',')
-        .trimResults()
-        .omitEmptyStrings()
-        .splitToList(value)
-        .stream()
+    return Splitter.on(',').trimResults().omitEmptyStrings().splitToList(value).stream()
         .flatMap(NameserversParameter::splitNameservers)
+        .map(DomainNameUtils::canonicalizeHostname)
         .collect(toImmutableSet());
   }
 

core/src/test/java/google/registry/beam/common/RegistryJpaWriteTest.java

@@ -33,6 +33,7 @@ import java.io.Serializable;
 import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
 import org.apache.beam.sdk.transforms.Create;
 import org.joda.time.DateTime;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -67,6 +68,7 @@ class RegistryJpaWriteTest implements Serializable {
         .containsExactlyElementsIn(contacts);
   }
 
+  @Disabled("b/263502442")
   @Test
   void testFailure_writeExistingEntity() {
     // RegistryJpaIO.Write actions should not write existing objects to the database because the

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -280,9 +280,9 @@ class Spec11PipelineTest {
   private void verifySaveToCloudSql() {
     tm().transact(
             () -> {
-              ImmutableList<Spec11ThreatMatch> sqlThreatMatches =
+              ImmutableList<Spec11ThreatMatch> spec11ThreatMatches =
                   Spec11ThreatMatchDao.loadEntriesByDate(tm(), new LocalDate(2020, 1, 27));
-              assertThat(sqlThreatMatches)
+              assertThat(spec11ThreatMatches)
                   .comparingElementsUsing(immutableObjectCorrespondence("id"))
                   .containsExactlyElementsIn(sqlThreatMatches);
             });

core/src/test/java/google/registry/testing/DatabaseHelper.java

@@ -1006,7 +1006,6 @@ public final class DatabaseHelper {
    *
    * <p>This was coded for testing RDE since its queries depend on the associated entries.
    *
-   *
    * @see #persistResource(ImmutableObject)
    */
   public static <R extends EppResource> R persistEppResource(final R resource) {

core/src/test/java/google/registry/tools/CreateDomainCommandTest.java

@@ -57,6 +57,23 @@ class CreateDomainCommandTest extends EppToolCommandTestCase<CreateDomainCommand
     eppVerifier.verifySent("domain_create_complete.xml");
   }
 
+  @Test
+  void testSuccess_completeWithCanonicalization() throws Exception {
+    runCommandForced(
+        "--client=NewRegistrar",
+        "--period=1",
+        "--nameservers=NS1.zdns.google,ns2.ZDNS.google,ns3.zdns.gOOglE,ns4.zdns.google",
+        "--registrant=crr-admin",
+        "--admins=crr-admin",
+        "--techs=crr-tech",
+        "--password=2fooBAR",
+        "--ds_records=1 2 2 9F86D081884C7D659A2FEAA0C55AD015A3BF4F1B2B0B822CD15D6C15B0F00A08,4 5 1"
+            + " A94A8FE5CCB19BA61C4C0873D391E987982FBBD3",
+        "--ds_records=60485 5  2  D4B7D520E7BB5F0F67674A0CCEB1E3E0614B93C4F9E99B8383F6A1E4469DA50A",
+        "example.tld");
+    eppVerifier.verifySent("domain_create_complete.xml");
+  }
+
   @Test
   void testSuccess_completeWithSquareBrackets() throws Exception {
     runCommandForced(
@@ -74,6 +91,23 @@ class CreateDomainCommandTest extends EppToolCommandTestCase<CreateDomainCommand
     eppVerifier.verifySent("domain_create_complete.xml");
   }
 
+  @Test
+  void testSuccess_completeWithSquareBracketsAndCanonicalization() throws Exception {
+    runCommandForced(
+        "--client=NewRegistrar",
+        "--period=1",
+        "--nameservers=NS[1-4].zdns.google",
+        "--registrant=crr-admin",
+        "--admins=crr-admin",
+        "--techs=crr-tech",
+        "--password=2fooBAR",
+        "--ds_records=1 2 2 9F86D081884C7D659A2FEAA0C55AD015A3BF4F1B2B0B822CD15D6C15B0F00A08,4 5 1"
+            + " A94A8FE5CCB19BA61C4C0873D391E987982FBBD3",
+        "--ds_records=60485 5  2  D4B7D520E7BB5F0F67674A0CCEB1E3E0614B93C4F9E99B8383F6A1E4469DA50A",
+        "example.tld");
+    eppVerifier.verifySent("domain_create_complete.xml");
+  }
+
   @Test
   void testSuccess_minimal() throws Exception {
     // Test that each optional field can be omitted. Also tests the auto-gen password.

core/src/test/java/google/registry/tools/UniformRapidSuspensionCommandTest.java

@@ -163,6 +163,27 @@ class UniformRapidSuspensionCommandTest
     assertInStdout("--restore_client_hold");
   }
 
+  @Test
+  void testCommand_bracketNameserverNotationWithCanonicalization() throws Exception {
+    persistDomainWithHosts(defaultDomain, defaultDsData, ns1, ns2);
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--hosts=URS[1-2].example.com",
+        "--dsdata=1 1 1 A94A8FE5CCB19BA61C4C0873D391E987982FBBD3",
+        "--renew_one_year=false");
+    eppVerifier
+        .expectRegistrarId("CharlestonRoad")
+        .expectSuperuser()
+        .verifySent("uniform_rapid_suspension.xml")
+        .verifyNoMoreSent();
+    assertInStdout("uniform_rapid_suspension --undo");
+    assertInStdout("--domain_name evil.tld");
+    assertInStdout("--hosts ns1.example.com,ns2.example.com");
+    assertInStdout("--dsdata 1 2 3 DEAD,4 5 6 BEEF");
+    assertNotInStdout("--locks_to_preserve");
+    assertNotInStdout("--restore_client_hold");
+  }
+
   @Test
   void testUndo_removesLocksReplacesHostsAndDsData() throws Exception {
     persistDomainWithHosts(defaultDomain, defaultDsData, urs1, urs2);

core/src/test/java/google/registry/tools/UpdateDomainCommandTest.java

@@ -93,10 +93,10 @@ class UpdateDomainCommandTest extends EppToolCommandTestCase<UpdateDomainCommand
   }
 
   @Test
-  void testSuccess_completeWithSquareBrackets() throws Exception {
+  void testSuccess_completeWithSquareBracketsAndCanonicalization() throws Exception {
     runCommandForced(
         "--client=NewRegistrar",
-        "--add_nameservers=ns[1-2].zdns.google",
+        "--add_nameservers=NS[1-2].zdns.google",
         "--add_admins=crr-admin2",
         "--add_techs=crr-tech2",
         "--add_statuses=serverDeleteProhibited",