Convert CountDomainsCommand to tm (#1092)

* Convert CountDomainsCommand to tm

As part of this, implement "select count(*)" queries in the QueryComposer.

* Replaced kludgy trick for objectify count

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -413,5 +413,10 @@ public class DatastoreTransactionManager implements TransactionManager {
     public Stream<T> stream() {
       return Streams.stream(buildQuery());
     }
+
+    @Override
+    public long count() {
+      return buildQuery().count();
+    }
   }
 }

core/src/main/java/google/registry/persistence/transaction/CriteriaQueryBuilder.java

@@ -41,11 +41,11 @@ public class CriteriaQueryBuilder<T> {
   }
 
   private final CriteriaQuery<T> query;
-  private final Root<T> root;
+  private final Root<?> root;
   private final ImmutableList.Builder<Predicate> predicates = new ImmutableList.Builder<>();
   private final ImmutableList.Builder<Order> orders = new ImmutableList.Builder<>();
 
-  private CriteriaQueryBuilder(CriteriaQuery<T> query, Root<T> root) {
+  private CriteriaQueryBuilder(CriteriaQuery<T> query, Root<?> root) {
     this.query = query;
     this.root = root;
   }
@@ -106,4 +106,13 @@ public class CriteriaQueryBuilder<T> {
     query = query.select(root);
     return new CriteriaQueryBuilder<>(query, root);
   }
+
+  /** Creates a "count" query for the table for the class. */
+  public static <T> CriteriaQueryBuilder<Long> createCount(EntityManager em, Class<T> clazz) {
+    CriteriaBuilder builder = em.getCriteriaBuilder();
+    CriteriaQuery<Long> query = builder.createQuery(Long.class);
+    Root<T> root = query.from(clazz);
+    query = query.select(builder.count(root));
+    return new CriteriaQueryBuilder<>(query, root);
+  }
 }

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -700,7 +700,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
     private TypedQuery<T> buildQuery() {
       CriteriaQueryBuilder<T> queryBuilder = CriteriaQueryBuilder.create(em, entityClass);
+      return addCriteria(queryBuilder);
+    }
 
+    private <U> TypedQuery<U> addCriteria(CriteriaQueryBuilder<U> queryBuilder) {
       for (WhereClause<?> pred : predicates) {
         pred.addToCriteriaQueryBuilder(queryBuilder);
       }
@@ -727,5 +730,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     public Stream<T> stream() {
       return buildQuery().getResultStream();
     }
+
+    @Override
+    public long count() {
+      CriteriaQueryBuilder<Long> queryBuilder = CriteriaQueryBuilder.createCount(em, entityClass);
+      return addCriteria(queryBuilder).getSingleResult();
+    }
   }
 }

core/src/main/java/google/registry/persistence/transaction/QueryComposer.java

@@ -88,6 +88,9 @@ public abstract class QueryComposer<T> {
   /** Returns the results of the query as a stream. */
   public abstract Stream<T> stream();
 
+  /** Returns the number of results of the query. */
+  public abstract long count();
+
   // We have to wrap the CriteriaQueryBuilder predicate factories in our own functions because at
   // the point where we pass them to the Comparator constructor, the compiler can't determine which
   // of the overloads to use since there is no "value" object for context.

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

@@ -14,12 +14,13 @@
 
 package google.registry.tools;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registries.assertTldsExist;
+import static google.registry.persistence.transaction.QueryComposer.Comparator;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.Iterables;
 import google.registry.model.domain.DomainBase;
 import google.registry.util.Clock;
 import java.util.List;
@@ -45,14 +46,12 @@ final class CountDomainsCommand implements CommandWithRemoteApi {
         .forEach(tld -> System.out.printf("%s,%d\n", tld, getCountForTld(tld, now)));
   }
 
-  private int getCountForTld(String tld, DateTime now) {
-    return Iterables.size(
-        ofy()
-            .load()
-            .type(DomainBase.class)
-            .filter("tld", tld)
-            .filter("deletionTime >", now)
-            .chunkAll()
-            .keys());
+  private long getCountForTld(String tld, DateTime now) {
+    return transactIfJpaTm(
+        () ->
+            tm().createQueryComposer(DomainBase.class)
+                .where("tld", Comparator.EQ, tld)
+                .where("deletionTime", Comparator.GT, now)
+                .count());
   }
 }

core/src/test/java/google/registry/persistence/transaction/QueryComposerTest.java

@@ -111,6 +111,17 @@ public class QueryComposerTest {
         .isEqualTo(alpha);
   }
 
+  @TestOfyAndSql
+  public void testCount() {
+    assertThat(
+            transactIfJpaTm(
+                () ->
+                    tm().createQueryComposer(TestEntity.class)
+                        .where("name", Comparator.GTE, "bravo")
+                        .count()))
+        .isEqualTo(2L);
+  }
+
   @TestOfyAndSql
   public void testGetSingleResult() {
     assertThat(

core/src/test/java/google/registry/tools/CountDomainsCommandTest.java

@@ -19,12 +19,14 @@ import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
 
 import google.registry.model.ofy.Ofy;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link CountDomainsCommand}. */
+@DualDatabaseTest
 public class CountDomainsCommandTest extends CommandTestCase<CountDomainsCommand> {
 
   @RegisterExtension public final InjectExtension inject = new InjectExtension();
@@ -36,7 +38,7 @@ public class CountDomainsCommandTest extends CommandTestCase<CountDomainsCommand
     createTlds("foo", "bar", "baz", "qux");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_singleTld() throws Exception {
     for (int i = 0; i < 51; i++) {
       persistActiveDomain(String.format("test-%d.foo", i));
@@ -48,7 +50,7 @@ public class CountDomainsCommandTest extends CommandTestCase<CountDomainsCommand
     assertStdoutIs("foo,51\n");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_multipleTlds() throws Exception {
     for (int i = 0; i < 29; i++) {
       persistActiveDomain(String.format("test-%d.foo", i));