From e4bf10ee9697b9c9b02c6e08e8ab20cad44b8a9a Mon Sep 17 00:00:00 2001
From: shouchengShen <50309488+Johnnyssc@users.noreply.github.com>
Date: Mon, 26 Aug 2024 21:04:34 +0800
Subject: [PATCH] [enhancement](Load)allow load data to the other partitions
 when some partitions are restoring (#39915)

If broker load or stream load task execute in one table that is
restoring data, load task will failed with Exception.
Exception info :"Table [xxx] is under restore" or "Table [xxx] is in
restore process, can't load into it".

But mostly restoreJob only effects some partitions in this table, not
all of them, so that the other partitions still need to load data
successfully.
To achieve this goal, before checking olap table state, check partition
state first.

cherry pick from master branch, pr has been merged:
https://github.com/apache/doris/pull/39595

Co-authored-by: shenshoucheng <shenshoucheng@jd.com>
---
 .../apache/doris/load/BrokerFileGroup.java    | 12 ++++++++++-
 .../transaction/DatabaseTransactionMgr.java   | 20 ++++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/load/BrokerFileGroup.java b/fe/fe-core/src/main/java/org/apache/doris/load/BrokerFileGroup.java
index 6e5ca07106b48e..ff27acadcac2f4 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/load/BrokerFileGroup.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/load/BrokerFileGroup.java
@@ -29,6 +29,7 @@
 import org.apache.doris.catalog.OlapTable;
 import org.apache.doris.catalog.OlapTable.OlapTableState;
 import org.apache.doris.catalog.Partition;
+import org.apache.doris.catalog.Partition.PartitionState;
 import org.apache.doris.catalog.Table;
 import org.apache.doris.common.DdlException;
 import org.apache.doris.common.Pair;
@@ -139,11 +140,20 @@ public void parse(Database db, DataDescription dataDescription) throws DdlExcept
                         throw new DdlException("Unknown partition '" + pName
                                 + "' in table '" + olapTable.getName() + "'");
                     }
+                    // partition which need load data
+                    if (partition.getState() == PartitionState.RESTORE) {
+                        throw new DdlException("Table [" + olapTable.getName()
+                                + "], Partition[" + partition.getName() + "] is under restore");
+                    }
                     partitionIds.add(partition.getId());
                 }
             }
 
-            if (olapTable.getState() == OlapTableState.RESTORE) {
+            boolean isPartitionRestoring = olapTable.getPartitions().stream().anyMatch(
+                    partition -> partition.getState() == PartitionState.RESTORE
+            );
+            // restore table
+            if (!isPartitionRestoring && olapTable.getState() == OlapTableState.RESTORE) {
                 throw new DdlException("Table [" + olapTable.getName() + "] is under restore");
             }
 
diff --git a/fe/fe-core/src/main/java/org/apache/doris/transaction/DatabaseTransactionMgr.java b/fe/fe-core/src/main/java/org/apache/doris/transaction/DatabaseTransactionMgr.java
index 0b5e7a18d043e2..636fe582ed11a9 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/transaction/DatabaseTransactionMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/transaction/DatabaseTransactionMgr.java
@@ -23,7 +23,9 @@
 import org.apache.doris.catalog.Env;
 import org.apache.doris.catalog.MaterializedIndex;
 import org.apache.doris.catalog.OlapTable;
+import org.apache.doris.catalog.OlapTable.OlapTableState;
 import org.apache.doris.catalog.Partition;
+import org.apache.doris.catalog.Partition.PartitionState;
 import org.apache.doris.catalog.PartitionInfo;
 import org.apache.doris.catalog.PartitionType;
 import org.apache.doris.catalog.Replica;
@@ -516,16 +518,24 @@ private void checkCommitStatus(List<Table> tableList, TransactionState transacti
                 continue;
             }
 
-            if (tbl.getState() == OlapTable.OlapTableState.RESTORE) {
-                throw new LoadException("Table " + tbl.getName() + " is in restore process. "
-                        + "Can not load into it");
-            }
-
             long partitionId = tabletMeta.getPartitionId();
             if (tbl.getPartition(partitionId) == null) {
                 // this can happen when partitionId == -1 (tablet being dropping)
                 // or partition really not exist.
                 continue;
+            } else if (tbl.getPartition(partitionId).getState() == PartitionState.RESTORE) {
+                // partition in restore process which can not load data
+                throw new LoadException("Table [" + tbl.getName() + "], Partition ["
+                        + tbl.getPartition(partitionId).getName() + "] is in restore process. Can not load into it");
+            }
+
+            boolean isPartitionRestoring = tbl.getPartitions().stream().anyMatch(
+                    partition -> partition.getState() == PartitionState.RESTORE
+            );
+            // restore table
+            if (!isPartitionRestoring && tbl.getState() == OlapTableState.RESTORE) {
+                throw new LoadException("Table " + tbl.getName() + " is in restore process. "
+                    + "Can not load into it");
             }
 
             if (!tableToPartition.containsKey(tableId)) {