Update SQL queries for reading and comming migration

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java

@@ -1264,9 +1264,10 @@ public interface CollectionDAO {
     int listDistinctCount();
 
     @Deprecated
-    @SqlQuery("SELECT DISTINCT fromFQN, toFQN FROM field_relationship LIMIT :limit OFFSET :offset")
+    @SqlQuery(
+        "SELECT DISTINCT fromFQN, toFQN FROM field_relationship WHERE fromFQNHash = '' or fromFQNHash is null or toFQNHash = '' or toFQNHash is null LIMIT :limit")
     @RegisterRowMapper(FieldRelationShipMapper.class)
-    List<Pair<String, String>> listDistinctWithOffset(@Bind("limit") int limit, @Bind("offset") int offset);
+    List<Pair<String, String>> migrationListDistinctWithOffset(@Bind("limit") int limit);
 
     @SqlQuery(
         "SELECT fromFQN, toFQN, json FROM field_relationship WHERE "
@@ -3404,9 +3405,10 @@ public interface CollectionDAO {
       }
     }
 
-    @SqlQuery("SELECT DISTINCT entityFQN FROM entity_extension_time_series LIMIT :limit OFFSET :offset")
+    @SqlQuery(
+        "SELECT DISTINCT entityFQN FROM entity_extension_time_series WHERE entityFQNHash = '' or entityFQNHash is null LIMIT :limit")
     @Deprecated
-    List<String> listDistinctWithOffset(@Bind("limit") int limit, @Bind("offset") int offset);
+    List<String> migrationListDistinctWithOffset(@Bind("limit") int limit);
   }
 
   class EntitiesCountRowMapper implements RowMapper<EntitiesCount> {

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityDAO.java

@@ -227,6 +227,10 @@ public interface EntityDAO<T extends EntityInterface> {
   @SqlQuery("SELECT json FROM <table> LIMIT :limit OFFSET :offset")
   List<String> listAfterWithOffset(@Define("table") String table, @Bind("limit") int limit, @Bind("offset") int offset);
 
+  @SqlQuery("SELECT json FROM <table> WHERE <nameHashColumn> = '' or <nameHashColumn> is null LIMIT :limit")
+  List<String> migrationListAfterWithOffset(
+      @Define("table") String table, @Define("nameHashColumn") String nameHashColumnName, @Bind("limit") int limit);
+
   @SqlQuery("SELECT json FROM <table> <cond> AND " + "ORDER BY <nameColumn> " + "LIMIT :limit " + "OFFSET :offset")
   List<String> listAfter(
       @Define("table") String table,
@@ -351,6 +355,11 @@ public interface EntityDAO<T extends EntityInterface> {
     return listAfterWithOffset(getTableName(), limit, offset);
   }
 
+  default List<String> migrationListAfterWithOffset(int limit) {
+    // No ordering
+    return migrationListAfterWithOffset(getTableName(), getNameHashColumn(), limit);
+  }
+
   default List<String> listAfter(ListFilter filter, int limit, int offset) {
     return listAfter(getTableName(), getNameHashColumn(), filter.getCondition(), limit, offset);
   }

openmetadata-service/src/main/java/org/openmetadata/service/migration/MigrationUtil.java

@@ -147,13 +147,14 @@ public class MigrationUtil {
       Handle handle, String updateSql, Class<T> clazz, EntityDAO<T> dao, boolean withName, int limitParam)
       throws IOException {
     LOG.debug("Starting Migration for table : {}", dao.getTableName());
-    int offset = 0;
+    while (true) {
-    int totalCount = dao.listTotalCount();
-    while (offset < totalCount) {
       PreparedBatch upsertBatch = handle.prepareBatch(updateSql);
       // Read from Database
-      List<String> jsons = dao.listAfterWithOffset(limitParam, offset);
+      List<String> jsons = dao.migrationListAfterWithOffset(limitParam);
-      offset = offset + limitParam;
+      LOG.debug("[{}]Read a Batch of Size: {}", dao.getTableName(), jsons.size());
+      if (jsons.isEmpty()) {
+        break;
+      }
       // Process Update
       for (String json : jsons) {
         // Update the Statements to Database
@@ -168,6 +169,7 @@ public class MigrationUtil {
           LOG.error("Failed in creating FQN Hash for Entity Name : {}", entity.getFullyQualifiedName(), ex);
         }
       }
+      LOG.debug("[{}]Committing a Batch of Size: {}", dao.getTableName(), jsons.size());
       upsertBatch.execute();
     }
     LOG.debug("End Migration for table : {}", dao.getTableName());
@@ -276,19 +278,14 @@ public class MigrationUtil {
   private static void updateFQNHashForFieldRelationship(
       Handle handle, String updateSql, CollectionDAO collectionDAO, int limitParam) {
     LOG.debug("Starting Migration for Field Relationship");
-    int offset = 0;
+    while (true) {
-    int totalCount;
-    try {
-      // This might result into exceptions if the column entityFQN is dropped once
-      totalCount = collectionDAO.fieldRelationshipDAO().listDistinctCount();
-    } catch (Exception ex) {
-      return;
-    }
-    if (totalCount > 0) {
-      while (offset < totalCount) {
       PreparedBatch upsertBatch = handle.prepareBatch(updateSql);
       List<Pair<String, String>> entityFQNPairList =
-            collectionDAO.fieldRelationshipDAO().listDistinctWithOffset(limitParam, offset);
+          collectionDAO.fieldRelationshipDAO().migrationListDistinctWithOffset(limitParam);
+      LOG.debug("[FieldRelationship] Read a Batch of Size: {}", entityFQNPairList.size());
+      if (entityFQNPairList.isEmpty()) {
+        break;
+      }
       for (Pair<String, String> entityFQNPair : entityFQNPairList) {
         try {
           String fromFQNHash = FullyQualifiedName.buildHash(entityFQNPair.getLeft());
@@ -304,9 +301,8 @@ public class MigrationUtil {
               "Failed in creating fromFQN : {} , toFQN : {}", entityFQNPair.getLeft(), entityFQNPair.getRight(), ex);
         }
       }
+      LOG.debug("[FieldRelationship] Committing a Batch of Size: {}", entityFQNPairList.size());
       upsertBatch.execute();
-        offset = offset + limitParam;
-      }
     }
     LOG.debug("End Migration for Field Relationship");
   }
@@ -314,32 +310,23 @@ public class MigrationUtil {
   private static void updateFQNHashEntityExtensionTimeSeries(
       Handle handle, String updateSql, CollectionDAO collectionDAO, int limitParam) {
     LOG.debug("Starting Migration for Entity Extension Time Series");
-    int offset = 0;
+    while (true) {
-    int totalCount;
-    try {
-      // This might result into exceptions if the column entityFQN is dropped once
-      totalCount = collectionDAO.entityExtensionTimeSeriesDao().listDistinctCount();
-    } catch (Exception ex) {
-      return;
-    }
-    if (totalCount > 0) {
-      while (offset < totalCount) {
       PreparedBatch upsertBatch = handle.prepareBatch(updateSql);
       List<String> entityFQNLists =
-            collectionDAO.entityExtensionTimeSeriesDao().listDistinctWithOffset(limitParam, offset);
+          collectionDAO.entityExtensionTimeSeriesDao().migrationListDistinctWithOffset(limitParam);
+      LOG.debug("[TimeSeries] Read a Batch of Size: {}", entityFQNLists.size());
+      if (entityFQNLists.isEmpty()) {
+        break;
+      }
       for (String entityFQN : entityFQNLists) {
         try {
-            upsertBatch
+          upsertBatch.bind("entityFQNHash", FullyQualifiedName.buildHash(entityFQN)).bind("entityFQN", entityFQN).add();
-                .bind("entityFQNHash", FullyQualifiedName.buildHash(entityFQN))
-                .bind("entityFQN", entityFQN)
-                .add();
         } catch (Exception ex) {
           LOG.error("Failed in creating EntityFQN : {}", entityFQN, ex);
         }
       }
+      LOG.debug("[TimeSeries] Committing a Batch of Size: {}", entityFQNLists.size());
       upsertBatch.execute();
-        offset = offset + limitParam;
-      }
     }
     LOG.debug("Ended Migration for Entity Extension Time Series");
   }