Set up deployment of the Spec11 pipeline with JPA TM (#716)

* Set up deployment of the Spec11 pipeline with JPA TM

* Remove unnecessarily pipeline options setting

* Use enviroment name in BeamJpaModuleTest

* Fix checkstyle error

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

@@ -54,37 +54,39 @@ public class BeamJpaModule {
 
   private static final String GCS_SCHEME = "gs://";
 
-  @Nullable private final String credentialFilePath;
+  @Nullable private final String sqlAccessInfoFile;
+  @Nullable private final String cloudKmsProjectId;
 
   /**
    * Constructs a new instance of {@link BeamJpaModule}.
    *
    * <p>Note: it is an unfortunately necessary antipattern to check for the validity of
-   * credentialFilePath in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
+   * sqlAccessInfoFile in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
    * Unfortunately, this is a restriction imposed upon us by Dagger. Specifically, because we use
    * this in at least one 1 {@link google.registry.tools.RegistryTool} command(s), it must be
    * instantiated in {@code google.registry.tools.RegistryToolComponent} for all possible commands;
    * Dagger doesn't permit it to ever be null. For the vast majority of commands, it will never be
    * used (so a null credential file path is fine in those cases).
    *
-   * @param credentialFilePath the path to a Cloud SQL credential file. This must refer to either a
+   * @param sqlAccessInfoFile the path to a Cloud SQL credential file. This must refer to either a
    *     real encrypted file on GCS as returned by {@link
    *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
    *     with credentials to a test database.
    */
-  public BeamJpaModule(@Nullable String credentialFilePath) {
+  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
-    this.credentialFilePath = credentialFilePath;
+    this.sqlAccessInfoFile = sqlAccessInfoFile;
+    this.cloudKmsProjectId = cloudKmsProjectId;
   }
 
   /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
   private boolean isCloudSqlCredential() {
-    return credentialFilePath.startsWith(GCS_SCHEME);
+    return sqlAccessInfoFile.startsWith(GCS_SCHEME);
   }
 
   @Provides
   @Singleton
   SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
-    checkArgument(!isNullOrEmpty(credentialFilePath), "Null or empty credentialFilePath");
+    checkArgument(!isNullOrEmpty(sqlAccessInfoFile), "Null or empty credentialFilePath");
     String line = readOnlyLineFromCredentialFile();
     if (isCloudSqlCredential()) {
       line = lazyDecryptor.get().decrypt(line);
@@ -101,7 +103,7 @@ public class BeamJpaModule {
 
   String readOnlyLineFromCredentialFile() {
     try {
-      ResourceId resourceId = FileSystems.matchSingleFileSpec(credentialFilePath).resourceId();
+      ResourceId resourceId = FileSystems.matchSingleFileSpec(sqlAccessInfoFile).resourceId();
       try (BufferedReader reader =
           new BufferedReader(
               new InputStreamReader(
@@ -141,8 +143,8 @@ public class BeamJpaModule {
 
   @Provides
   @Config("beamCloudKmsProjectId")
-  static String kmsProjectId() {
+  String kmsProjectId() {
-    return "domain-registry-dev";
+    return cloudKmsProjectId;
   }
 
   @Provides

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

@@ -24,7 +24,6 @@ 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;
@@ -227,7 +226,7 @@ public class InitSqlPipeline implements Serializable {
             transformId,
             options.getMaxConcurrentSqlWriters(),
             options.getSqlWriteBatchSize(),
-            new JpaSupplierFactory(credentialFileUrl, jpaGetter)));
+            new JpaSupplierFactory(credentialFileUrl, options.getCloudKmsProjectId(), jpaGetter)));
   }
 
   private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {
@@ -235,26 +234,4 @@ public class InitSqlPipeline implements Serializable {
       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

@@ -60,6 +60,13 @@ public interface InitSqlPipelineOptions extends GcpOptions {
 
   void setEnvironment(String environment);
 
+  @Description("The GCP project that contains the keyring used for decrypting the "
+      + "SQL credential file.")
+  @Nullable
+  String getCloudKmsProjectId();
+
+  void setCloudKmsProjectId(String cloudKmsProjectId);
+
   @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.")

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

@@ -0,0 +1,48 @@
+// 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 google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.Transforms.SerializableSupplier;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+
+public class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
+
+  private static final long serialVersionUID = 1L;
+
+  private final String credentialFileUrl;
+  @Nullable private final String cloudKmsProjectId;
+  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
+      jpaGetter;
+
+  public JpaSupplierFactory(
+      String credentialFileUrl,
+      @Nullable String cloudKmsProjectId,
+      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
+    this.credentialFileUrl = credentialFileUrl;
+    this.cloudKmsProjectId = cloudKmsProjectId;
+    this.jpaGetter = jpaGetter;
+  }
+
+  @Override
+  public JpaTransactionManager get() {
+    return jpaGetter.apply(
+        DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
+            .beamJpaModule(new BeamJpaModule(credentialFileUrl, cloudKmsProjectId))
+            .build());
+  }
+}

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

@@ -16,17 +16,29 @@ package google.registry.tools;
 
 import com.beust.jcommander.Parameters;
 import google.registry.beam.spec11.Spec11Pipeline;
+import google.registry.config.CredentialModule.LocalCredential;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.util.GoogleCredentialsBundle;
+import google.registry.util.Retrier;
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 /** Nomulus command that deploys the {@link Spec11Pipeline} template. */
 @Parameters(commandDescription = "Deploy the Spec11 pipeline to GCS.")
 public class DeploySpec11PipelineCommand implements Command {
 
-  @Inject Spec11Pipeline spec11Pipeline;
+  @Inject @Config("projectId") String projectId;
+  @Inject @Config("beamStagingUrl") String beamStagingUrl;
+  @Inject @Config("spec11TemplateUrl")String spec11TemplateUrl;
+  @Inject @Config("reportingBucketUrl")String reportingBucketUrl;
+  @Inject @LocalCredential GoogleCredentialsBundle googleCredentialsBundle;
+  @Inject Retrier retrier;
+  @Inject @Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile;
 
   @Override
   public void run() {
-    spec11Pipeline.deploy();
+    Spec11Pipeline pipeline = new Spec11Pipeline(projectId, beamStagingUrl, spec11TemplateUrl,
+        reportingBucketUrl, googleCredentialsBundle, retrier);
+    pipeline.deploy();
   }
 }
-

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

@@ -29,7 +29,7 @@ import com.google.appengine.tools.remoteapi.RemoteApiOptions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
-import google.registry.beam.initsql.BeamJpaModule;
+import google.registry.backup.AppEngineEnvironment;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -68,9 +68,8 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
 
   @Parameter(
       names = {"--sql_access_info"},
-      description =
+      description = "Name of a file containing space-separated SQL access info used when deploying "
-          "Name of a file containing space-separated SQL access info used when deploying "
+          + "Beam pipelines")
-              + "Beam pipelines")
   private String sqlAccessInfoFile = null;
 
   // Do not make this final - compile-time constant inlining may interfere with JCommander.
@@ -168,7 +167,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
     component =
         DaggerRegistryToolComponent.builder()
             .credentialFilePath(credentialJson)
-            .beamJpaModule(new BeamJpaModule(sqlAccessInfoFile))
+            .sqlAccessInfoFile(sqlAccessInfoFile)
             .build();
 
     // JCommander stores sub-commands as nested JCommander objects containing a list of user objects
@@ -179,7 +178,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
             Iterables.getOnlyElement(jcommander.getCommands().get(parsedCommand).getObjects());
     loggingParams.configureLogging();  // Must be called after parameters are parsed.
 
-    try {
+    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
       runCommand(command);
     } catch (RuntimeException ex) {
       if (Throwables.getRootCause(ex) instanceof LoginRequiredException) {

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

@@ -134,6 +134,9 @@ interface RegistryToolComponent {
     @BindsInstance
     Builder credentialFilePath(@Nullable @Config("credentialFilePath") String credentialFilePath);
 
+    @BindsInstance
+    Builder sqlAccessInfoFile(@Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile);
+
     Builder beamJpaModule(BeamJpaModule beamJpaModule);
 
     RegistryToolComponent build();

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

@@ -54,7 +54,7 @@ public final class BeamJpaExtension implements BeforeEachCallback, AfterEachCall
     if (beamJpaModule != null) {
       return beamJpaModule;
     }
-    return beamJpaModule = new BeamJpaModule(credentialFile.getAbsolutePath());
+    return beamJpaModule = new BeamJpaModule(credentialFile.getAbsolutePath(), null);
   }
 
   @Override

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

@@ -82,7 +82,8 @@ class BeamJpaModuleTest {
         DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
             .beamJpaModule(
                 new BeamJpaModule(
-                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0)))
+                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0),
+                    String.format("domain-registry-%s", environmentName)))
             .build()
             .cloudSqlJpaTransactionManager();
     assertThat(