[#12048] V9 migration verification script optimisation - fetch ReadNo…

…tifications for account comparison (#12905) * Eager load readNotifications for accounts when verifying --------- Co-authored-by: Zhang Ziqing <[email protected]>
TEAMMATES · Mar 18, 2024 · 868982f · 868982f
1 parent ffb07dd
commit 868982f
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 7 deletions.
diff --git a/src/client/java/teammates/client/scripts/sql/VerifyAccountAttributes.java b/src/client/java/teammates/client/scripts/sql/VerifyAccountAttributes.java
@@ -1,11 +1,21 @@
 package teammates.client.scripts.sql;
 
+// CHECKSTYLE.OFF:ImportOrder
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 
+import jakarta.persistence.TypedQuery;
+import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
+import jakarta.persistence.criteria.JoinType;
+import jakarta.persistence.criteria.Order;
+import jakarta.persistence.criteria.Root;
+
+import teammates.common.util.HibernateUtil;
 import teammates.storage.entity.Account;
 import teammates.storage.sqlentity.ReadNotification;
 
@@ -16,6 +26,8 @@
 public class VerifyAccountAttributes
         extends VerifyNonCourseEntityAttributesBaseScript<Account, teammates.storage.sqlentity.Account> {
 
+    private static final String READ_NOTIFICATION_FIELD = "readNotifications";
+
     public VerifyAccountAttributes() {
         super(Account.class,
                 teammates.storage.sqlentity.Account.class);
@@ -46,6 +58,28 @@ public boolean verifyAccountFields(teammates.storage.sqlentity.Account sqlEntity
 
     }
 
+    @Override
+    protected List<teammates.storage.sqlentity.Account> lookupSqlEntitiesByPageNumber(int pageNum) {
+        CriteriaBuilder cb = HibernateUtil.getCriteriaBuilder();
+        CriteriaQuery<teammates.storage.sqlentity.Account> pageQuery = cb.createQuery(sqlEntityClass);
+
+        // sort by id to maintain stable order.
+        Root<teammates.storage.sqlentity.Account> root = pageQuery.from(sqlEntityClass);
+        pageQuery.select(root);
+        List<Order> orderList = new LinkedList<>();
+        orderList.add(cb.asc(root.get("id")));
+        pageQuery.orderBy(orderList);
+
+        // perform query with pagination
+        TypedQuery<teammates.storage.sqlentity.Account> query = HibernateUtil.createQuery(pageQuery);
+        query.setFirstResult(calculateOffset(pageNum));
+        query.setMaxResults(CONST_SQL_FETCH_BASE_SIZE);
+
+        // Fetch read notifications eagerly with one join
+        root.fetch(READ_NOTIFICATION_FIELD, JoinType.LEFT);
+        return query.getResultList();
+    }
+
     // Used for sql data migration
     @Override
     public boolean equals(teammates.storage.sqlentity.Account sqlEntity, Account datastoreEntity) {

diff --git a/src/client/java/teammates/client/scripts/sql/VerifyNonCourseEntityAttributesBaseScript.java b/src/client/java/teammates/client/scripts/sql/VerifyNonCourseEntityAttributesBaseScript.java
@@ -33,14 +33,19 @@ public abstract class VerifyNonCourseEntityAttributesBaseScript<E extends teamma
      * for optimized batch-fetching.
     */
 
-    private static int constSqlFetchBaseSize = 1000;
+    /**
+     * Batch size to fetch per page.
+     */
+    protected static final int CONST_SQL_FETCH_BASE_SIZE = 1000;
 
     /** Datastore entity class. */
     protected Class<E> datastoreEntityClass;
 
     /** SQL entity class. */
     protected Class<T> sqlEntityClass;
 
+    private long entitiesVerified = 0;
+
     public VerifyNonCourseEntityAttributesBaseScript(
             Class<E> datastoreEntityClass, Class<T> sqlEntityClass) {
         this.datastoreEntityClass = datastoreEntityClass;
@@ -77,8 +82,8 @@ protected Map<String, E> lookupDataStoreEntities(List<String> datastoreEntitiesI
     /**
      * Calculate offset.
      */
-    private int calculateOffset(int pageNum) {
-        return (pageNum - 1) * constSqlFetchBaseSize;
+    protected int calculateOffset(int pageNum) {
+        return (pageNum - 1) * CONST_SQL_FETCH_BASE_SIZE;
     }
 
     /**
@@ -89,17 +94,22 @@ private int getNumPages() {
         CriteriaQuery<Long> countQuery = cb.createQuery(Long.class);
         countQuery.select(cb.count(countQuery.from(sqlEntityClass)));
         long countResults = HibernateUtil.createQuery(countQuery).getSingleResult().longValue();
-        int numPages = (int) (Math.ceil((double) countResults / (double) constSqlFetchBaseSize));
+        int numPages = (int) (Math.ceil((double) countResults / (double) CONST_SQL_FETCH_BASE_SIZE));
         log(String.format("Has %d entities with %d pages", countResults, numPages));
 
         return numPages;
     }
 
-    private List<T> lookupSqlEntitiesByPageNumber(int pageNum) {
+    /**
+     * Sort SQL entities by id in ascending order and return entities on page.
+     * @param pageNum page in a sorted entities tables
+     * @return list of SQL entities on page num
+     */
+    protected List<T> lookupSqlEntitiesByPageNumber(int pageNum) {
         CriteriaBuilder cb = HibernateUtil.getCriteriaBuilder();
         CriteriaQuery<T> pageQuery = cb.createQuery(sqlEntityClass);
 
-        // sort by createdAt to maintain stable order.
+        // sort by id to maintain stable order.
         Root<T> root = pageQuery.from(sqlEntityClass);
         pageQuery.select(root);
         List<Order> orderList = new LinkedList<>();
@@ -109,7 +119,7 @@ private List<T> lookupSqlEntitiesByPageNumber(int pageNum) {
         // perform query with pagination
         TypedQuery<T> query = HibernateUtil.createQuery(pageQuery);
         query.setFirstResult(calculateOffset(pageNum));
-        query.setMaxResults(constSqlFetchBaseSize);
+        query.setMaxResults(CONST_SQL_FETCH_BASE_SIZE);
 
         return query.getResultList();
     }
@@ -149,15 +159,18 @@ protected List<Map.Entry<T, E>> checkAllEntitiesForFailures() {
                     + (endTimeForDatastore - startTimeForDatastore) + " milliseconds");
 
             long startTimeForEquals = System.currentTimeMillis();
+            entitiesVerified += sqlEntities.size();
             for (T sqlEntity : sqlEntities) {
                 E datastoreEntity = datastoreEntities.get(generateID(sqlEntity));
                 if (datastoreEntity == null) {
+                    entitiesVerified -= 1;
                     failures.add(new AbstractMap.SimpleEntry<T, E>(sqlEntity, null));
                     continue;
                 }
 
                 boolean isEqual = equals(sqlEntity, datastoreEntity);
                 if (!isEqual) {
+                    entitiesVerified -= 1;
                     failures.add(new AbstractMap.SimpleEntry<T, E>(sqlEntity, datastoreEntity));
                     continue;
                 }
@@ -175,6 +188,7 @@ protected List<Map.Entry<T, E>> checkAllEntitiesForFailures() {
     protected void runCheckAllEntities(Class<T> sqlEntityClass,
             Class<E> datastoreEntityClass) {
         HibernateUtil.beginTransaction();
+        long checkStartTime = System.currentTimeMillis();
         List<Map.Entry<T, E>> failedEntities = checkAllEntitiesForFailures();
 
         System.out.println("========================================");
@@ -186,6 +200,10 @@ protected void runCheckAllEntities(Class<T> sqlEntityClass,
         } else {
             log("No errors detected");
         }
+
+        long checkEndTime = System.currentTimeMillis();
+        log("Entity took " + (checkEndTime - checkStartTime) + " milliseconds to verify");
+        log("Verified " + entitiesVerified + " SQL entities successfully");
         HibernateUtil.commitTransaction();
     }