Allow BEAM pipeline to choose JDBC isolation levels (#916)

* Allow BEAM pipeline to choose JDBC isolation levels

Some BEAM pipelines may only perform READ-ONLY (e.g., reporting) or
blind-write (datastore to sql data migration) operations, which do not
need the default TRANSACTION_SERIALIZABLE isolation level. In such
cases, a less strict level allows better performance.

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

@@ -30,6 +30,7 @@ import google.registry.keyring.kms.KmsModule;
 import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.JdbcJpaTm;
 import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
+import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.privileges.secretmanager.SecretManagerModule;
 import google.registry.util.UtilsModule;
@@ -57,6 +58,7 @@ public class BeamJpaModule {
 
   @Nullable private final String sqlAccessInfoFile;
   @Nullable private final String cloudKmsProjectId;
+  @Nullable private final TransactionIsolationLevel isolationOverride;
 
   /**
    * Constructs a new instance of {@link BeamJpaModule}.
@@ -73,10 +75,20 @@ public class BeamJpaModule {
    *     real encrypted file on GCS as returned by {@link
    *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
    *     with credentials to a test database.
+   * @param cloudKmsProjectId the GCP project where the credential decryption key can be found
+   * @param isolationOverride the desired Transaction Isolation level for all JDBC connections
    */
-  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
+  public BeamJpaModule(
+      @Nullable String sqlAccessInfoFile,
+      @Nullable String cloudKmsProjectId,
+      @Nullable TransactionIsolationLevel isolationOverride) {
     this.sqlAccessInfoFile = sqlAccessInfoFile;
     this.cloudKmsProjectId = cloudKmsProjectId;
+    this.isolationOverride = isolationOverride;
+  }
+
+  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
+    this(sqlAccessInfoFile, cloudKmsProjectId, null);
   }
 
   /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
@@ -154,6 +166,13 @@ public class BeamJpaModule {
     return "nomulus-tool-keyring";
   }
 
+  @Provides
+  @Config("beamIsolationOverride")
+  @Nullable
+  TransactionIsolationLevel providesIsolationOverride() {
+    return isolationOverride;
+  }
+
   @Provides
   @Config("beamHibernateHikariMaximumPoolSize")
   static int getBeamHibernateHikariMaximumPoolSize() {

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

@@ -27,6 +27,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
+import dagger.BindsOptionalOf;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
@@ -42,8 +43,12 @@ import google.registry.privileges.secretmanager.SqlUser.RobotUser;
 import google.registry.tools.AuthModule.CloudSqlClientCredential;
 import google.registry.util.Clock;
 import java.lang.annotation.Documented;
+import java.sql.Connection;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import javax.inject.Provider;
 import javax.inject.Qualifier;
 import javax.inject.Singleton;
 import javax.persistence.EntityManagerFactory;
@@ -52,7 +57,7 @@ import org.hibernate.cfg.Environment;
 
 /** Dagger module class for the persistence layer. */
 @Module
-public class PersistenceModule {
+public abstract class PersistenceModule {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   // This name must be the same as the one defined in persistence.xml.
@@ -102,25 +107,48 @@ public class PersistenceModule {
       @Config("cloudSqlJdbcUrl") String jdbcUrl,
       @Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
       @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
-    return createPartialSqlConfigs(jdbcUrl, instanceConnectionName, defaultConfigs);
+    return createPartialSqlConfigs(
+        jdbcUrl, instanceConnectionName, defaultConfigs, Optional.empty());
   }
 
+  /**
+   * Optionally overrides the isolation level in the config file.
+   *
+   * <p>The binding for {@link TransactionIsolationLevel} may be {@link Nullable}. As a result, it
+   * is a compile-time error to inject {@code Optional<TransactionIsolation>} (See {@link
+   * BindsOptionalOf} for more information). User should inject {@code
+   * Optional<Provider<TransactionIsolation>>} instead.
+   */
+  @BindsOptionalOf
+  @Config("beamIsolationOverride")
+  abstract TransactionIsolationLevel bindBeamIsolationOverride();
+
   @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);
+      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
+      @Config("beamIsolationOverride")
+          Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
+    return createPartialSqlConfigs(
+        jdbcUrl, instanceConnectionName, defaultConfigs, isolationOverride);
   }
 
-  private static ImmutableMap<String, String> createPartialSqlConfigs(
-      String jdbcUrl, String instanceConnectionName, ImmutableMap<String, String> defaultConfigs) {
+  @VisibleForTesting
+  static ImmutableMap<String, String> createPartialSqlConfigs(
+      String jdbcUrl,
+      String instanceConnectionName,
+      ImmutableMap<String, String> defaultConfigs,
+      Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
     HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
     overrides.put(Environment.URL, jdbcUrl);
     overrides.put(HIKARI_DS_SOCKET_FACTORY, "com.google.cloud.sql.postgres.SocketFactory");
     overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, instanceConnectionName);
+    isolationOverride
+        .map(Provider::get)
+        .ifPresent(override -> overrides.put(Environment.ISOLATION, override.name()));
     return ImmutableMap.copyOf(overrides);
   }
 
@@ -254,6 +282,37 @@ public class PersistenceModule {
     }
   }
 
+  /**
+   * Transaction isolation levels supported by Cloud SQL (mysql and postgresql).
+   *
+   * <p>Enum names may be used for property-based configuration, and must match the corresponding
+   * variable names in {@link Connection}.
+   */
+  public enum TransactionIsolationLevel {
+    TRANSACTION_READ_UNCOMMITTED,
+    TRANSACTION_READ_COMMITTED,
+    TRANSACTION_REPEATABLE_READ,
+    TRANSACTION_SERIALIZABLE;
+
+    private final int value;
+
+    TransactionIsolationLevel() {
+      try {
+        // name() is final in parent class (Enum.java), therefore safe to call in constructor.
+        value = Connection.class.getField(name()).getInt(null);
+      } catch (Exception e) {
+        throw new IllegalStateException(
+            String.format(
+                "%s Enum name %s has no matching public field in java.sql.Connection.",
+                getClass().getSimpleName(), name()));
+      }
+    }
+
+    public final int getValue() {
+      return value;
+    }
+  }
+
   /** Dagger qualifier for {@link JpaTransactionManager} used for App Engine application. */
   @Qualifier
   @Documented

core/src/test/java/google/registry/persistence/PersistenceModuleTest.java

@@ -15,10 +15,24 @@
 package google.registry.persistence;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth8.assertThat;
 
+import dagger.Component;
+import google.registry.beam.initsql.BeamJpaModule;
+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.TransactionIsolationLevel;
+import google.registry.privileges.secretmanager.SecretManagerModule;
 import google.registry.testing.DatastoreEntityExtension;
+import google.registry.util.UtilsModule;
+import java.util.Optional;
+import javax.inject.Provider;
+import javax.inject.Singleton;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
+import org.hibernate.cfg.Environment;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -29,7 +43,7 @@ import org.testcontainers.junit.jupiter.Testcontainers;
 
 /** Unit tests for {@link PersistenceModule}. */
 @Testcontainers
-public class PersistenceModuleTest {
+class PersistenceModuleTest {
 
   @Container
   private final PostgreSQLContainer database =
@@ -64,4 +78,61 @@ public class PersistenceModuleTest {
     assertThat(em.isOpen()).isTrue();
     em.close();
   }
+
+  @Test
+  void appengineIsolation() {
+    assertThat(PersistenceModule.provideDefaultDatabaseConfigs().get(Environment.ISOLATION))
+        .isEqualTo(TransactionIsolationLevel.TRANSACTION_SERIALIZABLE.name());
+  }
+
+  @Test
+  void beamIsolation_default() {
+    Optional<Provider<TransactionIsolationLevel>> injected =
+        DaggerPersistenceModuleTest_BeamConfigTestComponent.builder()
+            .beamJpaModule(new BeamJpaModule(null, null))
+            .build()
+            .getIsolationOverride();
+    assertThat(injected).isNotNull();
+    assertThat(injected.get().get()).isNull();
+    assertThat(
+            PersistenceModule.provideBeamPipelineCloudSqlConfigs(
+                    "", "", PersistenceModule.provideDefaultDatabaseConfigs(), injected)
+                .get(Environment.ISOLATION))
+        .isEqualTo(TransactionIsolationLevel.TRANSACTION_SERIALIZABLE.name());
+  }
+
+  @Test
+  void beamIsolation_override() {
+    Optional<Provider<TransactionIsolationLevel>> injected =
+        DaggerPersistenceModuleTest_BeamConfigTestComponent.builder()
+            .beamJpaModule(
+                new BeamJpaModule(
+                    null, null, TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED))
+            .build()
+            .getIsolationOverride();
+    assertThat(injected).isNotNull();
+    assertThat(injected.get().get())
+        .isEqualTo(TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED);
+    assertThat(
+            PersistenceModule.provideBeamPipelineCloudSqlConfigs(
+                    "", "", PersistenceModule.provideDefaultDatabaseConfigs(), injected)
+                .get(Environment.ISOLATION))
+        .isEqualTo(TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED.name());
+  }
+
+  @Singleton
+  @Component(
+      modules = {
+        BeamJpaModule.class,
+        ConfigModule.class,
+        CredentialModule.class,
+        KmsModule.class,
+        PersistenceModule.class,
+        SecretManagerModule.class,
+        UtilsModule.class
+      })
+  public interface BeamConfigTestComponent {
+    @Config("beamIsolationOverride")
+    Optional<Provider<TransactionIsolationLevel>> getIsolationOverride();
+  }
 }