Remove aggressive check in RegistryJpaIO.Write (#1889)

* Prevent saving duplicate rows in spec11 pipeline

* Chain applies together

core/src/main/java/google/registry/batch/CheckPackagesComplianceAction.java

@@ -15,8 +15,10 @@ package google.registry.batch;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.primitives.Ints;
 import google.registry.config.RegistryConfig.Config;
@@ -28,8 +30,10 @@ import google.registry.request.Action;
 import google.registry.request.Action.Service;
 import google.registry.request.auth.Auth;
 import google.registry.ui.server.SendEmailUtils;
+import google.registry.util.Clock;
 import java.util.Optional;
 import javax.inject.Inject;
+import org.joda.time.Days;
 
 /**
  * An action that checks all {@link PackagePromotion} objects for compliance with their max create
@@ -44,77 +48,164 @@ public class CheckPackagesComplianceAction implements Runnable {
   public static final String PATH = "/_dr/task/checkPackagesCompliance";
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private final SendEmailUtils sendEmailUtils;
-  private final String packageCreateLimitEmailSubjectText;
-  private final String packageCreateLimitEmailBodyText;
+  private final Clock clock;
+  private final String packageCreateLimitEmailSubject;
+  private final String packageDomainLimitWarningEmailSubject;
+  private final String packageDomainLimitUpgradeEmailSubject;
+  private final String packageCreateLimitEmailBody;
+  private final String packageDomainLimitWarningEmailBody;
+  private final String packageDomainLimitUpgradeEmailBody;
   private final String registrySupportEmail;
+  private static final int THIRTY_DAYS = 30;
+  private static final int FORTY_DAYS = 40;
 
   @Inject
   public CheckPackagesComplianceAction(
       SendEmailUtils sendEmailUtils,
-      @Config("packageCreateLimitEmailSubjectText") String packageCreateLimitEmailSubjectText,
-      @Config("packageCreateLimitEmailBodyText") String packageCreateLimitEmailBodyText,
+      Clock clock,
+      @Config("packageCreateLimitEmailSubject") String packageCreateLimitEmailSubject,
+      @Config("packageDomainLimitWarningEmailSubject") String packageDomainLimitWarningEmailSubject,
+      @Config("packageDomainLimitUpgradeEmailSubject") String packageDomainLimitUpgradeEmailSubject,
+      @Config("packageCreateLimitEmailBody") String packageCreateLimitEmailBody,
+      @Config("packageDomainLimitWarningEmailBody") String packageDomainLimitWarningEmailBody,
+      @Config("packageDomainLimitUpgradeEmailBody") String packageDomainLimitUpgradeEmailBody,
       @Config("registrySupportEmail") String registrySupportEmail) {
     this.sendEmailUtils = sendEmailUtils;
-    this.packageCreateLimitEmailSubjectText = packageCreateLimitEmailSubjectText;
-    this.packageCreateLimitEmailBodyText = packageCreateLimitEmailBodyText;
+    this.clock = clock;
+    this.packageCreateLimitEmailSubject = packageCreateLimitEmailSubject;
+    this.packageDomainLimitWarningEmailSubject = packageDomainLimitWarningEmailSubject;
+    this.packageDomainLimitUpgradeEmailSubject = packageDomainLimitUpgradeEmailSubject;
+    this.packageCreateLimitEmailBody = packageCreateLimitEmailBody;
+    this.packageDomainLimitWarningEmailBody = packageDomainLimitWarningEmailBody;
+    this.packageDomainLimitUpgradeEmailBody = packageDomainLimitUpgradeEmailBody;
     this.registrySupportEmail = registrySupportEmail;
   }
 
   @Override
   public void run() {
-    tm().transact(
-            () -> {
-              ImmutableList<PackagePromotion> packages = tm().loadAllOf(PackagePromotion.class);
-              ImmutableList.Builder<PackagePromotion> packagesOverCreateLimit =
-                  new ImmutableList.Builder<>();
-              for (PackagePromotion packagePromo : packages) {
-                Long creates =
-                    (Long)
-                        tm().query(
-                                "SELECT COUNT(*) FROM DomainHistory WHERE current_package_token ="
-                                    + " :token AND modificationTime >= :lastBilling AND type ="
-                                    + " 'DOMAIN_CREATE'")
-                            .setParameter("token", packagePromo.getToken().getKey().toString())
-                            .setParameter(
-                                "lastBilling", packagePromo.getNextBillingDate().minusYears(1))
-                            .getSingleResult();
-                if (creates > packagePromo.getMaxCreates()) {
-                  int overage = Ints.saturatedCast(creates) - packagePromo.getMaxCreates();
-                  logger.atInfo().log(
-                      "Package with package token %s has exceeded their max domain creation limit"
-                          + " by %d name(s).",
-                      packagePromo.getToken().getKey(), overage);
-                  packagesOverCreateLimit.add(packagePromo);
-                }
-              }
-              if (packagesOverCreateLimit.build().isEmpty()) {
-                logger.atInfo().log("Found no packages over their create limit.");
-              } else {
-                logger.atInfo().log(
-                    "Found %d packages over their create limit.",
-                    packagesOverCreateLimit.build().size());
-                for (PackagePromotion packagePromotion : packagesOverCreateLimit.build()) {
-                  AllocationToken packageToken = tm().loadByKey(packagePromotion.getToken());
-                  Optional<Registrar> registrar =
-                      Registrar.loadByRegistrarIdCached(
-                          packageToken.getAllowedRegistrarIds().iterator().next());
-                  if (registrar.isPresent()) {
-                    String body =
-                        String.format(
-                            packageCreateLimitEmailBodyText,
-                            registrar.get().getRegistrarName(),
-                            packageToken.getToken(),
-                            registrySupportEmail);
-                    sendNotification(
-                        packageToken, packageCreateLimitEmailSubjectText, body, registrar.get());
-                  } else {
-                    logger.atSevere().log(
-                        String.format(
-                            "Could not find registrar for package token %s", packageToken));
-                  }
-                }
-              }
-            });
+    tm().transact(this::checkPackages);
+  }
+
+  private void checkPackages() {
+    ImmutableList<PackagePromotion> packages = tm().loadAllOf(PackagePromotion.class);
+    ImmutableList.Builder<PackagePromotion> packagesOverCreateLimitBuilder =
+        new ImmutableList.Builder<>();
+    ImmutableList.Builder<PackagePromotion> packagesOverActiveDomainsLimitBuilder =
+        new ImmutableList.Builder<>();
+    for (PackagePromotion packagePromo : packages) {
+      Long creates =
+          (Long)
+              tm().query(
+                      "SELECT COUNT(*) FROM DomainHistory WHERE current_package_token ="
+                          + " :token AND modificationTime >= :lastBilling AND type ="
+                          + " 'DOMAIN_CREATE'")
+                  .setParameter("token", packagePromo.getToken().getKey().toString())
+                  .setParameter("lastBilling", packagePromo.getNextBillingDate().minusYears(1))
+                  .getSingleResult();
+      if (creates > packagePromo.getMaxCreates()) {
+        int overage = Ints.saturatedCast(creates) - packagePromo.getMaxCreates();
+        logger.atInfo().log(
+            "Package with package token %s has exceeded their max domain creation limit"
+                + " by %d name(s).",
+            packagePromo.getToken().getKey(), overage);
+        packagesOverCreateLimitBuilder.add(packagePromo);
+      }
+
+      Long activeDomains =
+          tm().query(
+                  "SELECT COUNT(*) FROM Domain WHERE currentPackageToken = :token"
+                      + " AND deletionTime = :endOfTime",
+                  Long.class)
+              .setParameter("token", packagePromo.getToken())
+              .setParameter("endOfTime", END_OF_TIME)
+              .getSingleResult();
+      if (activeDomains > packagePromo.getMaxDomains()) {
+        int overage = Ints.saturatedCast(activeDomains) - packagePromo.getMaxDomains();
+        logger.atInfo().log(
+            "Package with package token %s has exceed their max active domains limit by"
+                + " %d name(s).",
+            packagePromo.getToken().getKey(), overage);
+        packagesOverActiveDomainsLimitBuilder.add(packagePromo);
+      }
+    }
+    handlePackageCreationOverage(packagesOverCreateLimitBuilder.build());
+    handleActiveDomainOverage(packagesOverActiveDomainsLimitBuilder.build());
+  }
+
+  private void handlePackageCreationOverage(ImmutableList<PackagePromotion> overageList) {
+    if (overageList.isEmpty()) {
+      logger.atInfo().log("Found no packages over their create limit.");
+      return;
+    }
+    logger.atInfo().log("Found %d packages over their create limit.", overageList.size());
+    for (PackagePromotion packagePromotion : overageList) {
+      AllocationToken packageToken = tm().loadByKey(packagePromotion.getToken());
+      Optional<Registrar> registrar =
+          Registrar.loadByRegistrarIdCached(
+              Iterables.getOnlyElement(packageToken.getAllowedRegistrarIds()));
+      if (registrar.isPresent()) {
+        String body =
+            String.format(
+                packageCreateLimitEmailBody,
+                registrar.get().getRegistrarName(),
+                packageToken.getToken(),
+                registrySupportEmail);
+        sendNotification(packageToken, packageCreateLimitEmailSubject, body, registrar.get());
+      } else {
+        throw new IllegalStateException(
+            String.format("Could not find registrar for package token %s", packageToken));
+      }
+    }
+  }
+
+  private void handleActiveDomainOverage(ImmutableList<PackagePromotion> overageList) {
+    if (overageList.isEmpty()) {
+      logger.atInfo().log("Found no packages over their active domains limit.");
+      return;
+    }
+    logger.atInfo().log("Found %d packages over their active domains limit.", overageList.size());
+    for (PackagePromotion packagePromotion : overageList) {
+      int daysSinceLastNotification =
+          packagePromotion
+              .getLastNotificationSent()
+              .map(sentDate -> Days.daysBetween(sentDate, clock.nowUtc()).getDays())
+              .orElse(Integer.MAX_VALUE);
+      if (daysSinceLastNotification < THIRTY_DAYS) {
+        // Don't send an email if notification was already sent within the last 30
+        // days
+        continue;
+      } else if (daysSinceLastNotification < FORTY_DAYS) {
+        // Send an upgrade email if last email was between 30 and 40 days ago
+        sendActiveDomainOverageEmail(/* warning= */ false, packagePromotion);
+      } else {
+        // Send a warning email
+        sendActiveDomainOverageEmail(/* warning= */ true, packagePromotion);
+      }
+    }
+  }
+
+  private void sendActiveDomainOverageEmail(boolean warning, PackagePromotion packagePromotion) {
+    String emailSubject =
+        warning ? packageDomainLimitWarningEmailSubject : packageDomainLimitUpgradeEmailSubject;
+    String emailTemplate =
+        warning ? packageDomainLimitWarningEmailBody : packageDomainLimitUpgradeEmailBody;
+    AllocationToken packageToken = tm().loadByKey(packagePromotion.getToken());
+    Optional<Registrar> registrar =
+        Registrar.loadByRegistrarIdCached(
+            Iterables.getOnlyElement(packageToken.getAllowedRegistrarIds()));
+    if (registrar.isPresent()) {
+      String body =
+          String.format(
+              emailTemplate,
+              registrar.get().getRegistrarName(),
+              packageToken.getToken(),
+              registrySupportEmail);
+      sendNotification(packageToken, emailSubject, body, registrar.get());
+      tm().put(packagePromotion.asBuilder().setLastNotificationSent(clock.nowUtc()).build());
+    } else {
+      throw new IllegalStateException(
+          String.format("Could not find registrar for package token %s", packageToken));
+    }
   }
 
   private void sendNotification(

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/config/RegistryConfig.java

@@ -1317,15 +1317,41 @@ public final class RegistryConfig {
     }
 
     @Provides
-    @Config("packageCreateLimitEmailSubjectText")
-    public static String providePackageCreateLimitEmailSubjectText(RegistryConfigSettings config) {
-      return config.packageMonitoring.packageCreateLimitEmailSubjectText;
+    @Config("packageCreateLimitEmailSubject")
+    public static String providePackageCreateLimitEmailSubject(RegistryConfigSettings config) {
+      return config.packageMonitoring.packageCreateLimitEmailSubject;
     }
 
     @Provides
-    @Config("packageCreateLimitEmailBodyText")
-    public static String providePackageCreateLimitEmailBodyText(RegistryConfigSettings config) {
-      return config.packageMonitoring.packageCreateLimitEmailBodyText;
+    @Config("packageCreateLimitEmailBody")
+    public static String providePackageCreateLimitEmailBody(RegistryConfigSettings config) {
+      return config.packageMonitoring.packageCreateLimitEmailBody;
+    }
+
+    @Provides
+    @Config("packageDomainLimitWarningEmailSubject")
+    public static String providePackageDomainLimitWarningEmailSubject(
+        RegistryConfigSettings config) {
+      return config.packageMonitoring.packageDomainLimitWarningEmailSubject;
+    }
+
+    @Provides
+    @Config("packageDomainLimitWarningEmailBody")
+    public static String providePackageDomainLimitWarningEmailBody(RegistryConfigSettings config) {
+      return config.packageMonitoring.packageDomainLimitWarningEmailBody;
+    }
+
+    @Provides
+    @Config("packageDomainLimitUpgradeEmailSubject")
+    public static String providePackageDomainLimitUpgradeEmailSubject(
+        RegistryConfigSettings config) {
+      return config.packageMonitoring.packageDomainLimitUpgradeEmailSubject;
+    }
+
+    @Provides
+    @Config("packageDomainLimitUpgradeEmailBody")
+    public static String providePackageDomainLimitUpgradeEmailBody(RegistryConfigSettings config) {
+      return config.packageMonitoring.packageDomainLimitUpgradeEmailBody;
     }
   }
 

core/src/main/java/google/registry/config/RegistryConfigSettings.java

@@ -255,7 +255,11 @@ public class RegistryConfigSettings {
 
   /** Configuration for package compliance monitoring. */
   public static class PackageMonitoring {
-    public String packageCreateLimitEmailSubjectText;
-    public String packageCreateLimitEmailBodyText;
+    public String packageCreateLimitEmailSubject;
+    public String packageCreateLimitEmailBody;
+    public String packageDomainLimitWarningEmailSubject;
+    public String packageDomainLimitWarningEmailBody;
+    public String packageDomainLimitUpgradeEmailSubject;
+    public String packageDomainLimitUpgradeEmailBody;
   }
 }

core/src/main/java/google/registry/config/files/default-config.yaml

@@ -536,9 +536,9 @@ sslCertificateValidation:
 # Configuration options for the package compliance monitoring
 packageMonitoring:
   # Email subject text to notify partners their package has exceeded the limit for domain creates
-  packageCreateLimitEmailSubjectText: "NOTICE: Your Package Is Being Upgraded"
+  packageCreateLimitEmailSubject: "NOTICE: Your package is being upgraded"
   # Email body text template notify partners their package has exceeded the limit for domain creates
-  packageCreateLimitEmailBodyText: >
+  packageCreateLimitEmailBody: >
     Dear %1$s,
     
     We are contacting you to inform you that your package with the package token 
@@ -550,3 +550,37 @@ packageMonitoring:
     
     Regards,
     Example Registry
+
+  # Email subject text to notify partners their package has exceeded the limit for current active domains and warn them their package will be upgraded in 30 days
+  packageDomainLimitWarningEmailSubject: "WARNING: Your package has exceeded the domain limit"
+  # Email body text template to warn partners their package has exceeded the limit for active domains and will be upgraded in 30 days
+  packageDomainLimitWarningEmailBody: >
+    Dear %1$s,
+      
+    We are contacting you to inform you that your package with the package token 
+    %2$s has exceeded its limit for active domains.
+    Your package will be upgraded to the next tier in 30 days if the number of active domains does not return below the limit. 
+
+    If you have any questions or require additional support, please contact us
+    at %3$s.
+
+    Regards,
+    Example Registry
+
+  # Email subject text to notify partners their package has exceeded the limit
+  # for current active domains for more than 30 days and will be upgraded
+  packageDomainLimitUpgradeEmailSubject: "NOTICE: Your package is being upgraded"
+  # Email body text template to warn partners their package has exceeded the
+  # limit for active domains for more than 30 days and will be upgraded
+  packageDomainLimitUpgradeEmailBody: >
+    Dear %1$s,
+    
+    We are contacting you to inform you that your package with the package token 
+    %2$s has exceeded its limit for active domains.
+    Your package will now be upgraded to the next tier. 
+
+    If you have any questions or require additional support, please contact us
+    at %3$s.
+
+    Regards,
+    Example Registry

core/src/main/java/google/registry/env/common/backend/WEB-INF/web.xml

@@ -175,16 +175,6 @@
     <url-pattern>/_dr/cron/fanout</url-pattern>
   </servlet-mapping>
 
-  <!-- Pipeline GUI servlets. -->
-  <servlet>
-    <servlet-name>pipeline</servlet-name>
-    <servlet-class>com.google.appengine.tools.pipeline.impl.servlets.PipelineServlet</servlet-class>
-  </servlet>
-  <servlet-mapping>
-    <servlet-name>pipeline</servlet-name>
-    <url-pattern>/_ah/pipeline/*</url-pattern>
-  </servlet-mapping>
-
   <!-- Syncs registrars to the registrar spreadsheet.  -->
   <servlet-mapping>
     <servlet-name>backend-servlet</servlet-name>

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/batch/CheckPackagesComplianceActionTest.java

@@ -57,7 +57,13 @@ public class CheckPackagesComplianceActionTest {
   // This is the default creation time for test data.
   private final FakeClock clock = new FakeClock(DateTime.parse("2012-03-25TZ"));
   private static final String CREATE_LIMIT_EMAIL_SUBJECT = "create limit subject";
+  private static final String DOMAIN_LIMIT_WARNING_EMAIL_SUBJECT = "domain limit warning subject";
+  private static final String DOMAIN_LIMIT_UPGRADE_EMAIL_SUBJECT = "domain limit upgrade subject";
   private static final String CREATE_LIMIT_EMAIL_BODY = "create limit body %1$s %2$s %3$s";
+  private static final String DOMAIN_LIMIT_WARNING_EMAIL_BODY =
+      "domain limit warning body %1$s %2$s %3$s";
+  private static final String DOMAIN_LIMIT_UPGRADE_EMAIL_BODY =
+      "domain limit upgrade body %1$s %2$s %3$s";
   private static final String SUPPORT_EMAIL = "registry@test.com";
 
   @RegisterExtension
@@ -70,7 +76,6 @@ public class CheckPackagesComplianceActionTest {
   private final Logger loggerToIntercept =
       Logger.getLogger(CheckPackagesComplianceAction.class.getCanonicalName());
   private final SendEmailService emailService = mock(SendEmailService.class);
-
   private Contact contact;
   private PackagePromotion packagePromotion;
   private SendEmailUtils sendEmailUtils;
@@ -88,7 +93,15 @@ public class CheckPackagesComplianceActionTest {
     createTld("tld");
     action =
         new CheckPackagesComplianceAction(
-            sendEmailUtils, CREATE_LIMIT_EMAIL_SUBJECT, CREATE_LIMIT_EMAIL_BODY, SUPPORT_EMAIL);
+            sendEmailUtils,
+            clock,
+            CREATE_LIMIT_EMAIL_SUBJECT,
+            DOMAIN_LIMIT_WARNING_EMAIL_SUBJECT,
+            DOMAIN_LIMIT_UPGRADE_EMAIL_SUBJECT,
+            CREATE_LIMIT_EMAIL_BODY,
+            DOMAIN_LIMIT_WARNING_EMAIL_BODY,
+            DOMAIN_LIMIT_UPGRADE_EMAIL_BODY,
+            SUPPORT_EMAIL);
     token =
         persistResource(
             new AllocationToken.Builder()
@@ -110,7 +123,6 @@ public class CheckPackagesComplianceActionTest {
             .setLastNotificationSent(DateTime.parse("2010-11-12T05:00:00Z"))
             .build();
 
-    tm().transact(() -> tm().put(packagePromotion));
     contact = persistActiveContact("contact1234");
   }
 
@@ -121,6 +133,7 @@ public class CheckPackagesComplianceActionTest {
 
   @Test
   void testSuccess_noPackageOverCreateLimit() {
+    tm().transact(() -> tm().put(packagePromotion));
     persistEppResource(
         DatabaseHelper.newDomain("foo.tld", contact)
             .asBuilder()
@@ -136,6 +149,7 @@ public class CheckPackagesComplianceActionTest {
 
   @Test
   void testSuccess_onePackageOverCreateLimit() throws Exception {
+    tm().transact(() -> tm().put(packagePromotion));
     // Create limit is 1, creating 2 domains to go over the limit
     persistEppResource(
         DatabaseHelper.newDomain("foo.tld", contact)
@@ -168,6 +182,7 @@ public class CheckPackagesComplianceActionTest {
 
   @Test
   void testSuccess_multiplePackagesOverCreateLimit() {
+    tm().transact(() -> tm().put(packagePromotion));
     // Create limit is 1, creating 2 domains to go over the limit
     persistEppResource(
         DatabaseHelper.newDomain("foo.tld", contact)
@@ -234,6 +249,7 @@ public class CheckPackagesComplianceActionTest {
 
   @Test
   void testSuccess_onlyChecksCurrentBillingYear() {
+    tm().transact(() -> tm().put(packagePromotion));
     AllocationToken token2 =
         persistResource(
             new AllocationToken.Builder()
@@ -273,4 +289,278 @@ public class CheckPackagesComplianceActionTest {
         .hasLogAtLevelWithMessage(Level.INFO, "Found no packages over their create limit.");
     verifyNoInteractions(emailService);
   }
+
+  @Test
+  void testSuccess_noPackageOverActiveDomainsLimit() {
+    tm().transact(() -> tm().put(packagePromotion));
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    action.run();
+    verifyNoInteractions(emailService);
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found no packages over their active domains limit.");
+  }
+
+  @Test
+  void testSuccess_onePackageOverActiveDomainsLimit() {
+    packagePromotion = packagePromotion.asBuilder().setMaxCreates(4).setMaxDomains(1).build();
+    tm().transact(() -> tm().put(packagePromotion));
+    // Domains limit is 1, creating 2 domains to go over the limit
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    AllocationToken token2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("token")
+                .setTokenType(TokenType.PACKAGE)
+                .setCreationTimeForTest(DateTime.parse("2010-11-12T05:00:00Z"))
+                .setAllowedTlds(ImmutableSet.of("foo"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setRenewalPriceBehavior(RenewalPriceBehavior.SPECIFIED)
+                .setDiscountFraction(1)
+                .build());
+    PackagePromotion packagePromotion2 =
+        new PackagePromotion.Builder()
+            .setToken(token2)
+            .setMaxDomains(8)
+            .setMaxCreates(4)
+            .setPackagePrice(Money.of(CurrencyUnit.USD, 1000))
+            .setNextBillingDate(DateTime.parse("2012-11-12T05:00:00Z"))
+            .build();
+    tm().transact(() -> tm().put(packagePromotion2));
+    persistEppResource(
+        DatabaseHelper.newDomain("foo2.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token2.createVKey())
+            .build());
+
+    action.run();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found 1 packages over their active domains limit.");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token abc123 has exceed their max active domains limit by 1"
+                + " name(s).");
+    verify(emailService).sendEmail(emailCaptor.capture());
+    EmailMessage emailMessage = emailCaptor.getValue();
+    assertThat(emailMessage.subject()).isEqualTo(DOMAIN_LIMIT_WARNING_EMAIL_SUBJECT);
+    assertThat(emailMessage.body())
+        .isEqualTo(
+            String.format(
+                DOMAIN_LIMIT_WARNING_EMAIL_BODY, "The Registrar", "abc123", SUPPORT_EMAIL));
+    PackagePromotion packageAfterCheck =
+        tm().transact(() -> PackagePromotion.loadByTokenString(token.getToken()).get());
+    assertThat(packageAfterCheck.getLastNotificationSent().get()).isEqualTo(clock.nowUtc());
+  }
+
+  @Test
+  void testSuccess_multiplePackagesOverActiveDomainsLimit() {
+    tm().transact(
+            () -> tm().put(packagePromotion.asBuilder().setMaxDomains(1).setMaxCreates(4).build()));
+    // Domains limit is 1, creating 2 domains to go over the limit
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    AllocationToken token2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("token")
+                .setTokenType(TokenType.PACKAGE)
+                .setCreationTimeForTest(DateTime.parse("2010-11-12T05:00:00Z"))
+                .setAllowedTlds(ImmutableSet.of("foo"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setRenewalPriceBehavior(RenewalPriceBehavior.SPECIFIED)
+                .setDiscountFraction(1)
+                .build());
+    PackagePromotion packagePromotion2 =
+        new PackagePromotion.Builder()
+            .setToken(token2)
+            .setMaxDomains(1)
+            .setMaxCreates(5)
+            .setPackagePrice(Money.of(CurrencyUnit.USD, 1000))
+            .setNextBillingDate(DateTime.parse("2012-11-12T05:00:00Z"))
+            .build();
+    tm().transact(() -> tm().put(packagePromotion2));
+
+    persistEppResource(
+        DatabaseHelper.newDomain("foo2.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token2.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz2.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token2.createVKey())
+            .build());
+
+    action.run();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found 2 packages over their active domains limit.");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token abc123 has exceed their max active domains limit by 1"
+                + " name(s).");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token token has exceed their max active domains limit by 1"
+                + " name(s).");
+    verify(emailService, times(2)).sendEmail(any(EmailMessage.class));
+  }
+
+  @Test
+  void testSuccess_packageOverActiveDomainsLimitAlreadySentWarningEmail_DoesNotSendAgain() {
+    packagePromotion =
+        packagePromotion
+            .asBuilder()
+            .setMaxCreates(4)
+            .setMaxDomains(1)
+            .setLastNotificationSent(clock.nowUtc().minusDays(5))
+            .build();
+    tm().transact(() -> tm().put(packagePromotion));
+    // Domains limit is 1, creating 2 domains to go over the limit
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    action.run();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found 1 packages over their active domains limit.");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token abc123 has exceed their max active domains limit by 1"
+                + " name(s).");
+    verifyNoInteractions(emailService);
+    PackagePromotion packageAfterCheck =
+        tm().transact(() -> PackagePromotion.loadByTokenString(token.getToken()).get());
+    assertThat(packageAfterCheck.getLastNotificationSent().get())
+        .isEqualTo(clock.nowUtc().minusDays(5));
+  }
+
+  @Test
+  void testSuccess_packageOverActiveDomainsLimitAlreadySentWarningEmailOver40DaysAgo_SendsAgain() {
+    packagePromotion =
+        packagePromotion
+            .asBuilder()
+            .setMaxCreates(4)
+            .setMaxDomains(1)
+            .setLastNotificationSent(clock.nowUtc().minusDays(45))
+            .build();
+    tm().transact(() -> tm().put(packagePromotion));
+    // Domains limit is 1, creating 2 domains to go over the limit
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    action.run();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found 1 packages over their active domains limit.");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token abc123 has exceed their max active domains limit by 1"
+                + " name(s).");
+    verify(emailService).sendEmail(emailCaptor.capture());
+    EmailMessage emailMessage = emailCaptor.getValue();
+    assertThat(emailMessage.subject()).isEqualTo(DOMAIN_LIMIT_WARNING_EMAIL_SUBJECT);
+    assertThat(emailMessage.body())
+        .isEqualTo(
+            String.format(
+                DOMAIN_LIMIT_WARNING_EMAIL_BODY, "The Registrar", "abc123", SUPPORT_EMAIL));
+    PackagePromotion packageAfterCheck =
+        tm().transact(() -> PackagePromotion.loadByTokenString(token.getToken()).get());
+    assertThat(packageAfterCheck.getLastNotificationSent().get()).isEqualTo(clock.nowUtc());
+  }
+
+  @Test
+  void testSuccess_packageOverActiveDomainsLimitAlreadySentWarning30DaysAgo_SendsUpgradeEmail() {
+    packagePromotion =
+        packagePromotion
+            .asBuilder()
+            .setMaxCreates(4)
+            .setMaxDomains(1)
+            .setLastNotificationSent(clock.nowUtc().minusDays(31))
+            .build();
+    tm().transact(() -> tm().put(packagePromotion));
+    // Domains limit is 1, creating 2 domains to go over the limit
+    persistEppResource(
+        DatabaseHelper.newDomain("foo.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+    persistEppResource(
+        DatabaseHelper.newDomain("buzz.tld", contact)
+            .asBuilder()
+            .setCurrentPackageToken(token.createVKey())
+            .build());
+
+    action.run();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(Level.INFO, "Found 1 packages over their active domains limit.");
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Package with package token abc123 has exceed their max active domains limit by 1"
+                + " name(s).");
+    verify(emailService).sendEmail(emailCaptor.capture());
+    EmailMessage emailMessage = emailCaptor.getValue();
+    assertThat(emailMessage.subject()).isEqualTo(DOMAIN_LIMIT_UPGRADE_EMAIL_SUBJECT);
+    assertThat(emailMessage.body())
+        .isEqualTo(
+            String.format(
+                DOMAIN_LIMIT_UPGRADE_EMAIL_BODY, "The Registrar", "abc123", SUPPORT_EMAIL));
+    PackagePromotion packageAfterCheck =
+        tm().transact(() -> PackagePromotion.loadByTokenString(token.getToken()).get());
+    assertThat(packageAfterCheck.getLastNotificationSent().get()).isEqualTo(clock.nowUtc());
+  }
 }

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",