End-to-end Datastore to SQL pipeline (#707)

* End-to-end Datastore to SQL pipeline

Defined InitSqlPipeline that performs end-to-end migration from
a Datastore backup to a SQL database.

Also fixed/refined multiple tests related to this migration.

core/src/main/java/google/registry/backup/AppEngineEnvironment.java

@@ -22,21 +22,35 @@ import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 
 /**
- * Sets up a placeholder {@link Environment} on a non-AppEngine platform so that Datastore Entities
- * can be converted from/to Objectify entities. See {@code DatastoreEntityExtension} in test source
- * for more information.
+ * Sets up a fake {@link Environment} so that the following operations can be performed without the
+ * Datastore service:
+ *
+ * <ul>
+ *   <li>Create Objectify {@code Keys}.
+ *   <li>Instantiate Objectify objects.
+ *   <li>Convert Datastore {@code Entities} to their corresponding Objectify objects.
+ * </ul>
+ *
+ * <p>User has the option to specify their desired {@code appId} string, which forms part of an
+ * Objectify {@code Key} and is included in the equality check. This feature makes it easy to
+ * compare a migrated object in SQL with the original in Objectify.
+ *
+ * <p>Note that conversion from Objectify objects to Datastore {@code Entities} still requires the
+ * Datastore service.
  */
 public class AppEngineEnvironment implements Closeable {
 
-  private static final Environment PLACEHOLDER_ENV = createAppEngineEnvironment();
-
   private boolean isPlaceHolderNeeded;
 
   public AppEngineEnvironment() {
+    this("PlaceholderAppId");
+  }
+
+  public AppEngineEnvironment(String appId) {
     isPlaceHolderNeeded = ApiProxy.getCurrentEnvironment() == null;
     // isPlaceHolderNeeded may be true when we are invoked in a test with AppEngineRule.
     if (isPlaceHolderNeeded) {
-      ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+      ApiProxy.setEnvironmentForCurrentThread(createAppEngineEnvironment(appId));
     }
   }
 
@@ -48,7 +62,7 @@ public class AppEngineEnvironment implements Closeable {
   }
 
   /** Returns a placeholder {@link Environment} that can return hardcoded AppId and Attributes. */
-  private static Environment createAppEngineEnvironment() {
+  private static Environment createAppEngineEnvironment(String appId) {
     return (Environment)
         Proxy.newProxyInstance(
             Environment.class.getClassLoader(),
@@ -56,7 +70,7 @@ public class AppEngineEnvironment implements Closeable {
             (Object proxy, Method method, Object[] args) -> {
               switch (method.getName()) {
                 case "getAppId":
-                  return "PlaceholderAppId";
+                  return appId;
                 case "getAttributes":
                   return ImmutableMap.<String, Object>of();
                 default:

core/src/main/java/google/registry/beam/initsql/DomainBaseUtil.java

@@ -0,0 +1,75 @@
+// Copyright 2020 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.initsql;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.appengine.api.datastore.Entity;
+import java.util.Objects;
+
+/** Helper for manipulating {@code DomainBase} when migrating from Datastore to SQL database */
+final class DomainBaseUtil {
+
+  private DomainBaseUtil() {}
+
+  /**
+   * Removes {@link google.registry.model.billing.BillingEvent.Recurring}, {@link
+   * google.registry.model.poll.PollMessage PollMessages} and {@link
+   * google.registry.model.host.HostResource name servers} from a Datastore {@link Entity} that
+   * represents an Ofy {@link google.registry.model.domain.DomainBase}. This breaks the cycle of
+   * foreign key constraints between these entity kinds, allowing {@code DomainBases} to be inserted
+   * into the SQL database. See {@link InitSqlPipeline} for a use case, where the full {@code
+   * DomainBases} are written again during the last stage of the pipeline.
+   *
+   * <p>The returned object may be in bad state. Specifically, {@link
+   * google.registry.model.eppcommon.StatusValue#INACTIVE} is not added after name servers are
+   * removed. This only impacts tests.
+   *
+   * <p>This operation is performed on an Datastore {@link Entity} instead of Ofy Java object
+   * because Objectify requires access to a Datastore service when converting an Ofy object to a
+   * Datastore {@code Entity}. If we insist on working with Objectify objects, we face a few
+   * unsatisfactory options:
+   *
+   * <ul>
+   *   <li>Connect to our production Datastore, which incurs unnecessary security and code health
+   *       risk.
+   *   <li>Connect to a separate real Datastore instance, which is a waster and overkill.
+   *   <li>Use an in-memory test Datastore, which is a project health risk in that the test
+   *       Datastore would be added to Nomulus' production binary unless we create a separate
+   *       project for this pipeline.
+   * </ul>
+   *
+   * <p>Given our use case, operating on Datastore entities is the best option.
+   *
+   * @throws IllegalArgumentException if input does not represent a DomainBase
+   */
+  static Entity removeBillingAndPollAndHosts(Entity domainBase) {
+    checkNotNull(domainBase, "domainBase");
+    checkArgument(
+        Objects.equals(domainBase.getKind(), "DomainBase"),
+        "Expecting DomainBase, got %s",
+        domainBase.getKind());
+    Entity clone = domainBase.clone();
+    clone.removeProperty("autorenewBillingEvent");
+    clone.removeProperty("autorenewPollMessage");
+    clone.removeProperty("deletePollMessage");
+    clone.removeProperty("nsHosts");
+    domainBase.getProperties().keySet().stream()
+        .filter(s -> s.startsWith("transferData."))
+        .forEach(s -> clone.removeProperty(s));
+    return clone;
+  }
+}

core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java

@@ -0,0 +1,260 @@
+// Copyright 2020 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.initsql;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.backup.AppEngineEnvironment;
+import google.registry.backup.VersionedEntity;
+import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.Transforms.RemoveDomainBaseForeignKeys;
+import google.registry.beam.initsql.Transforms.SerializableSupplier;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.host.HostResource;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.RegistrarContact;
+import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import java.io.Serializable;
+import java.util.Collection;
+import java.util.Optional;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.Wait;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.joda.time.DateTime;
+
+/**
+ * A BEAM pipeline that populates a SQL database with data from a Datastore backup.
+ *
+ * <p>This pipeline migrates EPP resources and related entities that cross-reference each other. To
+ * avoid violating foreign key constraints, writes to SQL are ordered by entity kinds. In addition,
+ * the {@link DomainBase} kind is written twice (see details below). The write order is presented
+ * below. Although some kinds can be written concurrently, e.g. {@code ContactResource} and {@code
+ * RegistrarContact}, we do not expect any performance benefit since the limiting resource is the
+ * number of JDBC connections. Google internal users may refer to <a
+ * href="http://go/registry-r3-init-sql">the design doc</a> for more information.
+ *
+ * <ol>
+ *   <li>{@link Registry}: Assumes that {@code PremiumList} and {@code ReservedList} have been set
+ *       up in the SQL database.
+ *   <li>{@link Registrar}: Logically depends on {@code Registry}, Foreign key not modeled yet.
+ *   <li>{@link ContactResource}: references {@code Registrar}
+ *   <li>{@link RegistrarContact}: references {@code Registrar}.
+ *   <li>Cleansed {@link DomainBase}: with references to {@code BillingEvent}, {@code Recurring},
+ *       {@code Cancellation} and {@code HostResource} removed, still references {@code Registrar}
+ *       and {@code ContactResource}. The removal breaks circular Foreign Key references.
+ *   <li>{@link HostResource}: references {@code DomainBase}.
+ *   <li>{@link HistoryEntry}: maps to one of three SQL entity types and may reference {@code
+ *       Registrar}, {@code ContactResource}, {@code HostResource}, and {@code DomainBase}.
+ *   <li>{@link AllocationToken}: references {@code HistoryEntry}.
+ *   <li>{@link BillingEvent.Recurring}: references {@code Registrar}, {@code DomainBase} and {@code
+ *       HistoryEntry}.
+ *   <li>{@link BillingEvent.OneTime}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       BillingEvent.Recurring}, {@code HistoryEntry} and {@code AllocationToken}.
+ *   <li>{@link BillingEvent.Modification}: SQL model TBD. Will reference {@code Registrar}, {@code
+ *       DomainBase} and {@code BillingEvent.OneTime}.
+ *   <li>{@link BillingEvent.Cancellation}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       BillingEvent.Recurring}, {@code BillingEvent.OneTime}, and {@code HistoryEntry}.
+ *   <li>{@link PollMessage}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       ContactResource}, {@code HostResource}, and {@code HistoryEntry}.
+ *   <li>{@link DomainBase}, original copy from Datastore.
+ * </ol>
+ */
+public class InitSqlPipeline implements Serializable {
+
+  /**
+   * Datastore kinds to be written to the SQL database before the cleansed version of {@link
+   * DomainBase}.
+   */
+  // TODO(weiminyu): include Registry.class when it is modeled in JPA.
+  private static final ImmutableList<Class<?>> PHASE_ONE_ORDERED =
+      ImmutableList.of(Registrar.class, ContactResource.class);
+
+  /**
+   * Datastore kinds to be written to the SQL database after the cleansed version of {@link
+   * DomainBase}.
+   *
+   * <p>The following entities are missing from the list:
+   *
+   * <ul>
+   *   <li>Those not modeled in JPA yet, e.g., {@code BillingEvent.Modification}.
+   *   <li>Those waiting for sanitation, e.g., {@code HistoryEntry}, which would have duplicate keys
+   *       after converting to SQL model.
+   *   <li>Those that have foreign key constraints on the above.
+   * </ul>
+   */
+  // TODO(weiminyu): add more entities when available.
+  private static final ImmutableList<Class<?>> PHASE_TWO_ORDERED =
+      ImmutableList.of(HostResource.class);
+
+  private final InitSqlPipelineOptions options;
+
+  private final Pipeline pipeline;
+
+  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
+      jpaGetter;
+
+  InitSqlPipeline(InitSqlPipelineOptions options) {
+    this.options = options;
+    pipeline = Pipeline.create(options);
+    jpaGetter = JpaTransactionManagerComponent::cloudSqlJpaTransactionManager;
+  }
+
+  @VisibleForTesting
+  InitSqlPipeline(InitSqlPipelineOptions options, Pipeline pipeline) {
+    this.options = options;
+    this.pipeline = pipeline;
+    jpaGetter = JpaTransactionManagerComponent::localDbJpaTransactionManager;
+  }
+
+  public PipelineResult run() {
+    setupPipeline();
+    return pipeline.run();
+  }
+
+  @VisibleForTesting
+  void setupPipeline() {
+    PCollectionTuple datastoreSnapshot =
+        pipeline.apply(
+            "Load Datastore snapshot",
+            Transforms.loadDatastoreSnapshot(
+                options.getDatastoreExportDir(),
+                options.getCommitLogDir(),
+                DateTime.parse(options.getCommitLogStartTimestamp()),
+                DateTime.parse(options.getCommitLogEndTimestamp()),
+                ImmutableSet.<String>builder()
+                    .add("DomainBase")
+                    .addAll(toKindStrings(PHASE_ONE_ORDERED))
+                    .addAll(toKindStrings(PHASE_TWO_ORDERED))
+                    .build()));
+
+    // Set up the pipeline to write entity kinds from PHASE_ONE_ORDERED to SQL. Return a object
+    // that signals the completion of the phase.
+    PCollection<Void> blocker =
+        scheduleOnePhaseWrites(datastoreSnapshot, PHASE_ONE_ORDERED, Optional.empty(), null);
+    blocker =
+        writeToSql(
+            "DomainBase without circular foreign keys",
+            removeDomainBaseForeignKeys(datastoreSnapshot)
+                .apply("Wait on phase one", Wait.on(blocker)));
+    // Set up the pipeline to write entity kinds from PHASE_TWO_ORDERED to SQL. This phase won't
+    // start until all cleansed DomainBases have been written (started by line above).
+    scheduleOnePhaseWrites(
+        datastoreSnapshot, PHASE_TWO_ORDERED, Optional.of(blocker), "DomainBaseNoFkeys");
+  }
+
+  private PCollection<VersionedEntity> removeDomainBaseForeignKeys(
+      PCollectionTuple datastoreSnapshot) {
+    PCollection<VersionedEntity> domainBases =
+        datastoreSnapshot.get(Transforms.createTagForKind("DomainBase"));
+    return domainBases.apply(
+        "Remove circular foreign keys from DomainBase",
+        ParDo.of(new RemoveDomainBaseForeignKeys()));
+  }
+
+  /**
+   * Sets up the pipeline to write entities in {@code entityClasses} to SQL. Entities are written
+   * one kind at a time based on each kind's position in {@code entityClasses}. Concurrency exists
+   * within each kind.
+   *
+   * @param datastoreSnapshot the Datastore snapshot of all data to be migrated to SQL
+   * @param entityClasses the entity types in write order
+   * @param blockingPCollection the pipeline stage that blocks this phase
+   * @param blockingTag description of the stage (if exists) that blocks this phase. Needed for
+   *     generating unique transform ids
+   * @return the output {@code PCollection} from the writing of the last entity kind. Other parts of
+   *     the pipeline can {@link Wait} on this object
+   */
+  private PCollection<Void> scheduleOnePhaseWrites(
+      PCollectionTuple datastoreSnapshot,
+      Collection<Class<?>> entityClasses,
+      Optional<PCollection<Void>> blockingPCollection,
+      String blockingTag) {
+    checkArgument(!entityClasses.isEmpty(), "Each phase must have at least one kind.");
+    ImmutableList<TupleTag<VersionedEntity>> tags =
+        toKindStrings(entityClasses).stream()
+            .map(Transforms::createTagForKind)
+            .collect(ImmutableList.toImmutableList());
+
+    PCollection<Void> prev = blockingPCollection.orElse(null);
+    String prevTag = blockingTag;
+    for (TupleTag<VersionedEntity> tag : tags) {
+      PCollection<VersionedEntity> curr = datastoreSnapshot.get(tag);
+      if (prev != null) {
+        curr = curr.apply("Wait on " + prevTag, Wait.on(prev));
+      }
+      prev = writeToSql(tag.getId(), curr);
+      prevTag = tag.getId();
+    }
+    return prev;
+  }
+
+  private PCollection<Void> writeToSql(String transformId, PCollection<VersionedEntity> data) {
+    String credentialFileUrl =
+        options.getSqlCredentialUrlOverride() != null
+            ? options.getSqlCredentialUrlOverride()
+            : BackupPaths.getCloudSQLCredentialFilePatterns(options.getEnvironment()).get(0);
+
+    return data.apply(
+        "Write to sql: " + transformId,
+        Transforms.writeToSql(
+            transformId,
+            options.getMaxConcurrentSqlWriters(),
+            options.getSqlWriteBatchSize(),
+            new JpaSupplierFactory(credentialFileUrl, jpaGetter)));
+  }
+
+  private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {
+    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+      return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList());
+    }
+  }
+
+  static class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
+    private static final long serialVersionUID = 1L;
+
+    private String credentialFileUrl;
+    private SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter;
+
+    JpaSupplierFactory(
+        String credentialFileUrl,
+        SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
+      this.credentialFileUrl = credentialFileUrl;
+      this.jpaGetter = jpaGetter;
+    }
+
+    @Override
+    public JpaTransactionManager get() {
+      return jpaGetter.apply(
+          DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
+              .beamJpaModule(new BeamJpaModule(credentialFileUrl))
+              .build());
+    }
+  }
+}

core/src/main/java/google/registry/beam/initsql/InitSqlPipelineOptions.java

@@ -0,0 +1,84 @@
+// Copyright 2020 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.initsql;
+
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.Validation;
+
+/** Pipeline options for {@link InitSqlPipeline} */
+public interface InitSqlPipelineOptions extends GcpOptions {
+
+  @Description(
+      "Overrides the URL to the SQL credential file. " + "Required if environment is not provided.")
+  @Nullable
+  String getSqlCredentialUrlOverride();
+
+  void setSqlCredentialUrlOverride(String credentialUrlOverride);
+
+  @Description("The root directory of the export to load.")
+  String getDatastoreExportDir();
+
+  void setDatastoreExportDir(String datastoreExportDir);
+
+  @Description("The directory that contains all CommitLog files.")
+  String getCommitLogDir();
+
+  void setCommitLogDir(String commitLogDir);
+
+  @Description("The earliest CommitLogs to load, in ISO8601 format.")
+  @Validation.Required
+  String getCommitLogStartTimestamp();
+
+  void setCommitLogStartTimestamp(String commitLogStartTimestamp);
+
+  @Description("The latest CommitLogs to load, in ISO8601 format.")
+  @Validation.Required
+  String getCommitLogEndTimestamp();
+
+  void setCommitLogEndTimestamp(String commitLogEndTimestamp);
+
+  @Description(
+      "The deployed environment, alpha, crash, sandbox, or production. "
+          + "Not required only if sqlCredentialUrlOverride is provided.")
+  @Nullable
+  String getEnvironment();
+
+  void setEnvironment(String environment);
+
+  @Description(
+      "The maximum JDBC connection pool size on a VM. "
+          + "This value should be equal to or greater than the number of cores on the VM.")
+  @Default.Integer(4)
+  int getJdbcMaxPoolSize();
+
+  void setJdbcMaxPoolSize(int jdbcMaxPoolSize);
+
+  @Description(
+      "A hint to the pipeline runner of the maximum number of concurrent SQL writers to create. "
+          + "Note that multiple writers may run on the same VM and share the connection pool.")
+  @Default.Integer(4)
+  int getMaxConcurrentSqlWriters();
+
+  void setMaxConcurrentSqlWriters(int maxConcurrentSqlWriters);
+
+  @Description("The number of entities to be written to the SQL database in one transaction.")
+  @Default.Integer(20)
+  int getSqlWriteBatchSize();
+
+  void setSqlWriteBatchSize(int sqlWriteBatchSize);
+}

core/src/main/java/google/registry/beam/initsql/README.md

@@ -1,3 +1,17 @@
 ## Summary
 
-This package contains a BEAM pipeline that populates a Cloud SQL database from a Datastore backup.
+This package contains a BEAM pipeline that populates a Cloud SQL database from a
+Datastore backup. The pipeline uses an unsynchronized Datastore export and
+overlapping CommitLogs generated by the Nomulus server to recreate a consistent
+Datastore snapshot, and writes the data to a Cloud SQL instance.
+
+## Pipeline Visualization
+
+The golden flow graph of the InitSqlPipeline is saved both as a text-base
+[DOT file](../../../../../../test/resources/google/registry/beam/initsql/pipeline_golden.dot)
+and a
+[.png file](../../../../../../test/resources/google/registry/beam/initsql/pipeline_golden.png).
+A test compares the flow graph of the current pipeline with the golden graph,
+and will fail if changes are detected. When this happens, run the Gradle task
+':core:updateInitSqlPipelineGraph' to update the golden files and review the
+changes.

core/src/main/java/google/registry/beam/initsql/Transforms.java

@@ -17,8 +17,10 @@ package google.registry.beam.initsql;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Throwables.throwIfUnchecked;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
+import static google.registry.persistence.JpaRetries.isFailedTxnRetriable;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -29,14 +31,16 @@ import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
 
 import avro.shaded.com.google.common.collect.Iterators;
+import com.google.appengine.api.datastore.Entity;
+import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Streams;
 import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.CommitLogImports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.ofy.Ofy;
 import google.registry.persistence.transaction.JpaTransactionManager;
@@ -49,7 +53,6 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
-import javax.persistence.OptimisticLockException;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.Compression;
 import org.apache.beam.sdk.io.FileIO;
@@ -70,7 +73,6 @@ import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.TypeDescriptor;
@@ -225,7 +227,7 @@ public final class Transforms {
     return new PTransform<PCollection<String>, PCollection<Metadata>>() {
       @Override
       public PCollection<Metadata> expand(PCollection<String> input) {
-        return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.DISALLOW));
+        return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.ALLOW));
       }
     };
   }
@@ -263,6 +265,11 @@ public final class Transforms {
 
   /**
    * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
+   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
+   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
+   *
+   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
+   * job, we will not add features unless proven necessary.
    *
    * @param transformId a unique ID for an instance of the returned transform
    * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
@@ -270,22 +277,21 @@ public final class Transforms {
    * @param batchSize the number of entities to write in each operation
    * @param jpaSupplier supplier of a {@link JpaTransactionManager}
    */
-  public static PTransform<PCollection<VersionedEntity>, PDone> writeToSql(
+  public static PTransform<PCollection<VersionedEntity>, PCollection<Void>> writeToSql(
       String transformId,
       int maxWriters,
       int batchSize,
       SerializableSupplier<JpaTransactionManager> jpaSupplier) {
-    return new PTransform<PCollection<VersionedEntity>, PDone>() {
+    return new PTransform<PCollection<VersionedEntity>, PCollection<Void>>() {
       @Override
-      public PDone expand(PCollection<VersionedEntity> input) {
-        input
+      public PCollection<Void> expand(PCollection<VersionedEntity> input) {
+        return input
             .apply(
                 "Shard data for " + transformId,
                 MapElements.into(kvs(integers(), TypeDescriptor.of(VersionedEntity.class)))
                     .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
             .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
             .apply("Write in batch for " + transformId, ParDo.of(new SqlBatchWriter(jpaSupplier)));
-        return PDone.in(input.getPipeline());
       }
     };
   }
@@ -397,8 +403,10 @@ public final class Transforms {
     public void setup() {
       sleeper = new SystemSleeper();
 
-      ObjectifyService.initOfy();
-      ofy = ObjectifyService.ofy();
+      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+        ObjectifyService.initOfy();
+        ofy = ObjectifyService.ofy();
+      }
 
       synchronized (SqlBatchWriter.class) {
         if (instanceCount == 0) {
@@ -444,7 +452,10 @@ public final class Transforms {
           runnable.run();
           return;
         } catch (Throwable throwable) {
-          throwIfNotCausedBy(throwable, OptimisticLockException.class);
+          if (!isFailedTxnRetriable(throwable)) {
+            throwIfUnchecked(throwable);
+            throw new RuntimeException(throwable);
+          }
           int sleepMillis = (1 << attempt) * initialDelayMillis;
           int jitter =
               ThreadLocalRandom.current().nextInt((int) (sleepMillis * jitterRatio))
@@ -453,21 +464,28 @@ public final class Transforms {
         }
       }
     }
+  }
 
-    /**
-     * Rethrows {@code throwable} if it is not (and does not have a cause of) {@code causeType};
-     * otherwise returns with no side effects.
-     */
-    private void throwIfNotCausedBy(Throwable throwable, Class<? extends Throwable> causeType) {
-      Throwable t = throwable;
-      while (t != null) {
-        if (causeType.isInstance(t)) {
-          return;
-        }
-        t = t.getCause();
-      }
-      Throwables.throwIfUnchecked(t);
-      throw new RuntimeException(t);
+  /**
+   * Removes BillingEvents, {@link google.registry.model.poll.PollMessage PollMessages} and {@link
+   * google.registry.model.host.HostResource} from a {@link DomainBase}. These are circular foreign
+   * key constraints that prevent migration of {@code DomainBase} to SQL databases.
+   *
+   * <p>See {@link InitSqlPipeline} for more information.
+   */
+  static class RemoveDomainBaseForeignKeys extends DoFn<VersionedEntity, VersionedEntity> {
+
+    @ProcessElement
+    public void processElement(
+        @Element VersionedEntity domainBase, OutputReceiver<VersionedEntity> out) {
+      checkArgument(
+          domainBase.getEntity().isPresent(), "Unexpected delete entity %s", domainBase.key());
+      Entity outputEntity =
+          DomainBaseUtil.removeBillingAndPollAndHosts(domainBase.getEntity().get());
+      out.output(
+          VersionedEntity.from(
+              domainBase.commitTimeMills(),
+              EntityTranslator.convertToPb(outputEntity).toByteArray()));
     }
   }
 }

core/src/main/java/google/registry/persistence/JpaRetries.java

@@ -0,0 +1,58 @@
+// Copyright 2019 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.persistence;
+
+import com.google.common.base.Predicates;
+import com.google.common.collect.ImmutableSet;
+import java.sql.SQLException;
+import java.util.function.Predicate;
+import javax.persistence.OptimisticLockException;
+
+/** Helpers for identifying retriable database operations. */
+public final class JpaRetries {
+
+  private JpaRetries() {}
+
+  private static final ImmutableSet<String> RETRIABLE_TXN_SQL_STATE =
+      ImmutableSet.of(
+          "40001", // serialization_failure
+          "40P01", // deadlock_detected, PSQL-specific
+          "55006", // object_in_use, PSQL and DB2
+          "55P03" // lock_not_available, PSQL-specific
+          );
+
+  private static final Predicate<Throwable> RETRIABLE_TXN_PREDICATE =
+      Predicates.or(
+          OptimisticLockException.class::isInstance,
+          e ->
+              e instanceof SQLException
+                  && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
+
+  public static boolean isFailedTxnRetriable(Throwable throwable) {
+    Throwable t = throwable;
+    while (t != null) {
+      if (RETRIABLE_TXN_PREDICATE.test(t)) {
+        return true;
+      }
+      t = t.getCause();
+    }
+    return false;
+  }
+
+  public static boolean isFailedQueryRetriable(Throwable throwable) {
+    // TODO(weiminyu): check for more error codes.
+    return isFailedTxnRetriable(throwable);
+  }
+}