From e75c1c33eb81eb9dec43f282e6eee7e5a09d8515 Mon Sep 17 00:00:00 2001
From: Weimin Yu <weiminyu@google.com>
Date: Fri, 10 Sep 2021 21:23:07 -0400
Subject: [PATCH] Add a temporary fix to Hibernate detach in BEAM queries
 (#1314)

* Add a temporary fix to Hibernate detach in query

Make all queries in RegistryQuery (exclusively used by BEAM) use
EntityManager.clear() to detach entities. This is a temporary measure
that unblocks work in BEAM. We will revert the work once
JpaTransactionManager can detach entities properly for all types of
queries.

Also fixed regression bugs that broke query result streaming:
- The code that sets query fetch size was not carried over from
QueryComposer.
- The new ReadOnlyCheckingTypedQuery class did not override parent's
getResultStream() method, which calls getList().
---
 .../registry/beam/common/RegistryQuery.java   | 31 ++++++++++++++++---
 .../ReadOnlyCheckingTypedQuery.java           |  6 ++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/core/src/main/java/google/registry/beam/common/RegistryQuery.java b/core/src/main/java/google/registry/beam/common/RegistryQuery.java
index f61dc01d5..dba1cf5b3 100644
--- a/core/src/main/java/google/registry/beam/common/RegistryQuery.java
+++ b/core/src/main/java/google/registry/beam/common/RegistryQuery.java
@@ -16,6 +16,7 @@ package google.registry.beam.common;
 
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
+import google.registry.persistence.transaction.JpaTransactionManager;
 import java.io.Serializable;
 import java.util.Map;
 import java.util.function.Supplier;
@@ -28,6 +29,15 @@ import javax.persistence.criteria.CriteriaQuery;
 
 /** Interface for query instances used by {@link RegistryJpaIO.Read}. */
 public interface RegistryQuery<T> extends Serializable {
+
+  /**
+   * Number of JPA entities to fetch in each batch during a query.
+   *
+   * <p>With Hibernate, for result streaming to work, a query's fetchSize property must be set to a
+   * non-zero value.
+   */
+  int QUERY_FETCH_SIZE = 1000;
+
   Stream<T> stream();
 
   interface CriteriaQuerySupplier<T> extends Supplier<CriteriaQuery<T>>, Serializable {}
@@ -49,6 +59,7 @@ public interface RegistryQuery<T> extends Serializable {
       if (parameters != null) {
         parameters.forEach(query::setParameter);
       }
+      JpaTransactionManager.setQueryFetchSize(query, QUERY_FETCH_SIZE);
       @SuppressWarnings("unchecked")
       Stream<T> resultStream = query.getResultStream();
       return nativeQuery ? resultStream : resultStream.map(e -> detach(entityManager, e));
@@ -64,11 +75,14 @@ public interface RegistryQuery<T> extends Serializable {
   static <T> RegistryQuery<T> createQuery(
       String jpql, @Nullable Map<String, Object> parameters, Class<T> clazz) {
     return () -> {
-      TypedQuery<T> query = jpaTm().query(jpql, clazz);
+      // TODO(b/193662898): switch to jpaTm().query() when it can properly detach loaded entities.
+      EntityManager entityManager = jpaTm().getEntityManager();
+      TypedQuery<T> query = entityManager.createQuery(jpql, clazz);
       if (parameters != null) {
         parameters.forEach(query::setParameter);
       }
-      return query.getResultStream();
+      JpaTransactionManager.setQueryFetchSize(query, QUERY_FETCH_SIZE);
+      return query.getResultStream().map(e -> detach(entityManager, e));
     };
   }
 
@@ -82,7 +96,13 @@ public interface RegistryQuery<T> extends Serializable {
    * @param <T> Type of each row in the result set.
    */
   static <T> RegistryQuery<T> createQuery(CriteriaQuerySupplier<T> criteriaQuery) {
-    return () -> jpaTm().query(criteriaQuery.get()).getResultStream();
+    return () -> {
+      // TODO(b/193662898): switch to jpaTm().query() when it can properly detach loaded entities.
+      EntityManager entityManager = jpaTm().getEntityManager();
+      TypedQuery<T> query = entityManager.createQuery(criteriaQuery.get());
+      JpaTransactionManager.setQueryFetchSize(query, QUERY_FETCH_SIZE);
+      return query.getResultStream().map(e -> detach(entityManager, e));
+    };
   }
 
   /**
@@ -108,7 +128,10 @@ public interface RegistryQuery<T> extends Serializable {
       return;
     }
     try {
-      entityManager.detach(object);
+      // TODO(b/193662898): choose detach() or clear() based on the type of transaction.
+      // For context, EntityManager.detach() does not remove all metadata about loaded entities.
+      // See b/193925312 or https://hibernate.atlassian.net/browse/HHH-14735 for details.
+      entityManager.clear();
     } catch (IllegalArgumentException e) {
       // Not an entity. Do nothing.
     }
diff --git a/core/src/main/java/google/registry/persistence/transaction/ReadOnlyCheckingTypedQuery.java b/core/src/main/java/google/registry/persistence/transaction/ReadOnlyCheckingTypedQuery.java
index dcda3a136..f0077c835 100644
--- a/core/src/main/java/google/registry/persistence/transaction/ReadOnlyCheckingTypedQuery.java
+++ b/core/src/main/java/google/registry/persistence/transaction/ReadOnlyCheckingTypedQuery.java
@@ -21,6 +21,7 @@ import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Stream;
 import javax.persistence.FlushModeType;
 import javax.persistence.LockModeType;
 import javax.persistence.Parameter;
@@ -41,6 +42,11 @@ class ReadOnlyCheckingTypedQuery<T> implements TypedQuery<T> {
     return delegate.getResultList();
   }
 
+  @Override
+  public Stream<T> getResultStream() {
+    return delegate.getResultStream();
+  }
+
   @Override
   public T getSingleResult() {
     return delegate.getSingleResult();