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 6e5ca07106..ff27acadca 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.KeysType;
 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 class BrokerFileGroup implements Writable {
                         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 0b5e7a18d0..636fe582ed 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.DatabaseIf;
 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 @@ public class DatabaseTransactionMgr {
                 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)) {