Rewrite the JPA output connector for BEAM (#995)

* Rewrite the JPA output connector for BEAM

Following BEAM's IO connector style, added a RegistryJpaIO class to hold
IO connectors, and implemented the Write connector as a static inner
class in it. The JpaTransactionManager used by the Write connector
retrieves SQL credentials from the SecretManager.

Cleaned up SQL-related pipeline parameters.

Converted the InitSqlPipeline to use RegistryJpaIO.

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

@@ -0,0 +1,124 @@
+// 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.common;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.appengine.api.datastore.Entity;
+import com.google.common.collect.ImmutableList;
+import google.registry.backup.VersionedEntity;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.beam.initsql.BackupTestStore;
+import google.registry.beam.initsql.InitSqlTestUtils;
+import google.registry.beam.initsql.Transforms;
+import google.registry.model.ImmutableObject;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.registrar.Registrar;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectExtension;
+import java.io.Serializable;
+import java.nio.file.Path;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.transforms.Create;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Unit test for {@link RegistryJpaIO.Write}. */
+class RegistryJpaWriteTest implements Serializable {
+
+  private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
+
+  private final FakeClock fakeClock = new FakeClock(START_TIME);
+
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+
+  @RegisterExtension final transient InjectExtension injectRule = new InjectExtension();
+
+  @RegisterExtension
+  final transient JpaIntegrationTestExtension database =
+      new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();
+
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  private ImmutableList<Entity> contacts;
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    try (BackupTestStore store = new BackupTestStore(fakeClock)) {
+      injectRule.setStaticField(Ofy.class, "clock", fakeClock);
+
+      // Required for contacts created below.
+      Registrar ofyRegistrar = AppEngineExtension.makeRegistrar2();
+      store.insertOrUpdate(ofyRegistrar);
+      jpaTm().transact(() -> jpaTm().put(store.loadAsOfyEntity(ofyRegistrar)));
+
+      ImmutableList.Builder<Entity> builder = new ImmutableList.Builder<>();
+
+      for (int i = 0; i < 3; i++) {
+        ContactResource contact = DatabaseHelper.newContactResource("contact_" + i);
+        store.insertOrUpdate(contact);
+        builder.add(store.loadAsDatastoreEntity(contact));
+      }
+      contacts = builder.build();
+    }
+  }
+
+  @Test
+  void writeToSql_twoWriters() {
+    testPipeline
+        .apply(
+            Create.of(
+                contacts.stream()
+                    .map(InitSqlTestUtils::entityToBytes)
+                    .map(bytes -> VersionedEntity.from(0L, bytes))
+                    .collect(Collectors.toList())))
+        .apply(
+            RegistryJpaIO.<VersionedEntity>write()
+                .withName("ContactResource")
+                .withBatchSize(4)
+                .withShards(2)
+                .withJpaConverter(Transforms::convertVersionedEntityToSqlEntity));
+    testPipeline.run().waitUntilFinish();
+
+    ImmutableList<?> sqlContacts = jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class));
+    assertThat(sqlContacts)
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
+        .containsExactlyElementsIn(
+            contacts.stream()
+                .map(InitSqlTestUtils::datastoreToOfyEntity)
+                .map(ImmutableObject.class::cast)
+                .collect(ImmutableList.toImmutableList()));
+  }
+}

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

@@ -19,6 +19,7 @@ import static google.registry.beam.common.RegistryPipelineOptions.validateRegist
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import google.registry.config.RegistryEnvironment;
+import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.testing.SystemPropertyExtension;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.junit.jupiter.api.BeforeEach;
@@ -43,29 +44,42 @@ class RegistryPipelineOptionsTest {
 
   @Test
   void environment_fromArgs() {
-    assertThat(
-            PipelineOptionsFactory.fromArgs("--registryEnvironment=ALPHA")
-                .as(RegistryPipelineOptions.class)
-                .getRegistryEnvironment())
-        .isSameInstanceAs(RegistryEnvironment.ALPHA);
+    RegistryPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(
+                "--registryEnvironment=ALPHA", "--isolationOverride=TRANSACTION_SERIALIZABLE")
+            .withValidation()
+            .as(RegistryPipelineOptions.class);
+    assertThat(options.getRegistryEnvironment()).isSameInstanceAs(RegistryEnvironment.ALPHA);
+    assertThat(options.getIsolationOverride())
+        .isSameInstanceAs(TransactionIsolationLevel.TRANSACTION_SERIALIZABLE);
   }
 
   @Test
-  void environment_invalid() {
+  void environment_invalidEnvironment() {
     assertThrows(
         IllegalArgumentException.class,
         () ->
             PipelineOptionsFactory.fromArgs("--registryEnvironment=alpha")
+                .withValidation()
+                .as(RegistryPipelineOptions.class));
+  }
+
+  @Test
+  void environment_invalidIsolation() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
+            PipelineOptionsFactory.fromArgs("--isolationOverride=something_wrong")
+                .withValidation()
                 .as(RegistryPipelineOptions.class));
   }
 
   @Test
   void environment_undefined() {
-    assertThat(
-            PipelineOptionsFactory.create()
-                .as(RegistryPipelineOptions.class)
-                .getRegistryEnvironment())
-        .isNull();
+    RegistryPipelineOptions options =
+        PipelineOptionsFactory.fromArgs().withValidation().as(RegistryPipelineOptions.class);
+    assertThat(options.getRegistryEnvironment()).isNull();
+    assertThat(options.getIsolationOverride()).isNull();
   }
 
   @Test

core/src/test/java/google/registry/beam/initsql/BackupTestStore.java

@@ -47,7 +47,7 @@ import org.joda.time.format.DateTimeFormatter;
  * every transaction is invoked on this store, ensuring strictly increasing timestamps on causally
  * dependent transactions. In production, the same ordering is ensured by sleep and retry.
  */
-class BackupTestStore implements AutoCloseable {
+public final class BackupTestStore implements AutoCloseable {
 
   private static final DateTimeFormatter EXPORT_TIMESTAMP_FORMAT =
       DateTimeFormat.forPattern("yyyy-MM-dd'T'HH:mm:ss_SSS");
@@ -59,7 +59,7 @@ class BackupTestStore implements AutoCloseable {
 
   private CommitLogCheckpoint prevCommitLogCheckpoint;
 
-  BackupTestStore(FakeClock fakeClock) throws Exception {
+  public BackupTestStore(FakeClock fakeClock) throws Exception {
     this.fakeClock = fakeClock;
     this.appEngine =
         new AppEngineExtension.Builder()
@@ -88,7 +88,7 @@ class BackupTestStore implements AutoCloseable {
    * transaction.
    */
   @SafeVarargs
-  final long insertOrUpdate(Object... entities) {
+  public final long insertOrUpdate(Object... entities) {
     long timestamp = fakeClock.nowUtc().getMillis();
     tm().transact(() -> ofy().save().entities(entities).now());
     fakeClock.advanceOneMilli();
@@ -97,7 +97,7 @@ class BackupTestStore implements AutoCloseable {
 
   /** Deletes {@code entities} from the Datastore and returns the timestamp of this transaction. */
   @SafeVarargs
-  final long delete(Object... entities) {
+  public final long delete(Object... entities) {
     long timestamp = fakeClock.nowUtc().getMillis();
     tm().transact(() -> ofy().delete().entities(entities).now());
     fakeClock.advanceOneMilli();
@@ -111,7 +111,7 @@ class BackupTestStore implements AutoCloseable {
    * Objectify entity and want to find out the values of certain assign-on-persist properties. See
    * {@link VersionedEntity} for more information.
    */
-  Entity loadAsDatastoreEntity(Object ofyEntity) {
+  public Entity loadAsDatastoreEntity(Object ofyEntity) {
     try {
       return datastoreService.get(Key.create(ofyEntity).getRaw());
     } catch (EntityNotFoundException e) {
@@ -124,7 +124,7 @@ class BackupTestStore implements AutoCloseable {
    *
    * <p>See {@link #loadAsDatastoreEntity} and {@link VersionedEntity} for more information.
    */
-  Object loadAsOfyEntity(Object ofyEntity) {
+  public Object loadAsOfyEntity(Object ofyEntity) {
     try {
       return ofy().load().fromEntity(datastoreService.get(Key.create(ofyEntity).getRaw()));
     } catch (EntityNotFoundException e) {

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineGraphTest.java

@@ -38,7 +38,7 @@ class InitSqlPipelineGraphTest {
         "--commitLogEndTimestamp=2000-01-02TZ",
         "--datastoreExportDir=/somedir",
         "--commitLogDir=/someotherdir",
-        "--environment=alpha"
+        "--registryEnvironment=ALPHA"
       };
 
   private static final transient InitSqlPipelineOptions options =

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -317,8 +317,6 @@ class InitSqlPipelineTest {
   void runPipeline() {
     InitSqlPipelineOptions options =
         PipelineOptionsFactory.fromArgs(
-                "--sqlCredentialUrlOverride="
-                    + beamJpaExtension.getCredentialFile().getAbsolutePath(),
                 "--commitLogStartTimestamp=" + START_TIME,
                 "--commitLogEndTimestamp=" + fakeClock.nowUtc().plusMillis(1),
                 "--datastoreExportDir=" + exportDir.getAbsolutePath(),