Fix import-data and export-data correctness and error reporting

iotdb-client/cli/pom.xml

@@ -54,6 +54,11 @@
             <artifactId>iotdb-antlr</artifactId>
             <version>2.0.11-SNAPSHOT</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.iotdb</groupId>
+            <artifactId>iotdb-relational-grammar</artifactId>
+            <version>2.0.11-SNAPSHOT</version>
+        </dependency>
         <dependency>
             <groupId>org.apache.iotdb</groupId>
             <artifactId>node-commons</artifactId>

iotdb-client/cli/src/main/i18n/en/org/apache/iotdb/cli/i18n/CliMessages.java

@@ -77,6 +77,13 @@ public final class CliMessages {
           "Invalid args: Required values for option table not provided.";
 
   private CliMessages() {}
+  public static final String EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269 =
+      "Failed to insert CSV batch: %s";
+  public static final String MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6 =
+      "[%s] - Unexpected error occurred: %s";
+  public static final String
+      EXCEPTION_SQL_TSFILE_EXPORT_REQUIRES_A_QUERY_FROM_A_SINGLE_TABLE_IN_DATABASE_ARG_C9EB3768 =
+          "SQL/TsFile export requires a query from a single table in database %s.";
   // ---------------------------------------------------------------------------
   // Additional auto-collected messages
   // ---------------------------------------------------------------------------

iotdb-client/cli/src/main/i18n/zh/org/apache/iotdb/cli/i18n/CliMessages.java

@@ -77,6 +77,13 @@ public final class CliMessages {
           "参数无效：未提供 table 选项的必填值。";
 
   private CliMessages() {}
+  public static final String EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269 =
+      "CSV 批次写入失败：%s";
+  public static final String MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6 =
+      "[%s] - 发生意外错误：%s";
+  public static final String
+      EXCEPTION_SQL_TSFILE_EXPORT_REQUIRES_A_QUERY_FROM_A_SINGLE_TABLE_IN_DATABASE_ARG_C9EB3768 =
+          "SQL/TsFile 导出要求查询仅来自数据库 %s 中的一张表。";
   // ---------------------------------------------------------------------------
   // Additional auto-collected messages
   // ---------------------------------------------------------------------------

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractExportData.java

@@ -20,6 +20,7 @@
 package org.apache.iotdb.tool.data;
 
 import org.apache.iotdb.cli.utils.IoTPrinter;
+import org.apache.iotdb.isession.SessionDataSet;
 import org.apache.iotdb.rpc.IoTDBConnectionException;
 import org.apache.iotdb.rpc.RpcUtils;
 import org.apache.iotdb.rpc.StatementExecutionException;
@@ -30,6 +31,8 @@
 import java.time.Instant;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -40,7 +43,31 @@ public abstract class AbstractExportData extends AbstractDataTool {
   public abstract void init()
       throws IoTDBConnectionException, StatementExecutionException, TException;
 
-  public abstract void exportBySql(String sql, int index);
+  public abstract boolean exportBySql(String sql, int index);
+
+  protected static String formatSqlValue(String value, String type) {
+    if (value == null) {
+      return "NULL";
+    }
+    if ("TEXT".equalsIgnoreCase(type) || "STRING".equalsIgnoreCase(type)) {
+      return "'" + value.replace("'", "''") + "'";
+    }
+    return value;
+  }
+
+  protected static List<String> getCsvHeaders(SessionDataSet dataSet, boolean treeModel) {
+    List<String> headers = new ArrayList<>(dataSet.getColumnNames());
+    if (Boolean.TRUE.equals(needDataTypePrinted)) {
+      List<String> types = dataSet.getColumnTypes();
+      for (int i = 0; i < headers.size(); i++) {
+        String name = headers.get(i);
+        if (!"Time".equalsIgnoreCase(name) && !(treeModel && "Device".equals(name))) {
+          headers.set(i, String.format("%s(%s)", name, types.get(i)));
+        }
+      }
+    }
+    return headers;
+  }
 
   protected static void legalCheck(String sql) {
     String aggregatePattern =

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractImportData.java

@@ -19,13 +19,16 @@
 
 package org.apache.iotdb.tool.data;
 
+import org.apache.iotdb.cli.i18n.CliMessages;
 import org.apache.iotdb.cli.utils.IoTPrinter;
 import org.apache.iotdb.rpc.IoTDBConnectionException;
 import org.apache.iotdb.rpc.StatementExecutionException;
 import org.apache.iotdb.tool.common.Constants;
+import org.apache.iotdb.tool.common.ImportTsFileOperation;
 import org.apache.iotdb.tool.tsfile.ImportTsFileScanTool;
 
 import java.io.File;
+import java.io.IOException;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -34,10 +37,39 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public abstract class AbstractImportData extends AbstractDataTool implements Runnable {
 
   private static final IoTPrinter ioTPrinter = new IoTPrinter(System.out);
+  private static final AtomicBoolean importFailed = new AtomicBoolean();
+
+  @FunctionalInterface
+  protected interface DataInsertion {
+    void insert() throws IoTDBConnectionException, StatementExecutionException;
+  }
+
+  /** Retry connection failures without discarding the caller's batch when insertion fails. */
+  protected static void executeWithRetry(DataInsertion insertion, int retryTime) {
+    for (int attempt = 0; ; attempt++) {
+      try {
+        insertion.insert();
+        return;
+      } catch (IoTDBConnectionException e) {
+        if (attempt >= retryTime) {
+          throw new IllegalStateException(
+              String.format(
+                  CliMessages.EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269, e.getMessage()),
+              e);
+        }
+      } catch (StatementExecutionException e) {
+        throw new IllegalStateException(
+            String.format(
+                CliMessages.EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269, e.getMessage()),
+            e);
+      }
+    }
+  }
 
   public abstract void init()
       throws InterruptedException, IoTDBConnectionException, StatementExecutionException;
@@ -65,15 +97,20 @@ public void run() {
         }
       }
     } catch (Exception e) {
+      importFailed.set(true);
+      loadFileFailedNum.increment();
       ioTPrinter.println(
-          String.format("[%s] - Unexpected error occurred: %s", file.getName(), e.getMessage()));
+          String.format(
+              CliMessages.MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6,
+              file == null ? filePath : file.getName(),
+              e.getMessage()));
     }
   }
 
   protected abstract Runnable getAsyncImportRunnable();
 
   protected class ThreadManager {
-    public void asyncImportDataFiles() {
+    public boolean asyncImportDataFiles() {
       List<Thread> list = new ArrayList<>(threadNum);
       for (int i = 0; i < threadNum; i++) {
         Thread thread = new Thread(getAsyncImportRunnable());
@@ -85,16 +122,21 @@ public void asyncImportDataFiles() {
             try {
               thread.join();
             } catch (InterruptedException e) {
+              importFailed.set(true);
               Thread.currentThread().interrupt();
               ioTPrinter.println("ImportData thread join interrupted: " + e.getMessage());
             }
           });
-      ioTPrinter.println(Constants.IMPORT_COMPLETELY);
+      if (!importFailed.get()) {
+        ioTPrinter.println(Constants.IMPORT_COMPLETELY);
+      }
+      return !importFailed.get();
     }
   }
 
-  public static void init(AbstractImportData instance) {
-    instance.new ThreadManager().asyncImportDataFiles();
+  public static boolean init(AbstractImportData instance) {
+    importFailed.set(false);
+    return instance.new ThreadManager().asyncImportDataFiles();
   }
 
   protected abstract void importFromSqlFile(File file);
@@ -111,6 +153,7 @@ protected void processSuccessFile(String file) {
         processingLoadSuccessfulFileSuccessfulNum.increment();
         ioTPrinter.println("Processed success file [ " + file + " ] successfully!");
       } catch (final Exception processSuccessException) {
+        importFailed.set(true);
         ioTPrinter.println(
             "Failed to process success file [ "
                 + file
@@ -120,103 +163,48 @@ protected void processSuccessFile(String file) {
     }
   }
 
-  protected void processingFile(final String file, final boolean isSuccess) {
-    final String relativePath = file.substring(ImportTsFileScanTool.getSourceFullPathLength() + 1);
-    final Path sourcePath = Paths.get(file);
-
+  protected void processingFile(final String file, final boolean isSuccess) throws IOException {
+    final String relativePath = ImportTsFileScanTool.getRelativePath(file);
     final String target =
-        isSuccess
-            ? successDir
-            : failDir + File.separator + relativePath.replace(File.separator, "_");
-    final Path targetPath = Paths.get(target);
-
-    final String RESOURCE = ".resource";
-    Path sourceResourcePath = Paths.get(sourcePath + RESOURCE);
-    sourceResourcePath = Files.exists(sourceResourcePath) ? sourceResourcePath : null;
-    final Path targetResourcePath = Paths.get(target + RESOURCE);
-
-    final String MODS = ".mods";
-    Path sourceModsPath = Paths.get(sourcePath + MODS);
-    sourceModsPath = Files.exists(sourceModsPath) ? sourceModsPath : null;
-    final Path targetModsPath = Paths.get(target + MODS);
-
-    switch (isSuccess ? successOperation : failOperation) {
+        (isSuccess ? successDir : failDir)
+            + File.separator
+            + relativePath.replace(File.separator, "_");
+    final ImportTsFileOperation operation = isSuccess ? successOperation : failOperation;
+    processFile(Paths.get(file), Paths.get(target), operation);
+    for (String suffix : new String[] {".resource", ".mods"}) {
+      Path sidecar = Paths.get(file + suffix);
+      if (Files.exists(sidecar)) {
+        processFile(
+            sidecar,
+            Paths.get(target + suffix),
+            operation == ImportTsFileOperation.HARDLINK ? ImportTsFileOperation.CP : operation);
+      }
+    }
+  }
+
+  private static void processFile(Path source, Path target, ImportTsFileOperation operation)
+      throws IOException {
+    switch (operation) {
       case DELETE:
-        {
-          try {
-            Files.deleteIfExists(sourcePath);
-            if (null != sourceResourcePath) {
-              Files.deleteIfExists(sourceResourcePath);
-            }
-            if (null != sourceModsPath) {
-              Files.deleteIfExists(sourceModsPath);
-            }
-          } catch (final Exception e) {
-            ioTPrinter.println(String.format("Failed to delete file: %s", e.getMessage()));
-          }
-          break;
-        }
+        Files.deleteIfExists(source);
+        break;
       case CP:
-        {
-          try {
-            Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
-            if (null != sourceResourcePath) {
-              Files.copy(
-                  sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING);
-            }
-            if (null != sourceModsPath) {
-              Files.copy(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING);
-            }
-          } catch (final Exception e) {
-            ioTPrinter.println(String.format("Failed to copy file: %s", e.getMessage()));
-          }
-          break;
-        }
+        Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING);
+        break;
       case HARDLINK:
-        {
-          try {
-            Files.createLink(targetPath, sourcePath);
-          } catch (FileAlreadyExistsException e) {
-            ioTPrinter.println("Hardlink already exists: " + e.getMessage());
-          } catch (final Exception e) {
-            try {
-              Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
-            } catch (final Exception copyException) {
-              ioTPrinter.println(
-                  String.format("Failed to copy file: %s", copyException.getMessage()));
-            }
+        try {
+          Files.createLink(target, source);
+        } catch (FileAlreadyExistsException e) {
+          if (!Files.isSameFile(source, target)) {
+            Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING);
           }
-
-          try {
-            if (null != sourceResourcePath) {
-              Files.copy(
-                  sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING);
-            }
-            if (null != sourceModsPath) {
-              Files.copy(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING);
-            }
-          } catch (final Exception e) {
-            ioTPrinter.println(
-                String.format("Failed to copy resource or mods file: %s", e.getMessage()));
-          }
-          break;
+        } catch (IOException | UnsupportedOperationException e) {
+          Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING);
         }
+        break;
       case MV:
-        {
-          try {
-            Files.move(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
-            if (null != sourceResourcePath) {
-              Files.move(
-                  sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING);
-            }
-            if (null != sourceModsPath) {
-              Files.move(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING);
-            }
-          } catch (final Exception e) {
-            ioTPrinter.println(String.format("Failed to move file: %s", e.getMessage()));
-          }
-          break;
-        }
+        Files.move(source, target, StandardCopyOption.REPLACE_EXISTING);
+        break;
       default:
         break;
     }
@@ -232,6 +220,7 @@ protected void processFailFile(final String filePath, final Exception e) {
       }
 
       loadFileFailedNum.increment();
+      importFailed.set(true);
       ioTPrinter.println("Failed to import [ " + filePath + " ] file: " + e.getMessage());
 
       try {

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportData.java

@@ -236,15 +236,21 @@ public static void main(String[] args) {
           ioTPrinter.println(sql);
           String[] values = sql.trim().split(";");
           for (int i = 0; i < values.length; i++) {
-            exportData.exportBySql(values[i], i);
+            if (!exportData.exportBySql(values[i], i)) {
+              exitCode = Constants.CODE_ERROR;
+            }
           }
         } else {
-          exportData.exportBySql(null, 0);
+          if (!exportData.exportBySql(null, 0)) {
+            exitCode = Constants.CODE_ERROR;
+          }
         }
       } else {
         String[] values = queryCommand.trim().split(";");
         for (int i = 0; i < values.length; i++) {
-          exportData.exportBySql(values[i], i);
+          if (!exportData.exportBySql(values[i], i)) {
+            exitCode = Constants.CODE_ERROR;
+          }
         }
       }
     } catch (IOException e) {