Disambiguate injected Cloud SQL parameter names (#657)

* Disambiguate injected Cloud SQL parameter names

This allows us to also inject the BeamJpaModule into RegistryTool, which
allows us to use the SocketJpaTransactionManager in Beam pipelines.

Some side effects of this include:
- duplication of KMS connections -- one standard, one Beam
- duplication of the creation of the partial Hibernate SQL configs
- removal of ambiguity between credentialFileName, credentialFilename,
and credentialFilePath -- we now use the latter.
- Performing the credential null check when instantiating the SQL
connection rather than when instantiating the module object. See the code
comments for more details on this.

I verified that this compiles and the tests run successfully when
injecting a @SocketFactoryJpaTm into a Beam pipeline.

* Remove two unnecessary config points and change the name of two params

* Use @Config instead of @Named and change the pool size

* Replace non-visible link with code

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

@@ -19,15 +19,14 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.isNullOrEmpty;
 
 import com.google.common.base.Splitter;
-import com.google.common.collect.ImmutableList;
 import dagger.Binds;
 import dagger.Component;
 import dagger.Lazy;
 import dagger.Module;
 import dagger.Provides;
-import google.registry.beam.initsql.BeamJpaModule.BindModule;
 import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.keyring.kms.KmsModule;
 import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.JdbcJpaTm;
@@ -37,13 +36,14 @@ import google.registry.util.Clock;
 import google.registry.util.Sleeper;
 import google.registry.util.SystemClock;
 import google.registry.util.SystemSleeper;
+import google.registry.util.UtilsModule;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.nio.channels.Channels;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
-import javax.inject.Named;
+import javax.annotation.Nullable;
 import javax.inject.Singleton;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
@@ -53,27 +53,31 @@ import org.apache.beam.sdk.io.fs.ResourceId;
  *
  * <p>This module is intended for use in BEAM pipelines, and uses a BEAM utility to access GCS like
  * a regular file system.
- *
- * <p>Note that {@link google.registry.config.RegistryConfig.ConfigModule} cannot be used here,
- * since many bindings, especially KMS-related ones, are different.
  */
-@Module(includes = {BindModule.class})
-class BeamJpaModule {
+@Module
+public class BeamJpaModule {
 
   private static final String GCS_SCHEME = "gs://";
 
-  private final String credentialFilePath;
+  @Nullable private final String credentialFilePath;
 
   /**
    * 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.
+   * 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
    *     real encrypted file on GCS as returned by {@link
    *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
    *     with credentials to a test database.
    */
-  BeamJpaModule(String credentialFilePath) {
-    checkArgument(!isNullOrEmpty(credentialFilePath), "Null or empty credentialFilePath");
+  public BeamJpaModule(@Nullable String credentialFilePath) {
     this.credentialFilePath = credentialFilePath;
   }
 
@@ -85,6 +89,7 @@ class BeamJpaModule {
   @Provides
   @Singleton
   SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
+    checkArgument(!isNullOrEmpty(credentialFilePath), "Null or empty credentialFilePath");
     String line = readOnlyLineFromCredentialFile();
     if (isCloudSqlCredential()) {
       line = lazyDecryptor.get().decrypt(line);
@@ -114,13 +119,13 @@ class BeamJpaModule {
   }
 
   @Provides
-  @Config("cloudSqlJdbcUrl")
+  @Config("beamCloudSqlJdbcUrl")
   String provideJdbcUrl(SqlAccessInfo sqlAccessInfo) {
     return sqlAccessInfo.jdbcUrl();
   }
 
   @Provides
-  @Config("cloudSqlInstanceConnectionName")
+  @Config("beamCloudSqlInstanceConnectionName")
   String provideSqlInstanceName(SqlAccessInfo sqlAccessInfo) {
     return sqlAccessInfo
         .cloudSqlInstanceName()
@@ -128,39 +133,33 @@ class BeamJpaModule {
   }
 
   @Provides
-  @Config("cloudSqlUsername")
+  @Config("beamCloudSqlUsername")
   String provideSqlUsername(SqlAccessInfo sqlAccessInfo) {
     return sqlAccessInfo.user();
   }
 
   @Provides
-  @Config("cloudSqlPassword")
+  @Config("beamCloudSqlPassword")
   String provideSqlPassword(SqlAccessInfo sqlAccessInfo) {
     return sqlAccessInfo.password();
   }
 
   @Provides
-  @Config("cloudKmsProjectId")
+  @Config("beamCloudKmsProjectId")
   static String kmsProjectId() {
     return "domain-registry-dev";
   }
 
   @Provides
-  @Config("cloudKmsKeyRing")
+  @Config("beamCloudKmsKeyRing")
   static String keyRingName() {
     return "nomulus-tool-keyring";
   }
 
   @Provides
-  @Config("defaultCredentialOauthScopes")
-  static ImmutableList<String> defaultCredentialOauthScopes() {
-    return ImmutableList.of("https://www.googleapis.com/auth/cloud-platform");
-  }
-
-  @Provides
-  @Named("transientFailureRetries")
-  static int transientFailureRetries() {
-    return 12;
+  @Config("beamHibernateHikariMaximumPoolSize")
+  static int getBeamHibernateHikariMaximumPoolSize() {
+    return 4;
   }
 
   @Module
@@ -176,10 +175,12 @@ class BeamJpaModule {
   @Singleton
   @Component(
       modules = {
+        ConfigModule.class,
         CredentialModule.class,
         BeamJpaModule.class,
         KmsModule.class,
-        PersistenceModule.class
+        PersistenceModule.class,
+        UtilsModule.class
       })
   public interface JpaTransactionManagerComponent {
     @SocketFactoryJpaTm

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

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 
 import com.google.api.services.cloudkms.v1.model.DecryptRequest;
 import com.google.common.base.Strings;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.kms.KmsConnection;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
@@ -34,7 +35,7 @@ public class CloudSqlCredentialDecryptor {
   private final KmsConnection kmsConnection;
 
   @Inject
-  CloudSqlCredentialDecryptor(KmsConnection kmsConnection) {
+  CloudSqlCredentialDecryptor(@Config("beamKmsConnection") KmsConnection kmsConnection) {
     this.kmsConnection = kmsConnection;
   }
 

core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java

@@ -94,8 +94,7 @@ public class Spec11Pipeline implements Serializable {
       @Config("spec11TemplateUrl") String spec11TemplateUrl,
       @Config("reportingBucketUrl") String reportingBucketUrl,
       @LocalCredential GoogleCredentialsBundle googleCredentialsBundle,
-      Retrier retrier
-  ) {
+      Retrier retrier) {
     this.projectId = projectId;
     this.beamStagingUrl = beamStagingUrl;
     this.spec11TemplateUrl = spec11TemplateUrl;

core/src/main/java/google/registry/keyring/kms/KmsConnectionImpl.java

@@ -25,11 +25,9 @@ import com.google.api.services.cloudkms.v1.model.DecryptRequest;
 import com.google.api.services.cloudkms.v1.model.EncryptRequest;
 import com.google.api.services.cloudkms.v1.model.KeyRing;
 import com.google.api.services.cloudkms.v1.model.UpdateCryptoKeyPrimaryVersionRequest;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeyringException;
 import google.registry.util.Retrier;
 import java.io.IOException;
-import javax.inject.Inject;
 
 /** The {@link KmsConnection} which talks to Cloud KMS. */
 class KmsConnectionImpl implements KmsConnection {
@@ -44,12 +42,7 @@ class KmsConnectionImpl implements KmsConnection {
   private final String projectId;
   private final Retrier retrier;
 
-  @Inject
-  KmsConnectionImpl(
-      @Config("cloudKmsProjectId") String projectId,
-      @Config("cloudKmsKeyRing") String kmsKeyRingName,
-      Retrier retrier,
-      CloudKMS kms) {
+  KmsConnectionImpl(String projectId, String kmsKeyRingName, Retrier retrier, CloudKMS kms) {
     this.projectId = projectId;
     this.kmsKeyRingName = kmsKeyRingName;
     this.retrier = retrier;

core/src/main/java/google/registry/keyring/kms/KmsKeyring.java

@@ -21,6 +21,7 @@ import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
 import com.googlecode.objectify.Key;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
@@ -86,7 +87,7 @@ public class KmsKeyring implements Keyring {
   private final KmsConnection kmsConnection;
 
   @Inject
-  KmsKeyring(KmsConnection kmsConnection) {
+  KmsKeyring(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
     this.kmsConnection = kmsConnection;
   }
 

core/src/main/java/google/registry/keyring/kms/KmsModule.java

@@ -24,6 +24,7 @@ import google.registry.config.CredentialModule.DefaultCredential;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.Keyring;
 import google.registry.util.GoogleCredentialsBundle;
+import google.registry.util.Retrier;
 
 /** Dagger module for Cloud KMS. */
 @Module
@@ -32,9 +33,22 @@ public abstract class KmsModule {
   public static final String NAME = "KMS";
 
   @Provides
+  @Config("defaultKms")
   static CloudKMS provideKms(
       @DefaultCredential GoogleCredentialsBundle credentialsBundle,
       @Config("cloudKmsProjectId") String projectId) {
+    return createKms(credentialsBundle, projectId);
+  }
+
+  @Provides
+  @Config("beamKms")
+  static CloudKMS provideBeamKms(
+      @DefaultCredential GoogleCredentialsBundle credentialsBundle,
+      @Config("beamCloudKmsProjectId") String projectId) {
+    return createKms(credentialsBundle, projectId);
+  }
+
+  private static CloudKMS createKms(GoogleCredentialsBundle credentialsBundle, String projectId) {
     return new CloudKMS.Builder(
             credentialsBundle.getHttpTransport(),
             credentialsBundle.getJsonFactory(),
@@ -43,11 +57,28 @@ public abstract class KmsModule {
         .build();
   }
 
+  @Provides
+  @Config("defaultKmsConnection")
+  static KmsConnection provideKmsConnection(
+      @Config("cloudKmsProjectId") String projectId,
+      @Config("cloudKmsKeyRing") String keyringName,
+      Retrier retrier,
+      @Config("defaultKms") CloudKMS defaultKms) {
+    return new KmsConnectionImpl(projectId, keyringName, retrier, defaultKms);
+  }
+
+  @Provides
+  @Config("beamKmsConnection")
+  static KmsConnection provideBeamKmsConnection(
+      @Config("beamCloudKmsProjectId") String projectId,
+      @Config("beamCloudKmsKeyRing") String keyringName,
+      Retrier retrier,
+      @Config("beamKms") CloudKMS defaultKms) {
+    return new KmsConnectionImpl(projectId, keyringName, retrier, defaultKms);
+  }
+
   @Binds
   @IntoMap
   @StringKey(NAME)
   abstract Keyring provideKeyring(KmsKeyring keyring);
-
-  @Binds
-  abstract KmsConnection provideKmsConnection(KmsConnectionImpl kmsConnectionImpl);
 }

core/src/main/java/google/registry/keyring/kms/KmsUpdater.java

@@ -39,6 +39,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.collect.ImmutableMap;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.kms.KmsKeyring.PrivateKeyLabel;
 import google.registry.keyring.kms.KmsKeyring.PublicKeyLabel;
@@ -64,7 +65,7 @@ public final class KmsUpdater {
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(KmsConnection kmsConnection) {
+  public KmsUpdater(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
     this.kmsConnection = kmsConnection;
 
     // Use LinkedHashMap to preserve insertion order on update() to simplify testing and debugging

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

@@ -94,6 +94,21 @@ public class PersistenceModule {
       @Config("cloudSqlJdbcUrl") String jdbcUrl,
       @Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
       @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
+    return createPartialSqlConfigs(jdbcUrl, instanceConnectionName, defaultConfigs);
+  }
+
+  @Provides
+  @Singleton
+  @BeamPipelineCloudSqlConfigs
+  static ImmutableMap<String, String> provideBeamPipelineCloudSqlConfigs(
+      @Config("beamCloudSqlJdbcUrl") String jdbcUrl,
+      @Config("beamCloudSqlInstanceConnectionName") String instanceConnectionName,
+      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
+    return createPartialSqlConfigs(jdbcUrl, instanceConnectionName, defaultConfigs);
+  }
+
+  private static ImmutableMap<String, String> createPartialSqlConfigs(
+      String jdbcUrl, String instanceConnectionName, ImmutableMap<String, String> defaultConfigs) {
     HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
     overrides.put(Environment.URL, jdbcUrl);
     overrides.put(HIKARI_DS_SOCKET_FACTORY, "com.google.cloud.sql.postgres.SocketFactory");
@@ -135,13 +150,15 @@ public class PersistenceModule {
   @Singleton
   @SocketFactoryJpaTm
   static JpaTransactionManager provideSocketFactoryJpaTm(
-      @Config("cloudSqlUsername") String username,
-      @Config("cloudSqlPassword") String password,
-      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
+      @Config("beamCloudSqlUsername") String username,
+      @Config("beamCloudSqlPassword") String password,
+      @Config("beamHibernateHikariMaximumPoolSize") int hikariMaximumPoolSize,
+      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
       Clock clock) {
     HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, password);
+    overrides.put(HIKARI_MAXIMUM_POOL_SIZE, String.valueOf(hikariMaximumPoolSize));
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
@@ -149,9 +166,9 @@ public class PersistenceModule {
   @Singleton
   @JdbcJpaTm
   static JpaTransactionManager provideLocalJpaTm(
-      @Config("cloudSqlJdbcUrl") String jdbcUrl,
-      @Config("cloudSqlUsername") String username,
-      @Config("cloudSqlPassword") String password,
+      @Config("beamCloudSqlJdbcUrl") String jdbcUrl,
+      @Config("beamCloudSqlUsername") String username,
+      @Config("beamCloudSqlPassword") String password,
       @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
       Clock clock) {
     HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
@@ -218,6 +235,11 @@ public class PersistenceModule {
   @Documented
   @interface PartialCloudSqlConfigs {}
 
+  /** Dagger qualifier for the Cloud SQL configs used by Beam pipelines. */
+  @Qualifier
+  @Documented
+  @interface BeamPipelineCloudSqlConfigs {}
+
   /** Dagger qualifier for the default Hibernate configurations. */
   // TODO(shicong): Change annotations in this class to none public or put them in a top level
   //  package

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

@@ -50,7 +50,6 @@ import java.lang.annotation.RetentionPolicy;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import javax.annotation.Nullable;
-import javax.inject.Named;
 import javax.inject.Qualifier;
 import javax.inject.Singleton;
 
@@ -151,10 +150,10 @@ public class AuthModule {
   public static String provideLocalCredentialJson(
       Lazy<GoogleClientSecrets> clientSecrets,
       @StoredCredential Lazy<Credential> credential,
-      @Nullable @Named("credentialFileName") String credentialFilename) {
+      @Nullable @Config("credentialFilePath") String credentialFilePath) {
     try {
-      if (credentialFilename != null) {
-        return new String(Files.readAllBytes(Paths.get(credentialFilename)), UTF_8);
+      if (credentialFilePath != null) {
+        return new String(Files.readAllBytes(Paths.get(credentialFilePath)), UTF_8);
       } else {
         return new Gson()
             .toJson(

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

@@ -29,6 +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.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -153,11 +154,15 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
       return;
     }
 
-    checkState(RegistryToolEnvironment.get() == environment,
+    checkState(
+        RegistryToolEnvironment.get() == environment,
         "RegistryToolEnvironment argument pre-processing kludge failed.");
 
     component =
-        DaggerRegistryToolComponent.builder().credentialFilename(credentialJson).build();
+        DaggerRegistryToolComponent.builder()
+            .credentialFilePath(credentialJson)
+            .beamJpaModule(new BeamJpaModule(credentialJson))
+            .build();
 
     // JCommander stores sub-commands as nested JCommander objects containing a list of user objects
     // to be populated.  Extract the subcommand by getting the JCommander wrapper and then

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

@@ -18,8 +18,10 @@ import dagger.BindsInstance;
 import dagger.Component;
 import dagger.Lazy;
 import google.registry.batch.BatchModule;
+import google.registry.beam.initsql.BeamJpaModule;
 import google.registry.bigquery.BigqueryModule;
 import google.registry.config.CredentialModule.LocalCredentialJson;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.dns.writer.VoidDnsWriterModule;
 import google.registry.dns.writer.clouddns.CloudDnsWriterModule;
@@ -42,7 +44,6 @@ import google.registry.tools.AuthModule.LocalCredentialModule;
 import google.registry.util.UtilsModule;
 import google.registry.whois.WhoisModule;
 import javax.annotation.Nullable;
-import javax.inject.Named;
 import javax.inject.Singleton;
 
 /**
@@ -57,6 +58,7 @@ import javax.inject.Singleton;
       AppEngineAdminApiModule.class,
       AuthModule.class,
       BatchModule.class,
+      BeamJpaModule.class,
       BigqueryModule.class,
       ConfigModule.class,
       CloudDnsWriterModule.class,
@@ -130,7 +132,9 @@ interface RegistryToolComponent {
   @Component.Builder
   interface Builder {
     @BindsInstance
-    Builder credentialFilename(@Nullable @Named("credentialFileName") String credentialFilename);
+    Builder credentialFilePath(@Nullable @Config("credentialFilePath") String credentialFilePath);
+
+    Builder beamJpaModule(BeamJpaModule beamJpaModule);
 
     RegistryToolComponent build();
   }