Add a transactional parameter to @action

This also removes inner transactions on some Registrar methods with some resultant refactoring.
google · Nov 13, 2023 · 61eb640 · 61eb640
1 parent 9e3c589
commit 61eb640
Show file tree

Hide file tree

Showing 15 changed files with 264 additions and 117 deletions.
diff --git a/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java b/core/src/main/java/google/registry/batch/DeleteLoadTestDataAction.java
@@ -89,12 +89,9 @@ public void run() {
         !RegistryEnvironment.get().equals(PRODUCTION),
         "This action is not safe to run on PRODUCTION.");
 
-    tm().transact(
-            () -> {
-              LOAD_TEST_REGISTRARS.forEach(this::deletePollMessages);
-              tm().loadAllOfStream(Contact.class).forEach(this::deleteContact);
-              tm().loadAllOfStream(Host.class).forEach(this::deleteHost);
-            });
+    LOAD_TEST_REGISTRARS.forEach(this::deletePollMessages);
+    tm().loadAllOfStream(Contact.class).forEach(this::deleteContact);
+    tm().loadAllOfStream(Host.class).forEach(this::deleteHost);
   }
 
   private void deletePollMessages(String registrarId) {

diff --git a/core/src/main/java/google/registry/model/registrar/Registrar.java b/core/src/main/java/google/registry/model/registrar/Registrar.java
@@ -585,7 +585,7 @@ public Optional<RegistrarPoc> getWhoisAbuseContact() {
   }
 
   private ImmutableSet<RegistrarPoc> getContactPocs() {
-    return tm().transact(
+    return tm().reTransact(
             () ->
                 tm().query("FROM RegistrarPoc WHERE registrarId = :registrarId", RegistrarPoc.class)
                     .setParameter("registrarId", registrarId)
@@ -944,7 +944,7 @@ public static String checkValidEmail(String email) {
 
   /** Loads all registrar entities directly from the database. */
   public static Iterable<Registrar> loadAll() {
-    return tm().transact(() -> tm().loadAllOf(Registrar.class));
+    return tm().loadAllOf(Registrar.class);
   }
 
   /** Loads all registrar entities using an in-memory cache. */
@@ -962,7 +962,7 @@ public static ImmutableSet<VKey<Registrar>> loadAllKeysCached() {
   /** Loads and returns a registrar entity by its id directly from the database. */
   public static Optional<Registrar> loadByRegistrarId(String registrarId) {
     checkArgument(!Strings.isNullOrEmpty(registrarId), "registrarId must be specified");
-    return tm().transact(() -> tm().loadByKeyIfPresent(createVKey(registrarId)));
+    return tm().reTransact(() -> tm().loadByKeyIfPresent(createVKey(registrarId)));
   }
 
   /**

diff --git a/core/src/main/java/google/registry/rdap/UpdateRegistrarRdapBaseUrlsAction.java b/core/src/main/java/google/registry/rdap/UpdateRegistrarRdapBaseUrlsAction.java
@@ -69,7 +69,7 @@ public final class UpdateRegistrarRdapBaseUrlsAction implements Runnable {
   public void run() {
     try {
       ImmutableMap<String, String> ianaIdsToUrls = getIanaIdsToUrls();
-      tm().transact(() -> processAllRegistrars(ianaIdsToUrls));
+      processAllRegistrars(ianaIdsToUrls);
     } catch (Exception e) {
       throw new InternalServerErrorException("Error when retrieving RDAP base URL CSV file", e);
     }

diff --git a/core/src/main/java/google/registry/request/Action.java b/core/src/main/java/google/registry/request/Action.java
@@ -59,6 +59,9 @@ public String getServiceId() {
   /** HTTP methods that request processor should allow. */
   Method[] method() default Method.GET;
 
+  /** Whether to create a wrapping DB transaction around the execution of the action. */
+  boolean transactional() default true;
+
   /**
    * Indicates request processor should print "OK" to the HTTP client on success.
    *

diff --git a/core/src/main/java/google/registry/request/RequestHandler.java b/core/src/main/java/google/registry/request/RequestHandler.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_METHOD_NOT_ALLOWED;
@@ -155,7 +156,12 @@ public void handleRequest(HttpServletRequest req, HttpServletResponse rsp) throw
     boolean success = true;
     DateTime startTime = clock.nowUtc();
     try {
-      route.get().instantiator().apply(component).run();
+      Runnable action = route.get().instantiator().apply(component);
+      if (route.get().action().transactional()) {
+        tm().transact(action::run);
+      } else {
+        action.run();
+      }
       if (route.get().action().automaticallyPrintOk()) {
         rsp.setContentType(PLAIN_TEXT_UTF_8.toString());
         rsp.getWriter().write("OK\n");

diff --git a/core/src/main/java/google/registry/tools/ConfirmingCommand.java b/core/src/main/java/google/registry/tools/ConfirmingCommand.java
@@ -14,8 +14,10 @@
 
 package google.registry.tools;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.tools.CommandUtilities.promptForYes;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.commons.lang3.exception.ExceptionUtils.rethrow;
 
 import com.beust.jcommander.Parameter;
 import com.google.common.base.Strings;
@@ -44,20 +46,46 @@ protected ConfirmingCommand() {
 
   @Override
   public final void run() throws Exception {
-    if (checkExecutionState()) {
-      init();
-      printLineIfNotEmpty(prompt(), printStream);
-      if (dontRunCommand()) {
-        // This typically happens when all of the work is accomplished inside of prompt(), so do
-        // nothing further.
-        return;
-      } else if (force || promptForYes("Perform this command?")) {
-        printStream.println("Running ... ");
-        printStream.println(execute());
-        printLineIfNotEmpty(postExecute(), printStream);
-      } else {
-        printStream.println("Command aborted.");
-      }
+    if (!checkExecutionState()) {
+      return;
+    }
+
+    tm().transact(
+            () -> {
+              try {
+                init();
+                printLineIfNotEmpty(prompt(), printStream);
+              } catch (Exception e) {
+                rethrow(e);
+              }
+            });
+
+    if (dontRunCommand()) {
+      // This typically happens when all of the work is accomplished inside of
+      // prompt(), so do
+      // nothing further.
+      return;
+    } else if (force || promptForYes("Perform this command?")) {
+      printStream.println("Running ... ");
+      tm().transact(
+              () -> {
+                try {
+                  printStream.println(execute());
+                } catch (Exception e) {
+                  rethrow(e);
+                }
+              });
+
+      tm().transact(
+              () -> {
+                try {
+                  printLineIfNotEmpty(postExecute(), printStream);
+                } catch (Exception e) {
+                  rethrow(e);
+                }
+              });
+    } else {
+      printStream.println("Command aborted.");
     }
     printStream.close();
     errorPrintStream.close();

diff --git a/core/src/main/java/google/registry/tools/WhoisQueryCommand.java b/core/src/main/java/google/registry/tools/WhoisQueryCommand.java
@@ -14,6 +14,8 @@
 
 package google.registry.tools;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
@@ -43,6 +45,8 @@ final class WhoisQueryCommand implements Command {
 
   @Override
   public void run() {
-    System.out.println(whois.lookup(Joiner.on(' ').join(mainParameters), unicode, fullOutput));
+    System.out.println(
+        tm().transact(
+                () -> whois.lookup(Joiner.on(' ').join(mainParameters), unicode, fullOutput)));
   }
 }
diff --git a/core/src/main/java/google/registry/ui/server/registrar/ConsoleOteSetupAction.java b/core/src/main/java/google/registry/ui/server/registrar/ConsoleOteSetupAction.java
@@ -53,7 +53,9 @@
     service = Action.Service.DEFAULT,
     path = ConsoleOteSetupAction.PATH,
     method = {Method.POST, Method.GET},
-    auth = Auth.AUTH_PUBLIC_LEGACY)
+    auth = Auth.AUTH_PUBLIC_LEGACY,
+    transactional = false // Needs to explicitly run more than one transaction sequentially.
+    )
 public final class ConsoleOteSetupAction extends HtmlAction {
 
   public static final String PATH = "/registrar-ote-setup";

diff --git a/core/src/main/java/google/registry/ui/server/registrar/RegistrarSettingsAction.java b/core/src/main/java/google/registry/ui/server/registrar/RegistrarSettingsAction.java
@@ -78,7 +78,9 @@
     service = Action.Service.DEFAULT,
     path = RegistrarSettingsAction.PATH,
     method = Action.Method.POST,
-    auth = Auth.AUTH_PUBLIC_LOGGED_IN)
+    auth = Auth.AUTH_PUBLIC_LOGGED_IN,
+    transactional = false // Needs to explicitly run more than one transaction sequentially.
+    )
 public class RegistrarSettingsAction implements Runnable, JsonActionRunner.JsonAction {
 
   public static final String PATH = "/registrar-settings";
@@ -157,9 +159,9 @@ public Map<String, Object> handleJsonRequest(Map<String, ?> input) {
     try {
       switch (op) {
         case "update":
-          return update(args, registrarId).toJsonResponse();
+          return update(args, registrarId);
         case "read":
-          return read(registrarId).toJsonResponse();
+          return read(registrarId);
         default:
           throw new IllegalArgumentException("Unknown or unsupported operation: " + op);
       }
@@ -215,8 +217,11 @@ static EmailInfo create(
     }
   }
 
-  private RegistrarResult read(String registrarId) {
-    return RegistrarResult.create("Success", loadRegistrarUnchecked(registrarId));
+  private Map<String, Object> read(String registrarId) {
+    return tm().transact(
+            () ->
+                RegistrarResult.create("Success", loadRegistrarUnchecked(registrarId))
+                    .toJsonResponse());
   }
 
   private Registrar loadRegistrarUnchecked(String registrarId) {
@@ -227,15 +232,19 @@ private Registrar loadRegistrarUnchecked(String registrarId) {
     }
   }
 
-  private RegistrarResult update(final Map<String, ?> args, String registrarId) {
+  private Map<String, Object> update(final Map<String, ?> args, String registrarId) {
     // Email the updates
     sendExternalUpdatesIfNecessary(tm().transact(() -> saveUpdates(args, registrarId)));
     // Reload the result outside the transaction to get the most recent version
-    return RegistrarResult.create("Saved " + registrarId, loadRegistrarUnchecked(registrarId));
+    return tm().transact(
+            () ->
+                RegistrarResult.create("Saved " + registrarId, loadRegistrarUnchecked(registrarId))
+                    .toJsonResponse());
   }
 
   /** Saves the updates and returns info needed for the update email */
   private EmailInfo saveUpdates(final Map<String, ?> args, String registrarId) {
+    tm().assertInTransaction();
     // We load the registrar here rather than outside the transaction - to make
     // sure we have the latest version. This one is loaded inside the transaction, so it's
     // guaranteed to not change before we update it.

diff --git a/core/src/test/java/google/registry/export/SyncGroupMembersActionTest.java b/core/src/test/java/google/registry/export/SyncGroupMembersActionTest.java
@@ -69,7 +69,7 @@ private void runAction() {
     action.gSuiteDomainName = "domain-registry.example";
     action.response = response;
     action.retrier = new Retrier(new FakeSleeper(new FakeClock()), 3);
-    action.run();
+    tm().transact(action::run);
   }
 
   @Test
@@ -182,7 +182,11 @@ void test_doPost_addsAndRemovesContacts_acrossMultipleRegistrars() throws Except
         "[email protected]",
         Role.MEMBER);
     verify(response).setStatus(SC_OK);
-    assertThat(Iterables.filter(Registrar.loadAll(), Registrar::getContactsRequireSyncing))
+    assertThat(
+            tm().transact(
+                    () ->
+                        Iterables.filter(
+                            Registrar.loadAll(), Registrar::getContactsRequireSyncing)))
         .isEmpty();
   }
 

diff --git a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
@@ -2544,11 +2544,13 @@ void testFailure_disabledRegistrarCantCreateDomain() {
   private void doFailingTest_invalidRegistrarState(State registrarState) {
     persistContactsAndHosts();
     persistResource(
-        Registrar.loadByRegistrarId("TheRegistrar")
-            .get()
-            .asBuilder()
-            .setState(registrarState)
-            .build());
+        tm().transact(
+                () ->
+                    Registrar.loadByRegistrarId("TheRegistrar")
+                        .get()
+                        .asBuilder()
+                        .setState(registrarState)
+                        .build()));
     EppException thrown =
         assertThrows(RegistrarMustBeActiveForThisOperationException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();

diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java
@@ -977,9 +977,9 @@ public static <R extends ImmutableObject> R persistResource(final R resource) {
     assertWithMessage("Attempting to persist a Builder is almost certainly an error in test code")
         .that(resource)
         .isNotInstanceOf(Buildable.Builder.class);
-    tm().transact(() -> tm().put(resource));
+    tm().reTransact(() -> tm().put(resource));
     maybeAdvanceClock();
-    return tm().transact(() -> tm().loadByEntity(resource));
+    return tm().reTransact(() -> tm().loadByEntity(resource));
   }
 
   /** Persists the specified resources to the DB. */
@@ -1161,7 +1161,7 @@ public static ImmutableMap<String, PremiumEntry> loadPremiumEntries(PremiumList
   /** Loads and returns the registrar with the given client ID, or throws IAE if not present. */
   public static Registrar loadRegistrar(String registrarId) {
     return checkArgumentPresent(
-        Registrar.loadByRegistrarId(registrarId),
+        tm().reTransact(() -> Registrar.loadByRegistrarId(registrarId)),
         "Error in tests: Registrar %s does not exist",
         registrarId);
   }

diff --git a/core/src/test/java/google/registry/webdriver/RegistrarConsoleScreenshotTest.java b/core/src/test/java/google/registry/webdriver/RegistrarConsoleScreenshotTest.java
@@ -18,6 +18,7 @@
 import static google.registry.persistence.transaction.JpaTransactionManagerExtension.makeRegistrar2;
 import static google.registry.persistence.transaction.JpaTransactionManagerExtension.makeRegistrarContact2;
 import static google.registry.persistence.transaction.JpaTransactionManagerExtension.makeRegistrarContact3;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.server.Fixture.BASIC;
 import static google.registry.server.Route.route;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -172,13 +173,15 @@ void settingsContactEdit_setRegistryLockPassword() throws Throwable {
     driver.diffPage("contact_view");
 
     RegistrarPoc contact =
-        loadRegistrar("TheRegistrar").getContacts().stream()
-            .filter(c -> "[email protected]".equals(c.getEmailAddress()))
-            .findFirst()
-            .get();
-          assertThat(contact.verifyRegistryLockPassword("password")).isTrue();
-          assertThat(contact.getRegistryLockEmailAddress())
-              .isEqualTo(Optional.of("[email protected]"));
+        tm().transact(
+                () ->
+                    loadRegistrar("TheRegistrar").getContacts().stream()
+                        .filter(c -> "[email protected]".equals(c.getEmailAddress()))
+                        .findFirst()
+                        .get());
+    assertThat(contact.verifyRegistryLockPassword("password")).isTrue();
+    assertThat(contact.getRegistryLockEmailAddress())
+        .isEqualTo(Optional.of("[email protected]"));
   }
 
   @RetryingTest(3)

diff --git a/core/src/test/java/google/registry/webdriver/RegistrarConsoleWebTest.java b/core/src/test/java/google/registry/webdriver/RegistrarConsoleWebTest.java
@@ -15,6 +15,7 @@
 package google.registry.webdriver;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.server.Fixture.BASIC;
 import static google.registry.server.Route.route;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
@@ -176,7 +177,8 @@ void testWhoisSettingsEdit() throws Throwable {
   void testContactSettingsView() throws Throwable {
     driver.get(server.getUrl("/registrar#contact-settings"));
     driver.waitForDisplayedElement(By.id("reg-app-btn-add"));
-    ImmutableList<RegistrarPoc> contacts = loadRegistrar("TheRegistrar").getContacts().asList();
+    ImmutableList<RegistrarPoc> contacts =
+        tm().transact(() -> loadRegistrar("TheRegistrar").getContacts().asList());
     for (RegistrarPoc contact : contacts) {
       assertEltTextPresent(By.id("contacts[0].name"), contact.getName());
       assertEltTextPresent(By.id("contacts[0].emailAddress"), contact.getEmailAddress());