[To dev/1.3] Fix automatic LOAD retries during region recovery

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java

@@ -27,6 +27,8 @@
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.consensus.ConfigRegionId;
 import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.path.PatternTreeMap;
 import org.apache.iotdb.commons.schema.SchemaConstant;
@@ -591,6 +593,18 @@ private Analysis executeTabletConversionOnException(
 
   private boolean setTemporaryUnavailableStatusIfNecessary(
       final Analysis analysis, final Throwable throwable) {
+    // Region unavailability can persist beyond recovery. Retain the file for retry; permanent
+    // region loss still requires operator intervention.
+    if (throwable instanceof IoTDBRuntimeException
+        && ((IoTDBRuntimeException) throwable).getErrorCode()
+            == TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode()) {
+      analysis.setFinishQueryAfterAnalyze(true);
+      analysis.setFailStatus(
+          RpcUtils.getStatus(
+              TSStatusCode.LOAD_TEMPORARY_UNAVAILABLE_EXCEPTION, throwable.getMessage()));
+      analysis.setStatement(loadTsFileStatement);
+      return true;
+    }
     if (isTemporaryUnavailableDueToPipeSchemaNotReady(throwable)) {
       setFailAnalysisForTemporaryUnavailablePipeSchema(analysis, throwable);
       return true;
@@ -849,6 +863,7 @@ private void doAutoCreateAndVerify()
         if (isTemporaryUnavailableDueToPipeSchemaNotReady(e)) {
           throw e;
         }
+        rethrowIfUnavailableRegion(e);
         LOGGER.warn("Auto create or verify schema error.", e);
         throw new SemanticException(
             String.format(
@@ -858,6 +873,7 @@ private void doAutoCreateAndVerify()
         if (e.getCause() instanceof LoadAnalyzeTypeMismatchException && isConvertOnTypeMismatch) {
           throw (LoadAnalyzeTypeMismatchException) e.getCause();
         }
+        rethrowIfUnavailableRegion(e);
         LOGGER.warn("Auto create or verify schema error.", e);
         throw new SemanticException(
             String.format(
@@ -866,6 +882,27 @@ private void doAutoCreateAndVerify()
       }
     }
 
+    private void rethrowIfUnavailableRegion(final Exception e) {
+      // Preserve retryable region errors through the analyzer's schema-error wrapper.
+      Throwable current = e;
+      while (current != null) {
+        if (current instanceof IoTDBRuntimeException
+            && ((IoTDBRuntimeException) current).getErrorCode()
+                == TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode()) {
+          throw (IoTDBRuntimeException) current;
+        }
+        if (current instanceof IoTDBException
+            && ((IoTDBException) current).getErrorCode()
+                == TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode()) {
+          throw new IoTDBRuntimeException(
+              current.getMessage(),
+              current,
+              TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode());
+        }
+        current = current.getCause();
+      }
+    }
+
     private void makeSureNoDuplicatedMeasurementsInDevices() throws LoadAnalyzeException {
       for (final Map.Entry<IDeviceID, Set<MeasurementSchema>> entry :
           schemaCache.getDevice2TimeSeries().entrySet()) {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzerDiagnosticsTest.java

@@ -21,17 +21,27 @@
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
 import org.apache.iotdb.db.exception.LoadAnalyzeException;
+import org.apache.iotdb.db.exception.LoadAnalyzeTypeMismatchException;
+import org.apache.iotdb.db.exception.sql.SemanticException;
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.plan.analyze.Analysis;
 import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
 import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaFetcher;
+import org.apache.iotdb.db.queryengine.plan.analyze.schema.SchemaValidator;
 import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement;
 import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter;
 import org.apache.iotdb.rpc.TSStatusCode;
 
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.file.metadata.PlainDeviceID;
+import org.apache.tsfile.read.TsFileSequenceReader;
+import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
@@ -45,16 +55,20 @@
 import org.powermock.modules.junit4.PowerMockRunner;
 
 import java.io.File;
+import java.io.FileOutputStream;
+import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.Optional;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.Mockito.when;
 
 @RunWith(PowerMockRunner.class)
 @PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
-@PrepareForTest(LoadTsFileAnalyzer.class)
+@PrepareForTest({LoadTsFileAnalyzer.class, SchemaValidator.class})
 public class LoadTsFileAnalyzerDiagnosticsTest {
 
   @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
@@ -126,6 +140,153 @@ public void testConversionStatusKeepsMetadataWithoutMutatingOriginal() throws Ex
     Assert.assertTrue(status.getMessage().contains("missing WRITE_DATA privilege"));
   }
 
+  @Test
+  public void testUnavailableSchemaRegionsShouldRemainRetryable() throws Exception {
+    final TSStatus status = analyzeAfterSchemaFailure(TSStatusCode.NO_AVAILABLE_REGION_GROUP);
+
+    Assert.assertEquals(
+        TSStatusCode.LOAD_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode(), status.getCode());
+    Assert.assertTrue(status.getMessage().contains("schema request failed"));
+  }
+
+  @Test
+  public void testPermanentSchemaErrorShouldNotBecomeRetryable() throws Exception {
+    final TSStatus status = analyzeAfterSchemaFailure(TSStatusCode.METADATA_ERROR);
+
+    Assert.assertEquals(TSStatusCode.LOAD_FILE_ERROR.getStatusCode(), status.getCode());
+    Assert.assertTrue(status.getMessage().contains("schema request failed"));
+  }
+
+  @Test
+  public void testWrappedUnavailableSchemaRegionsShouldRemainRetryable() throws Exception {
+    final TSStatus status =
+        analyzeAfterSchemaFailure(
+            new SemanticException(
+                new IoTDBException(
+                    "schema request failed",
+                    TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode())));
+
+    Assert.assertEquals(
+        TSStatusCode.LOAD_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode(), status.getCode());
+    Assert.assertTrue(status.getMessage().contains("schema request failed"));
+  }
+
+  @Test
+  public void testWrappedPermanentSchemaErrorShouldNotBecomeRetryable() throws Exception {
+    final TSStatus status =
+        analyzeAfterSchemaFailure(
+            new SemanticException(
+                new IoTDBException(
+                    "schema request failed", TSStatusCode.METADATA_ERROR.getStatusCode())));
+
+    Assert.assertEquals(TSStatusCode.LOAD_FILE_ERROR.getStatusCode(), status.getCode());
+    Assert.assertTrue(status.getMessage().contains("schema request failed"));
+  }
+
+  @Test
+  public void testWrappedTypeMismatchShouldKeepConversionPriority() throws Exception {
+    final LoadAnalyzeTypeMismatchException mismatch =
+        new LoadAnalyzeTypeMismatchException("schema type mismatch");
+    mismatch.initCause(
+        new IoTDBRuntimeException(
+            "schema request failed", TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode()));
+    when(converter.convertForTreeModel(any()))
+        .thenReturn(Optional.of(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())));
+
+    final TSStatus status = analyzeAfterSchemaFailure(new SemanticException(mismatch));
+
+    // Successful conversion leaves the analysis failure status unset in 1.3.
+    Assert.assertNull(status);
+    Mockito.verify(converter).convertForTreeModel(any());
+  }
+
+  @Test
+  public void testUnavailableRegionsDuringSingleFileAnalysisShouldRetry() throws Exception {
+    final File file = temporaryFolder.newFile("regions-recovering.tsfile");
+    try (final FileOutputStream output = new FileOutputStream(file)) {
+      output.write(1);
+    }
+    PowerMockito.whenNew(TsFileSequenceReader.class)
+        .withArguments(file.getAbsolutePath())
+        .thenThrow(
+            new IoTDBRuntimeException(
+                "region request failed", TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode()));
+
+    try (final LoadTsFileAnalyzer analyzer =
+        new LoadTsFileAnalyzer(
+            LoadTsFileStatement.createUnchecked(file.getPath()),
+            new MPPQueryContext(new QueryId("load_region_recovery")),
+            Mockito.mock(IPartitionFetcher.class),
+            Mockito.mock(ISchemaFetcher.class))) {
+      final Analysis analysis = analyzer.analyzeFileByFile(new Analysis());
+
+      Assert.assertTrue(analysis.isFinishQueryAfterAnalyze());
+      Assert.assertEquals(
+          TSStatusCode.LOAD_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode(),
+          analysis.getFailStatus().getCode());
+      Assert.assertEquals("region request failed", analysis.getFailStatus().getMessage());
+      Assert.assertTrue(file.exists());
+    }
+  }
+
+  private TSStatus analyzeAfterSchemaFailure(final TSStatusCode code) throws Exception {
+    return analyzeAfterSchemaFailure(
+        new IoTDBRuntimeException("schema request failed", code.getStatusCode()));
+  }
+
+  private TSStatus analyzeAfterSchemaFailure(final RuntimeException exception) throws Exception {
+    PowerMockito.mockStatic(SchemaValidator.class);
+    when(SchemaValidator.validate(
+            any(),
+            anyList(),
+            anyList(),
+            anyList(),
+            anyList(),
+            anyList(),
+            anyList(),
+            anyBoolean(),
+            any()))
+        .thenThrow(exception);
+
+    final File file = temporaryFolder.newFile("schema-failure.tsfile");
+    final LoadTsFileStatement statement = LoadTsFileStatement.createUnchecked(file.getPath());
+    statement.setAutoCreateDatabase(false);
+
+    try (final LoadTsFileAnalyzer analyzer =
+        new LoadTsFileAnalyzer(
+            statement,
+            new MPPQueryContext(new QueryId("load_schema_recovery")),
+            Mockito.mock(IPartitionFetcher.class),
+            Mockito.mock(ISchemaFetcher.class))) {
+      final Field verifierField =
+          analyzer.getClass().getDeclaredField("schemaAutoCreatorAndVerifier");
+      verifierField.setAccessible(true);
+      final Object verifier = verifierField.get(analyzer);
+      final Field schemaCacheField = verifier.getClass().getDeclaredField("schemaCache");
+      schemaCacheField.setAccessible(true);
+      final Object schemaCache = schemaCacheField.get(verifier);
+      final IDeviceID device = new PlainDeviceID("root.sg.d");
+      final Method addTimeSeries =
+          schemaCache
+              .getClass()
+              .getDeclaredMethod("addTimeSeries", IDeviceID.class, MeasurementSchema.class);
+      addTimeSeries.setAccessible(true);
+      addTimeSeries.invoke(schemaCache, device, new MeasurementSchema("s", TSDataType.INT64));
+      final Method addIsAlignedCache =
+          schemaCache
+              .getClass()
+              .getDeclaredMethod(
+                  "addIsAlignedCache", IDeviceID.class, boolean.class, boolean.class);
+      addIsAlignedCache.setAccessible(true);
+      addIsAlignedCache.invoke(schemaCache, device, false, true);
+
+      final Analysis analysis = analyzer.analyzeFileByFile(new Analysis());
+      Assert.assertTrue(analysis.isFinishQueryAfterAnalyze());
+      Assert.assertTrue(file.exists());
+      return analysis.getFailStatus();
+    }
+  }
+
   private TSStatus convertAfterAnalysisFailure(final File file) throws Exception {
     try (final LoadTsFileAnalyzer analyzer =
         new LoadTsFileAnalyzer(