From 34b3dd037ac77482a8e12d23596ac58376bf9aa7 Mon Sep 17 00:00:00 2001
From: Ben McIlwain <mcilwain@google.com>
Date: Wed, 5 Jul 2023 15:09:20 -0400
Subject: [PATCH] Add minor refactoring follow-up for
 RefreshDnsForAllDomainsAction (#2063)

This is a follow-on to comments in PR #2037. It makes the main loop cleaner and
also removes ambiguities around database handling when the first query is run
with the cursor still empty because no results have been found yet.
---
 .../server/RefreshDnsForAllDomainsAction.java | 37 +++++++++++--------
 .../RefreshDnsForAllDomainsActionTest.java    |  4 +-
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/core/src/main/java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java b/core/src/main/java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java
index de89ffe95..d3275b6a0 100644
--- a/core/src/main/java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java
+++ b/core/src/main/java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java
@@ -32,7 +32,9 @@ import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import java.util.Optional;
 import java.util.Random;
+import javax.annotation.Nullable;
 import javax.inject.Inject;
+import javax.persistence.TypedQuery;
 import org.apache.arrow.util.VisibleForTesting;
 import org.apache.http.HttpStatus;
 import org.joda.time.Duration;
@@ -85,11 +87,14 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
     assertTldsExist(tlds);
     checkArgument(batchSize > 0, "Must specify a positive number for batch size");
     int smearMinutes = tm().transact(this::calculateSmearMinutes);
-    ImmutableList<String> previousBatch = ImmutableList.of("");
+
+    ImmutableList<String> domainsBatch;
+    @Nullable String lastInPreviousBatch = null;
     do {
-      String lastInPreviousBatch = getLast(previousBatch);
-      previousBatch = tm().transact(() -> refreshBatch(lastInPreviousBatch, smearMinutes));
-    } while (previousBatch.size() == batchSize);
+      Optional<String> lastInPreviousBatchOpt = Optional.ofNullable(lastInPreviousBatch);
+      domainsBatch = tm().transact(() -> refreshBatch(lastInPreviousBatchOpt, smearMinutes));
+      lastInPreviousBatch = domainsBatch.isEmpty() ? null : getLast(domainsBatch);
+    } while (domainsBatch.size() == batchSize);
   }
 
   /**
@@ -107,22 +112,22 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
     return Math.max(activeDomains.intValue() / 1000, 1);
   }
 
-  private ImmutableList<String> getBatch(String lastInPreviousBatch) {
-    return tm().query(
+  private ImmutableList<String> getBatch(Optional<String> lastInPreviousBatch) {
+    String sql =
+        String.format(
             "SELECT domainName FROM Domain WHERE tld IN (:tlds) AND"
-                + " deletionTime = :endOfTime  AND domainName >"
-                + " :lastInPreviousBatch ORDER BY domainName ASC",
-            String.class)
-        .setParameter("tlds", tlds)
-        .setParameter("endOfTime", END_OF_TIME)
-        .setParameter("lastInPreviousBatch", lastInPreviousBatch)
-        .setMaxResults(batchSize)
-        .getResultStream()
-        .collect(toImmutableList());
+                + " deletionTime = :endOfTime %s ORDER BY domainName ASC",
+            lastInPreviousBatch.isPresent() ? "AND domainName > :lastInPreviousBatch" : "");
+    TypedQuery<String> query =
+        tm().query(sql, String.class)
+            .setParameter("tlds", tlds)
+            .setParameter("endOfTime", END_OF_TIME);
+    lastInPreviousBatch.ifPresent(l -> query.setParameter("lastInPreviousBatch", l));
+    return query.setMaxResults(batchSize).getResultStream().collect(toImmutableList());
   }
 
   @VisibleForTesting
-  ImmutableList<String> refreshBatch(String lastInPreviousBatch, int smearMinutes) {
+  ImmutableList<String> refreshBatch(Optional<String> lastInPreviousBatch, int smearMinutes) {
     ImmutableList<String> domainBatch = getBatch(lastInPreviousBatch);
     try {
       // Smear the task execution time over the next N minutes.
diff --git a/core/src/test/java/google/registry/tools/server/RefreshDnsForAllDomainsActionTest.java b/core/src/test/java/google/registry/tools/server/RefreshDnsForAllDomainsActionTest.java
index 9d6c4b15c..61fbd18af 100644
--- a/core/src/test/java/google/registry/tools/server/RefreshDnsForAllDomainsActionTest.java
+++ b/core/src/test/java/google/registry/tools/server/RefreshDnsForAllDomainsActionTest.java
@@ -75,8 +75,8 @@ public class RefreshDnsForAllDomainsActionTest {
     action =
         new RefreshDnsForAllDomainsAction(
             response, ImmutableSet.of("bar"), Optional.of(1), new Random());
-    tm().transact(() -> action.refreshBatch("", 1000));
-    tm().transact(() -> action.refreshBatch("", 1000));
+    tm().transact(() -> action.refreshBatch(Optional.empty(), 1000));
+    tm().transact(() -> action.refreshBatch(Optional.empty(), 1000));
     ImmutableList<DnsRefreshRequest> refreshRequests =
         tm().transact(
                 () ->