Enable JpaTransactionManager in all environment (#352)

* Enable JpaTransactionManager in all environment

* Refactor to have a single place to create JpaTm

core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java

@@ -14,16 +14,13 @@
 
 package google.registry.model.transaction;
 
-import static google.registry.config.RegistryEnvironment.ALPHA;
-import static google.registry.config.RegistryEnvironment.CRASH;
-import static google.registry.config.RegistryEnvironment.SANDBOX;
 
 import com.google.appengine.api.utils.SystemProperty;
 import com.google.appengine.api.utils.SystemProperty.Environment.Value;
 import com.google.common.annotations.VisibleForTesting;
-import google.registry.config.RegistryEnvironment;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.DaggerPersistenceComponent;
+import google.registry.tools.RegistryToolEnvironment;
 
 /** Factory class to create {@link TransactionManager} instance. */
 // TODO: Rename this to PersistenceFactory and move to persistence package.
@@ -35,8 +32,11 @@ public class TransactionManagerFactory {
   private TransactionManagerFactory() {}
 
   private static JpaTransactionManager createJpaTransactionManager() {
-    if (shouldEnableJpaTm() && isInAppEngine()) {
+    if (isInAppEngine()) {
       return DaggerPersistenceComponent.create().appEngineJpaTransactionManager();
+    } else if (RegistryToolEnvironment.isInRegistryTool()
+        && RegistryToolEnvironment.isJpaTmEnabled()) {
+      return DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
     } else {
       return DummyJpaTransactionManager.create();
     }
@@ -49,23 +49,6 @@ public class TransactionManagerFactory {
     return new DatastoreTransactionManager(null);
   }
 
-  /**
-   * Sets jpaTm to the implementation for Nomulus tool. Note that this method should be only used by
-   * {@link google.registry.tools.RegistryCli} to initialize jpaTm.
-   */
-  public static void initForTool() {
-    if (shouldEnableJpaTm()) {
-      jpaTm = DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
-    }
-  }
-
-  // TODO(shicong): Enable JpaTm for all environments and remove this function
-  private static boolean shouldEnableJpaTm() {
-    return RegistryEnvironment.get() == ALPHA
-        || RegistryEnvironment.get() == CRASH
-        || RegistryEnvironment.get() == SANDBOX;
-  }
-
   /**
    * This function uses App Engine API to determine if the current runtime environment is App
    * Engine.

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

@@ -31,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
-import google.registry.model.transaction.TransactionManagerFactory;
 import google.registry.tools.AuthModule.LoginRequiredException;
 import google.registry.tools.params.ParameterFactory;
 import java.io.ByteArrayInputStream;
@@ -237,7 +236,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
     }
 
     if (command instanceof CommandWithCloudSql) {
-      TransactionManagerFactory.initForTool();
+      RegistryToolEnvironment.get().enableJpaTm();
     }
 
     command.run();

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

@@ -25,7 +25,7 @@ import google.registry.config.RegistryEnvironment;
 import google.registry.config.SystemPropertySetter;
 
 /** Enum of production environments, used for the {@code --environment} flag. */
-enum RegistryToolEnvironment {
+public enum RegistryToolEnvironment {
   PRODUCTION(RegistryEnvironment.PRODUCTION),
   ALPHA(RegistryEnvironment.ALPHA),
   CRASH(RegistryEnvironment.CRASH),
@@ -39,6 +39,7 @@ enum RegistryToolEnvironment {
 
   private static final ImmutableList<String> FLAGS = ImmutableList.of("-e", "--environment");
   private static RegistryToolEnvironment instance;
+  private static boolean isJpaTmEnabled = false;
   private final RegistryEnvironment actualEnvironment;
   private final ImmutableMap<String, String> extraProperties;
 
@@ -72,7 +73,7 @@ enum RegistryToolEnvironment {
    *
    * <p>This should be called after {@link #parseFromArgs(String[])}.
    */
-  static RegistryToolEnvironment get() {
+  public static RegistryToolEnvironment get() {
     checkState(instance != null, "No RegistryToolEnvironment has been set up");
     return instance;
   }
@@ -98,6 +99,26 @@ enum RegistryToolEnvironment {
     }
   }
 
+  /** Returns true if the RegistryToolEnvironment is set up. */
+  public static boolean isInRegistryTool() {
+    return instance != null;
+  }
+
+  /**
+   * Sets the flag to indicate that the running command needs JpaTransactionManager to be enabled.
+   */
+  public static void enableJpaTm() {
+    isJpaTmEnabled = true;
+  }
+
+  /**
+   * Returns true if the JpaTransactionManager is enabled. Note that JpaTm is actually enabled in
+   * {@link google.registry.model.transaction.TransactionManagerFactory} by reading this flag.
+   */
+  public static boolean isJpaTmEnabled() {
+    return isJpaTmEnabled;
+  }
+
   /** Extracts value from command-line arguments associated with any {@code flags}. */
   private static String getFlagValue(String[] args, Iterable<String> flags) {
     for (String flag : flags) {