Add a Beam pipeline to generate RDE deposit (part 1) (#1219)

This is the first part of the RdeStagingAction SQL migration where the
mapper logic is implemented in Beam.

A few helper methods are added to convert the DomainContent, HostBase
and ContactBase to their respective terminal child classes. This is
necessary and possible because the child classes do not have extra
fields and the base classes exist only to be embedded to other entities
(such as the various HistoryEntry entities). The conversion is necessary
because most of our code expects the terminal classes, such as the
RdeMarshaller's various marshallXXX() methods. The alternative would be
to change all the call sites, which seems to be much more disruptive.

Unfortunately there is is no good way to do this conversion than just
creating a builder and setting every fields there is.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1219)
<!-- Reviewable:end -->

core/src/main/java/google/registry/beam/rde/RdePipeline.java

@@ -0,0 +1,260 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.beam.rde;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.model.EppResourceUtils.loadAtPointInTimeAsync;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+import com.google.common.io.BaseEncoding;
+import google.registry.beam.common.RegistryJpaIO;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.host.HostResource;
+import google.registry.model.rde.RdeMode;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.Registrar.Type;
+import google.registry.persistence.VKey;
+import google.registry.rde.DepositFragment;
+import google.registry.rde.PendingDeposit;
+import google.registry.rde.PendingDeposit.PendingDepositCoder;
+import google.registry.rde.RdeFragmenter;
+import google.registry.rde.RdeMarshaller;
+import google.registry.xml.ValidationMode;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Supplier;
+import javax.persistence.Entity;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.FlatMapElements;
+import org.apache.beam.sdk.transforms.Flatten;
+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.PCollectionList;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
+import org.joda.time.DateTime;
+
+/**
+ * Definition of a Dataflow Flex template, which generates RDE/BRDA deposits.
+ *
+ * <p>To stage this template locally, run the {@code stage_beam_pipeline.sh} shell script.
+ *
+ * <p>Then, you can run the staged template via the API client library, gCloud or a raw REST call.
+ *
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
+ */
+public class RdePipeline implements Serializable {
+
+  private final RdeMarshaller marshaller;
+  private final ImmutableSetMultimap<String, PendingDeposit> pendings;
+
+  // Registrars to be excluded from data escrow. Not including the sandbox-only OTE type so that
+  // if sneaks into production we would get an extra signal.
+  private static final ImmutableSet<Type> IGNORED_REGISTRAR_TYPES =
+      Sets.immutableEnumSet(Registrar.Type.MONITORING, Registrar.Type.TEST);
+
+  private static final String EPP_RESOURCE_QUERY =
+      "SELECT id FROM %entity% "
+          + "WHERE COALESCE(creationClientId, '') NOT LIKE 'prober-%' "
+          + "AND COALESCE(currentSponsorClientId, '') NOT LIKE 'prober-%' "
+          + "AND COALESCE(lastEppUpdateClientId, '') NOT LIKE 'prober-%'";
+
+  public static String createEppResourceQuery(Class<? extends EppResource> clazz) {
+    return EPP_RESOURCE_QUERY.replace("%entity%", clazz.getAnnotation(Entity.class).name())
+        + (clazz.equals(DomainBase.class) ? " AND tld in (:tlds)" : "");
+  }
+
+  RdePipeline(RdePipelineOptions options) throws IOException, ClassNotFoundException {
+    this.marshaller = new RdeMarshaller(ValidationMode.valueOf(options.getValidationMode()));
+    this.pendings = decodePendings(options.getPendings());
+  }
+
+  @VisibleForTesting
+  PipelineResult run(Pipeline pipeline) {
+    createFragments(pipeline);
+    return pipeline.run();
+  }
+
+  PipelineResult run() {
+    return run(Pipeline.create());
+  }
+
+  PCollection<KV<PendingDeposit, DepositFragment>> createFragments(Pipeline pipeline) {
+    PCollection<KV<PendingDeposit, DepositFragment>> fragments =
+        PCollectionList.of(processRegistrars(pipeline))
+            .and(processNonRegistrarEntities(pipeline, DomainBase.class))
+            .and(processNonRegistrarEntities(pipeline, ContactResource.class))
+            .and(processNonRegistrarEntities(pipeline, HostResource.class))
+            .apply(Flatten.pCollections())
+            .setCoder(
+                KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)));
+    return fragments;
+  }
+
+  PCollection<KV<PendingDeposit, DepositFragment>> processRegistrars(Pipeline pipeline) {
+    return pipeline
+        .apply(
+            "Read all production Registrar entities",
+            RegistryJpaIO.read(
+                "SELECT clientIdentifier FROM Registrar WHERE type NOT IN (:types)",
+                ImmutableMap.of("types", IGNORED_REGISTRAR_TYPES),
+                String.class,
+                // TODO: consider adding coders for entities and pass them directly instead of using
+                // VKeys.
+                id -> VKey.createSql(Registrar.class, id)))
+        .apply(
+            "Marshal Registrar into DepositFragment",
+            FlatMapElements.into(
+                    TypeDescriptors.kvs(
+                        TypeDescriptor.of(PendingDeposit.class),
+                        TypeDescriptor.of(DepositFragment.class)))
+                .via(
+                    (VKey<Registrar> key) -> {
+                      Registrar registrar = jpaTm().transact(() -> jpaTm().loadByKey(key));
+                      DepositFragment fragment = marshaller.marshalRegistrar(registrar);
+                      return pendings.values().stream()
+                          .map(pending -> KV.of(pending, fragment))
+                          .collect(toImmutableSet());
+                    }));
+  }
+
+  @SuppressWarnings("deprecation") // Reshuffle is still recommended by Dataflow.
+  <T extends EppResource>
+      PCollection<KV<PendingDeposit, DepositFragment>> processNonRegistrarEntities(
+          Pipeline pipeline, Class<T> clazz) {
+    return createInputs(pipeline, clazz)
+        .apply("Marshal " + clazz.getSimpleName() + " into DepositFragment", mapToFragments(clazz))
+        .setCoder(KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)))
+        .apply(
+            "Reshuffle KV<PendingDeposit, DepositFragment> of "
+                + clazz.getSimpleName()
+                + " to prevent fusion",
+            Reshuffle.of());
+  }
+
+  <T extends EppResource> PCollection<VKey<T>> createInputs(Pipeline pipeline, Class<T> clazz) {
+    return pipeline.apply(
+        "Read all production " + clazz.getSimpleName() + " entities",
+        RegistryJpaIO.read(
+            createEppResourceQuery(clazz),
+            clazz.equals(DomainBase.class)
+                ? ImmutableMap.of("tlds", pendings.keySet())
+                : ImmutableMap.of(),
+            String.class,
+            // TODO: consider adding coders for entities and pass them directly instead of using
+            // VKeys.
+            x -> VKey.create(clazz, x)));
+  }
+
+  <T extends EppResource>
+      FlatMapElements<VKey<T>, KV<PendingDeposit, DepositFragment>> mapToFragments(Class<T> clazz) {
+    return FlatMapElements.into(
+            TypeDescriptors.kvs(
+                TypeDescriptor.of(PendingDeposit.class), TypeDescriptor.of(DepositFragment.class)))
+        .via(
+            (VKey<T> key) -> {
+              T resource = jpaTm().transact(() -> jpaTm().loadByKey(key));
+              // The set of all TLDs to which this resource should be emitted.
+              ImmutableSet<String> tlds =
+                  clazz.equals(DomainBase.class)
+                      ? ImmutableSet.of(((DomainBase) resource).getTld())
+                      : pendings.keySet();
+              // Get the set of all point-in-time watermarks we need, to minimize rewinding.
+              ImmutableSet<DateTime> dates =
+                  tlds.stream()
+                      .map(pendings::get)
+                      .flatMap(ImmutableSet::stream)
+                      .map(PendingDeposit::watermark)
+                      .collect(toImmutableSet());
+              // Launch asynchronous fetches of point-in-time representations of resource.
+              ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes =
+                  ImmutableMap.copyOf(
+                      Maps.asMap(dates, input -> loadAtPointInTimeAsync(resource, input)));
+              // Convert resource to an XML fragment for each watermark/mode pair lazily and cache
+              // the result.
+              RdeFragmenter fragmenter = new RdeFragmenter(resourceAtTimes, marshaller);
+              List<KV<PendingDeposit, DepositFragment>> results = new ArrayList<>();
+              for (String tld : tlds) {
+                for (PendingDeposit pending : pendings.get(tld)) {
+                  // Hosts and contacts don't get included in BRDA deposits.
+                  if (pending.mode() == RdeMode.THIN && !clazz.equals(DomainBase.class)) {
+                    continue;
+                  }
+                  Optional<DepositFragment> fragment =
+                      fragmenter.marshal(pending.watermark(), pending.mode());
+                  fragment.ifPresent(
+                      depositFragment -> results.add(KV.of(pending, depositFragment)));
+                }
+              }
+              return results;
+            });
+  }
+
+  /**
+   * Decodes the pipeline option extracted from the URL parameter sent by the pipeline launcher to
+   * the original TLD to pending deposit map.
+   */
+  @SuppressWarnings("unchecked")
+  static ImmutableSetMultimap<String, PendingDeposit> decodePendings(String encodedPending)
+      throws IOException, ClassNotFoundException {
+    try (ObjectInputStream ois =
+        new ObjectInputStream(
+            new ByteArrayInputStream(
+                BaseEncoding.base64Url().omitPadding().decode(encodedPending)))) {
+      return (ImmutableSetMultimap<String, PendingDeposit>) ois.readObject();
+    }
+  }
+
+  /**
+   * Encodes the TLD to pending deposit map in an URL safe string that is sent to the pipeline
+   * worker by the pipeline launcher as a pipeline option.
+   */
+  public static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
+      throws IOException {
+    try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
+      ObjectOutputStream oos = new ObjectOutputStream(baos);
+      oos.writeObject(pendings);
+      oos.flush();
+      return BaseEncoding.base64Url().omitPadding().encode(baos.toByteArray());
+    }
+  }
+
+  public static void main(String[] args) throws IOException, ClassNotFoundException {
+    PipelineOptionsFactory.register(RdePipelineOptions.class);
+    RdePipelineOptions options = PipelineOptionsFactory.fromArgs(args).as(RdePipelineOptions.class);
+    new RdePipeline(options).run();
+  }
+}

core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java

@@ -0,0 +1,32 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.beam.rde;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the spec11 pipeline. */
+public interface RdePipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The Base64-encoded serialized map of TLDs to PendingDeposit.")
+  String getPendings();
+
+  void setPendings(String value);
+
+  @Description("The validation mode (LENIENT|STRICT) that the RDE marshaller uses.")
+  String getValidationMode();
+
+  void setValidationMode(String value);
+}

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

@@ -398,7 +398,7 @@ public final class EppResourceUtils {
                 .orElse(null);
     if (resourceAtPointInTime == null) {
       logger.atSevere().log(
-          "Couldn't load resource at % for key %s, falling back to resource %s.",
+          "Couldn't load resource at %s for key %s, falling back to resource %s.",
           timestamp, resource.createVKey(), resource);
       return resource;
     }

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -110,7 +110,8 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getContactBase();
+    return getContactBase()
+        .map(contactBase -> new ContactResource.Builder().copyFrom(contactBase).build());
   }
 
   @PostLoad

core/src/main/java/google/registry/model/contact/ContactResource.java

@@ -79,5 +79,26 @@ public class ContactResource extends ContactBase
     private Builder(ContactResource instance) {
       super(instance);
     }
+
+    public Builder copyFrom(ContactBase contactBase) {
+      return this.setAuthInfo(contactBase.getAuthInfo())
+          .setContactId(contactBase.getContactId())
+          .setCreationClientId(contactBase.getCreationClientId())
+          .setCreationTime(contactBase.getCreationTime())
+          .setDeletionTime(contactBase.getDeletionTime())
+          .setDisclose(contactBase.getDisclose())
+          .setEmailAddress(contactBase.getEmailAddress())
+          .setFaxNumber(contactBase.getFaxNumber())
+          .setInternationalizedPostalInfo(contactBase.getInternationalizedPostalInfo())
+          .setLastTransferTime(contactBase.getLastTransferTime())
+          .setLastEppUpdateClientId(contactBase.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(contactBase.getLastEppUpdateTime())
+          .setLocalizedPostalInfo(contactBase.getLocalizedPostalInfo())
+          .setPersistedCurrentSponsorClientId(contactBase.getPersistedCurrentSponsorClientId())
+          .setRepoId(contactBase.getRepoId())
+          .setStatusValues(contactBase.getStatusValues())
+          .setTransferData(contactBase.getTransferData())
+          .setVoiceNumber(contactBase.getVoiceNumber());
+    }
   }
 }

core/src/main/java/google/registry/model/domain/DomainBase.java

@@ -181,5 +181,34 @@ public class DomainBase extends DomainContent
     Builder(DomainBase instance) {
       super(instance);
     }
+
+    public Builder copyFrom(DomainContent domainContent) {
+      return this.setAuthInfo(domainContent.getAuthInfo())
+          .setAutorenewPollMessage(domainContent.getAutorenewPollMessage())
+          .setAutorenewBillingEvent(domainContent.getAutorenewBillingEvent())
+          .setAutorenewEndTime(domainContent.getAutorenewEndTime())
+          .setContacts(domainContent.getContacts())
+          .setCreationClientId(domainContent.getCreationClientId())
+          .setCreationTime(domainContent.getCreationTime())
+          .setDomainName(domainContent.getDomainName())
+          .setDeletePollMessage(domainContent.getDeletePollMessage())
+          .setDsData(domainContent.getDsData())
+          .setDeletionTime(domainContent.getDeletionTime())
+          .setGracePeriods(domainContent.getGracePeriods())
+          .setIdnTableName(domainContent.getIdnTableName())
+          .setLastTransferTime(domainContent.getLastTransferTime())
+          .setLaunchNotice(domainContent.getLaunchNotice())
+          .setLastEppUpdateClientId(domainContent.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(domainContent.getLastEppUpdateTime())
+          .setNameservers(domainContent.getNameservers())
+          .setPersistedCurrentSponsorClientId(domainContent.getPersistedCurrentSponsorClientId())
+          .setRegistrant(domainContent.getRegistrant())
+          .setRegistrationExpirationTime(domainContent.getRegistrationExpirationTime())
+          .setRepoId(domainContent.getRepoId())
+          .setSmdId(domainContent.getSmdId())
+          .setSubordinateHosts(domainContent.getSubordinateHosts())
+          .setStatusValues(domainContent.getStatusValues())
+          .setTransferData(domainContent.getTransferData());
+    }
   }
 }

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -252,11 +252,18 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getDomainContent();
+    return getDomainContent()
+        .map(domainContent -> new DomainBase.Builder().copyFrom(domainContent).build());
   }
 
   @PostLoad
   void postLoad() {
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(domainTransactionRecords);
+    Hibernate.initialize(nsHosts);
+    Hibernate.initialize(dsDataHistories);
+    Hibernate.initialize(gracePeriodHistories);
+
     if (domainContent != null) {
       domainContent.nsHosts = nullToEmptyImmutableCopy(nsHosts);
       domainContent.gracePeriods =
@@ -278,12 +285,6 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         }
       }
     }
-
-    // TODO(b/188044616): Determine why Eager loading doesn't work here.
-    Hibernate.initialize(domainTransactionRecords);
-    Hibernate.initialize(nsHosts);
-    Hibernate.initialize(dsDataHistories);
-    Hibernate.initialize(gracePeriodHistories);
   }
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type

core/src/main/java/google/registry/model/host/HostBase.java

@@ -137,7 +137,7 @@ public class HostBase extends EppResource {
   }
 
   @Override
-  public Builder asBuilder() {
+  public Builder<? extends HostBase, ?> asBuilder() {
     return new Builder<>(clone(this));
   }
 

core/src/main/java/google/registry/model/host/HostHistory.java

@@ -111,7 +111,7 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getHostBase();
+    return getHostBase().map(hostBase -> new HostResource.Builder().copyFrom(hostBase).build());
   }
 
   @PostLoad

core/src/main/java/google/registry/model/host/HostResource.java

@@ -63,5 +63,21 @@ public class HostResource extends HostBase
     private Builder(HostResource instance) {
       super(instance);
     }
+
+    public Builder copyFrom(HostBase hostBase) {
+      return this.setCreationClientId(hostBase.getCreationClientId())
+          .setCreationTime(hostBase.getCreationTime())
+          .setDeletionTime(hostBase.getDeletionTime())
+          .setHostName(hostBase.getHostName())
+          .setInetAddresses(hostBase.getInetAddresses())
+          .setLastTransferTime(hostBase.getLastTransferTime())
+          .setLastSuperordinateChange(hostBase.getLastSuperordinateChange())
+          .setLastEppUpdateClientId(hostBase.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(hostBase.getLastEppUpdateTime())
+          .setPersistedCurrentSponsorClientId(hostBase.getPersistedCurrentSponsorClientId())
+          .setRepoId(hostBase.getRepoId())
+          .setSuperordinateDomain(hostBase.getSuperordinateDomain())
+          .setStatusValues(hostBase.getStatusValues());
+    }
   }
 }

core/src/main/java/google/registry/rde/PendingDeposit.java

@@ -17,8 +17,17 @@ package google.registry.rde;
 import com.google.auto.value.AutoValue;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.rde.RdeMode;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.Serializable;
 import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.BooleanCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
@@ -69,17 +78,9 @@ public abstract class PendingDeposit implements Serializable {
   @Nullable
   public abstract Integer revision();
 
-  static PendingDeposit create(
+  public static PendingDeposit create(
       String tld, DateTime watermark, RdeMode mode, CursorType cursor, Duration interval) {
-    return new AutoValue_PendingDeposit(
-        false,
-        tld,
-        watermark,
-        mode,
-        cursor,
-        interval,
-        null,
-        null);
+    return new AutoValue_PendingDeposit(false, tld, watermark, mode, cursor, interval, null, null);
   }
 
   static PendingDeposit createInManualOperation(
@@ -89,15 +90,52 @@ public abstract class PendingDeposit implements Serializable {
       String directoryWithTrailingSlash,
       @Nullable Integer revision) {
     return new AutoValue_PendingDeposit(
-        true,
-        tld,
-        watermark,
-        mode,
-        null,
-        null,
-        directoryWithTrailingSlash,
-        revision);
+        true, tld, watermark, mode, null, null, directoryWithTrailingSlash, revision);
   }
 
   PendingDeposit() {}
+
+  /**
+   * A deterministic coder for {@link PendingDeposit} used during a GroupBy transform.
+   *
+   * <p>We cannot use a {@link SerializableCoder} directly because it does not guarantee
+   * determinism, which is required by GroupBy.
+   */
+  public static class PendingDepositCoder extends AtomicCoder<PendingDeposit> {
+
+    private PendingDepositCoder() {
+      super();
+    }
+
+    private static final PendingDepositCoder INSTANCE = new PendingDepositCoder();
+
+    public static PendingDepositCoder of() {
+      return INSTANCE;
+    }
+
+    @Override
+    public void encode(PendingDeposit value, OutputStream outStream) throws IOException {
+      BooleanCoder.of().encode(value.manual(), outStream);
+      StringUtf8Coder.of().encode(value.tld(), outStream);
+      SerializableCoder.of(DateTime.class).encode(value.watermark(), outStream);
+      SerializableCoder.of(RdeMode.class).encode(value.mode(), outStream);
+      NullableCoder.of(SerializableCoder.of(CursorType.class)).encode(value.cursor(), outStream);
+      NullableCoder.of(SerializableCoder.of(Duration.class)).encode(value.interval(), outStream);
+      NullableCoder.of(StringUtf8Coder.of()).encode(value.directoryWithTrailingSlash(), outStream);
+      NullableCoder.of(VarIntCoder.of()).encode(value.revision(), outStream);
+    }
+
+    @Override
+    public PendingDeposit decode(InputStream inStream) throws IOException {
+      return new AutoValue_PendingDeposit(
+          BooleanCoder.of().decode(inStream),
+          StringUtf8Coder.of().decode(inStream),
+          SerializableCoder.of(DateTime.class).decode(inStream),
+          SerializableCoder.of(RdeMode.class).decode(inStream),
+          NullableCoder.of(SerializableCoder.of(CursorType.class)).decode(inStream),
+          NullableCoder.of(SerializableCoder.of(Duration.class)).decode(inStream),
+          NullableCoder.of(StringUtf8Coder.of()).decode(inStream),
+          NullableCoder.of(VarIntCoder.of()).decode(inStream));
+    }
+  }
 }

core/src/main/java/google/registry/rde/RdeFragmenter.java

@@ -0,0 +1,109 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.rde;
+
+import static google.registry.model.EppResourceUtils.loadAtPointInTime;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.host.HostResource;
+import google.registry.model.rde.RdeMode;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.function.Supplier;
+import org.joda.time.DateTime;
+
+/** Loading cache that turns a resource into XML for the various points in time and modes. */
+public class RdeFragmenter {
+  private final Map<WatermarkModePair, Optional<DepositFragment>> cache = new HashMap<>();
+  private final ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes;
+  private final RdeMarshaller marshaller;
+
+  long cacheHits = 0;
+  long resourcesNotFound = 0;
+  long resourcesFound = 0;
+
+  public RdeFragmenter(
+      ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes, RdeMarshaller marshaller) {
+    this.resourceAtTimes = resourceAtTimes;
+    this.marshaller = marshaller;
+  }
+
+  public Optional<DepositFragment> marshal(DateTime watermark, RdeMode mode) {
+    Optional<DepositFragment> result = cache.get(WatermarkModePair.create(watermark, mode));
+    if (result != null) {
+      cacheHits++;
+      return result;
+    }
+    EppResource resource = resourceAtTimes.get(watermark).get();
+    if (resource == null) {
+      result = Optional.empty();
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      resourcesNotFound++;
+      return result;
+    }
+    resourcesFound++;
+    if (resource instanceof DomainBase) {
+      result = Optional.of(marshaller.marshalDomain((DomainBase) resource, mode));
+      cache.put(WatermarkModePair.create(watermark, mode), result);
+      return result;
+    } else if (resource instanceof ContactResource) {
+      result = Optional.of(marshaller.marshalContact((ContactResource) resource));
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      return result;
+    } else if (resource instanceof HostResource) {
+      HostResource host = (HostResource) resource;
+      result =
+          Optional.of(
+              host.isSubordinate()
+                  ? marshaller.marshalSubordinateHost(
+                      host,
+                      // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
+                      // us.
+                      Objects.requireNonNull(
+                          loadAtPointInTime(
+                              tm().loadByKey(host.getSuperordinateDomain()), watermark)))
+                  : marshaller.marshalExternalHost(host));
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      return result;
+    } else {
+      throw new IllegalStateException(
+          String.format(
+              "Resource %s of type %s cannot be converted to XML.",
+              resource, resource.getClass().getSimpleName()));
+    }
+  }
+
+  /** Map key for {@link RdeFragmenter} cache. */
+  @AutoValue
+  abstract static class WatermarkModePair {
+    abstract DateTime watermark();
+
+    abstract RdeMode mode();
+
+    static WatermarkModePair create(DateTime watermark, RdeMode mode) {
+      return new AutoValue_RdeFragmenter_WatermarkModePair(watermark, mode);
+    }
+  }
+}

core/src/main/java/google/registry/rde/RdeStagingMapper.java

@@ -16,12 +16,10 @@ package google.registry.rde;
 
 import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.EppResourceUtils.loadAtPointInTime;
 import static google.registry.model.EppResourceUtils.loadAtPointInTimeAsync;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.tools.mapreduce.Mapper;
-import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
@@ -34,9 +32,6 @@ import google.registry.model.host.HostResource;
 import google.registry.model.rde.RdeMode;
 import google.registry.model.registrar.Registrar;
 import google.registry.xml.ValidationMode;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;
@@ -129,7 +124,7 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
         ImmutableMap.copyOf(Maps.asMap(dates, input -> loadAtPointInTimeAsync(resource, input)));
 
     // Convert resource to an XML fragment for each watermark/mode pair lazily and cache the result.
-    Fragmenter fragmenter = new Fragmenter(resourceAtTimes);
+    RdeFragmenter fragmenter = new RdeFragmenter(resourceAtTimes, marshaller);
 
     // Emit resource as an XML fragment for all TLDs and modes pending deposit.
     long resourcesEmitted = 0;
@@ -157,74 +152,4 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
     // Avoid running out of memory.
     tm().clearSessionCache();
   }
-
-  /** Loading cache that turns a resource into XML for the various points in time and modes. */
-  private class Fragmenter {
-    private final Map<WatermarkModePair, Optional<DepositFragment>> cache = new HashMap<>();
-    private final ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes;
-
-    long cacheHits = 0;
-    long resourcesNotFound = 0;
-    long resourcesFound = 0;
-
-    Fragmenter(ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes) {
-      this.resourceAtTimes = resourceAtTimes;
-    }
-
-    Optional<DepositFragment> marshal(DateTime watermark, RdeMode mode) {
-      Optional<DepositFragment> result = cache.get(WatermarkModePair.create(watermark, mode));
-      if (result != null) {
-        cacheHits++;
-        return result;
-      }
-      EppResource resource = resourceAtTimes.get(watermark).get();
-      if (resource == null) {
-        result = Optional.empty();
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        resourcesNotFound++;
-        return result;
-      }
-      resourcesFound++;
-      if (resource instanceof DomainBase) {
-        result = Optional.of(marshaller.marshalDomain((DomainBase) resource, mode));
-        cache.put(WatermarkModePair.create(watermark, mode), result);
-        return result;
-      } else if (resource instanceof ContactResource) {
-        result = Optional.of(marshaller.marshalContact((ContactResource) resource));
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        return result;
-      } else if (resource instanceof HostResource) {
-        HostResource host = (HostResource) resource;
-        result =
-            Optional.of(
-                host.isSubordinate()
-                    ? marshaller.marshalSubordinateHost(
-                        host,
-                        // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
-                        // us.
-                        Objects.requireNonNull(
-                            loadAtPointInTime(
-                                tm().loadByKey(host.getSuperordinateDomain()), watermark)))
-                    : marshaller.marshalExternalHost(host));
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        return result;
-      } else {
-        throw new AssertionError(resource.toString());
-      }
-    }
-  }
-
-  /** Map key for {@link Fragmenter} cache. */
-  @AutoValue
-  abstract static class WatermarkModePair {
-    abstract DateTime watermark();
-    abstract RdeMode mode();
-
-    static WatermarkModePair create(DateTime watermark, RdeMode mode) {
-      return new AutoValue_RdeStagingMapper_WatermarkModePair(watermark, mode);
-    }
-  }
 }

core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json

@@ -0,0 +1,23 @@
+{
+  "name": "RDE/BRDA Deposit Generation",
+  "description": "An Apache Beam pipeline generates RDE or BRDA deposits and deposits them to GCS with GhostRyde encryption.",
+  "parameters": [
+    {
+      "name": "pendings",
+      "label": "The pendings deposits to generate.",
+      "helpText": "A TLD to PendingDeposit map that is serialized and Base64 URL-safe encoded.",
+      "regexes": [
+        "A-Za-z0-9\\-_"
+      ]
+    },
+    {
+      "name": "validationMode",
+      "label": "How strict the marshaller validates the given EPP resources.",
+      "helpText": "If set to LENIENT the marshaller will not warn about missing data on the EPP resources.",
+      "is_optional": true,
+      "regexes": [
+        "^STRICT|LENIENT$"
+      ]
+    }
+  ]
+}