This is an automated email from the ASF dual-hosted git repository.
CritasWang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 0acff497175 Fix import-data and export-data correctness and error
reporting (#18828)
0acff497175 is described below
commit 0acff4971750046b6973c654fe16d7f5c3e9a333
Author: Caideyipi <[email protected]>
AuthorDate: Thu Oct 22 12:31:33 2026 +0800
Fix import-data and export-data correctness and error reporting (#18828)
* Fix import-data and export-data correctness and error reporting
* Declare relational grammar dependency for CLI export
* Fix Windows relative source archival test
* Fix archive collisions and typed CSV column names
* Handle import failures per file and share safe TsFile archival
* Fix standalone TsFile import status and parenthesized CSV columns
* Fix ambiguous table CSV imports and custom query validation
---
iotdb-client/cli/pom.xml | 5 +
.../en/org/apache/iotdb/cli/i18n/CliMessages.java | 16 +
.../zh/org/apache/iotdb/cli/i18n/CliMessages.java | 16 +
.../org/apache/iotdb/tool/common/Constants.java | 5 +-
.../iotdb/tool/common/ImportTsFileOperation.java | 59 ++
.../apache/iotdb/tool/data/AbstractExportData.java | 29 +-
.../apache/iotdb/tool/data/AbstractImportData.java | 197 ++--
.../org/apache/iotdb/tool/data/ExportData.java | 12 +-
.../apache/iotdb/tool/data/ExportDataTable.java | 164 ++-
.../org/apache/iotdb/tool/data/ExportDataTree.java | 131 +--
.../org/apache/iotdb/tool/data/ImportData.java | 3 +-
.../apache/iotdb/tool/data/ImportDataTable.java | 86 +-
.../org/apache/iotdb/tool/data/ImportDataTree.java | 88 +-
.../org/apache/iotdb/tool/tsfile/ImportTsFile.java | 9 +-
.../apache/iotdb/tool/tsfile/ImportTsFileBase.java | 126 +--
.../iotdb/tool/tsfile/ImportTsFileScanTool.java | 8 +
.../iotdb/tool/data/DataToolRegressionTest.java | 1042 ++++++++++++++++++++
.../apache/iotdb/tool/tsfile/ImportTsFileTest.java | 205 ++++
scripts/tools/windows/export-data.bat | 4 +-
scripts/tools/windows/export-tsfile.bat | 4 +-
scripts/tools/windows/import-data.bat | 4 +-
scripts/tools/windows/schema/export-schema.bat | 4 +-
scripts/tools/windows/schema/import-schema.bat | 4 +-
23 files changed, 1796 insertions(+), 425 deletions(-)
diff --git a/iotdb-client/cli/pom.xml b/iotdb-client/cli/pom.xml
index 5de247a084a..a1e00508740 100644
--- a/iotdb-client/cli/pom.xml
+++ b/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>
diff --git
a/iotdb-client/cli/src/main/i18n/en/org/apache/iotdb/cli/i18n/CliMessages.java
b/iotdb-client/cli/src/main/i18n/en/org/apache/iotdb/cli/i18n/CliMessages.java
index 0b0f30d970d..ec00a7dc559 100644
---
a/iotdb-client/cli/src/main/i18n/en/org/apache/iotdb/cli/i18n/CliMessages.java
+++
b/iotdb-client/cli/src/main/i18n/en/org/apache/iotdb/cli/i18n/CliMessages.java
@@ -76,7 +76,23 @@ public final class CliMessages {
MESSAGE_INVALID_ARGS_REQUIRED_VALUES_FOR_OPTION_TABLE_NOT_PROVIDED_4BC3FCFA =
"Invalid args: Required values for option table not provided.";
+ // Data import/export
+ 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.";
+ public static final String
+
MESSAGE_QUERY_TO_EXPORT_TABLE_SQL_TSFILE_QUERIES_REQUIRE_A_SINGLE_TABLE_WITHOUT_JOINS_SUBQUERIES_OR_CTES_OVERRIDES_TABLE_START_TIME_AND_END_TIME_OPTIONAL_CC12D400
=
+ "Query to export. Table SQL/TsFile queries require a single table
without joins, subqueries or CTEs. Overrides -table, -start_time and -end_time.
(optional)";
+ public static final String
+ MESSAGE_MULTIPLE_CSV_HEADERS_RESOLVE_TO_THE_SAME_TARGET_COLUMN_99189563 =
+ "Multiple CSV headers resolve to the same target column.";
+
private CliMessages() {}
+
//
---------------------------------------------------------------------------
// Additional auto-collected messages
//
---------------------------------------------------------------------------
diff --git
a/iotdb-client/cli/src/main/i18n/zh/org/apache/iotdb/cli/i18n/CliMessages.java
b/iotdb-client/cli/src/main/i18n/zh/org/apache/iotdb/cli/i18n/CliMessages.java
index a260732ab55..8321a28d08e 100644
---
a/iotdb-client/cli/src/main/i18n/zh/org/apache/iotdb/cli/i18n/CliMessages.java
+++
b/iotdb-client/cli/src/main/i18n/zh/org/apache/iotdb/cli/i18n/CliMessages.java
@@ -76,7 +76,23 @@ public final class CliMessages {
MESSAGE_INVALID_ARGS_REQUIRED_VALUES_FOR_OPTION_TABLE_NOT_PROVIDED_4BC3FCFA =
"参数无效：未提供 table 选项的必填值。";
+ // Data import/export
+ 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 中的一张表。";
+ public static final String
+
MESSAGE_QUERY_TO_EXPORT_TABLE_SQL_TSFILE_QUERIES_REQUIRE_A_SINGLE_TABLE_WITHOUT_JOINS_SUBQUERIES_OR_CTES_OVERRIDES_TABLE_START_TIME_AND_END_TIME_OPTIONAL_CC12D400
=
+ "导出查询。表模型 SQL/TsFile 查询要求仅包含一张表，不支持联接、子查询或 CTE。指定后忽略
-table、-start_time 和 -end_time。（可选）";
+ public static final String
+ MESSAGE_MULTIPLE_CSV_HEADERS_RESOLVE_TO_THE_SAME_TARGET_COLUMN_99189563 =
+ "多个 CSV 表头映射到同一目标列。";
+
private CliMessages() {}
+
//
---------------------------------------------------------------------------
// Additional auto-collected messages
//
---------------------------------------------------------------------------
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/Constants.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/Constants.java
index ed2d96dc860..b6999dea427 100644
--- a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/Constants.java
+++ b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/Constants.java
@@ -19,6 +19,8 @@
package org.apache.iotdb.tool.common;
+import org.apache.iotdb.cli.i18n.CliMessages;
+
import org.apache.tsfile.enums.TSDataType;
import java.text.SimpleDateFormat;
@@ -222,7 +224,8 @@ public class Constants {
public static final String QUERY_COMMAND_NAME = "query";
public static final String QUERY_COMMAND_ARGS_NAME = "query_command";
public static final String QUERY_COMMAND_DESC =
- "The query command that you want to execute.If sql_dialect is table The
'q' parameter is only applicable to export types of CSV, and is not available
for other types.If the '- q' parameter is not empty, then the parameters'
creatTime ',' EndTime 'and' table 'are not effective.(optional)";
+ CliMessages
+
.MESSAGE_QUERY_TO_EXPORT_TABLE_SQL_TSFILE_QUERIES_REQUIRE_A_SINGLE_TABLE_WITHOUT_JOINS_SUBQUERIES_OR_CTES_OVERRIDES_TABLE_START_TIME_AND_END_TIME_OPTIONAL_CC12D400;
public static final String TARGET_FILE_ARGS = "pfn";
public static final String TARGET_FILE_NAME = "prefix_file_name";
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/ImportTsFileOperation.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/ImportTsFileOperation.java
index 3769d855d16..9934efbef6f 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/ImportTsFileOperation.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/ImportTsFileOperation.java
@@ -21,6 +21,15 @@ package org.apache.iotdb.tool.common;
import org.apache.iotdb.cli.utils.IoTPrinter;
+import java.io.IOException;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.LinkOption;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+
public enum ImportTsFileOperation {
NONE,
MV,
@@ -29,6 +38,56 @@ public enum ImportTsFileOperation {
DELETE,
;
+ /** Process a TsFile together with its sidecars without replacing existing
archive files. */
+ public void processFileGroup(Path source, Path target) throws IOException {
+ List<String> suffixes = new ArrayList<>();
+ suffixes.add("");
+ for (String suffix : new String[] {".resource", ".mods"}) {
+ if (Files.exists(Paths.get(source + suffix))) {
+ suffixes.add(suffix);
+ }
+ }
+ if (this == CP || this == HARDLINK || this == MV) {
+ // Check every destination before moving any file in the TsFile/sidecar
group.
+ for (String suffix : suffixes) {
+ Path destination = Paths.get(target + suffix);
+ if (Files.exists(destination, LinkOption.NOFOLLOW_LINKS)) {
+ throw new FileAlreadyExistsException(destination.toString());
+ }
+ }
+ Files.createDirectories(target.getParent());
+ }
+ for (String suffix : suffixes) {
+ ImportTsFileOperation operation = this == HARDLINK && !suffix.isEmpty()
? CP : this;
+ operation.processFile(Paths.get(source + suffix), Paths.get(target +
suffix));
+ }
+ }
+
+ private void processFile(Path source, Path target) throws IOException {
+ switch (this) {
+ case DELETE:
+ Files.deleteIfExists(source);
+ break;
+ case CP:
+ Files.copy(source, target);
+ break;
+ case HARDLINK:
+ try {
+ Files.createLink(target, source);
+ } catch (FileAlreadyExistsException e) {
+ throw e;
+ } catch (IOException | UnsupportedOperationException e) {
+ Files.copy(source, target);
+ }
+ break;
+ case MV:
+ Files.move(source, target);
+ break;
+ default:
+ break;
+ }
+ }
+
public static boolean isValidOperation(String operation) {
return "none".equalsIgnoreCase(operation)
|| "mv".equalsIgnoreCase(operation)
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractExportData.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractExportData.java
index f1d506a72b4..e3c33950d42 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractExportData.java
+++
b/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 org.apache.thrift.TException;
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 =
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractImportData.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractImportData.java
index fa80c8feb20..8aef206c61b 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractImportData.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractImportData.java
@@ -19,25 +19,54 @@
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.nio.file.FileAlreadyExistsException;
-import java.nio.file.Files;
+import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
-import java.nio.file.StandardCopyOption;
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;
@@ -45,35 +74,55 @@ public abstract class AbstractImportData extends
AbstractDataTool implements Run
@Override
public void run() {
String filePath = "";
- File file = null;
- try {
- if (Constants.TSFILE_SUFFIXS.equalsIgnoreCase(fileType)) {
- while ((filePath = ImportTsFileScanTool.pollFromQueue()) != null) {
- file = new File(filePath);
+ if (Constants.TSFILE_SUFFIXS.equalsIgnoreCase(fileType)) {
+ while ((filePath = ImportTsFileScanTool.pollFromQueue()) != null) {
+ File file = new File(filePath);
+ try {
if (file.getName().endsWith(Constants.TSFILE_SUFFIXS)) {
importFromTsFile(file);
}
+ } catch (Exception e) {
+ handleUnexpectedError(file, e);
}
- } else {
- while ((filePath = ImportDataScanTool.pollFromQueue()) != null) {
- file = new File(filePath);
+ }
+ } else {
+ while ((filePath = ImportDataScanTool.pollFromQueue()) != null) {
+ File file = new File(filePath);
+ try {
if (file.getName().endsWith(Constants.SQL_SUFFIXS)) {
importFromSqlFile(file);
} else {
importFromCsvFile(file);
}
+ } catch (Exception e) {
+ handleUnexpectedError(file, e);
}
}
- } catch (Exception e) {
- ioTPrinter.println(
- String.format("[%s] - Unexpected error occurred: %s",
file.getName(), e.getMessage()));
}
}
+ private static void handleUnexpectedError(File file, Exception e) {
+ markFileFailed();
+ ioTPrinter.println(
+ String.format(
+ CliMessages.MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6,
+ file.getName(),
+ e.getMessage()));
+ }
+
+ protected static void markImportFailed() {
+ importFailed.set(true);
+ }
+
+ protected static void markFileFailed() {
+ markImportFailed();
+ loadFileFailedNum.increment();
+ }
+
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 +134,21 @@ public abstract class AbstractImportData extends
AbstractDataTool implements Run
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 +165,7 @@ public abstract class AbstractImportData extends
AbstractDataTool implements Run
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,106 +175,11 @@ public abstract class AbstractImportData extends
AbstractDataTool implements Run
}
}
- protected void processingFile(final String file, final boolean isSuccess) {
- final String relativePath =
file.substring(ImportTsFileScanTool.getSourceFullPathLength() + 1);
- final Path sourcePath = Paths.get(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) {
- 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;
- }
- 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;
- }
- 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 {
- 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;
- }
- 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;
- }
- default:
- break;
- }
+ protected void processingFile(final String file, final boolean isSuccess)
throws IOException {
+ final String relativePath = ImportTsFileScanTool.getRelativePath(file);
+ final Path target = Paths.get(isSuccess ? successDir :
failDir).resolve(relativePath);
+ final ImportTsFileOperation operation = isSuccess ? successOperation :
failOperation;
+ operation.processFileGroup(Paths.get(file), target);
}
protected void processFailFile(final String filePath, final Exception e) {
@@ -232,6 +192,7 @@ public abstract class AbstractImportData extends
AbstractDataTool implements Run
}
loadFileFailedNum.increment();
+ importFailed.set(true);
ioTPrinter.println("Failed to import [ " + filePath + " ] file: " +
e.getMessage());
try {
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportData.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportData.java
index 322f10e059b..63a2fdf45c7 100644
--- a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportData.java
+++ b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportData.java
@@ -236,15 +236,21 @@ public class ExportData extends AbstractDataTool {
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) {
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTable.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTable.java
index 1694e86c9f0..d60b5d559b6 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTable.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTable.java
@@ -19,7 +19,12 @@
package org.apache.iotdb.tool.data;
+import org.apache.iotdb.cli.i18n.CliMessages;
import org.apache.iotdb.cli.utils.IoTPrinter;
+import
org.apache.iotdb.commons.queryengine.plan.relational.sql.parser.CaseInsensitiveStream;
+import org.apache.iotdb.db.relational.grammar.sql.RelationalSqlBaseVisitor;
+import org.apache.iotdb.db.relational.grammar.sql.RelationalSqlLexer;
+import org.apache.iotdb.db.relational.grammar.sql.RelationalSqlParser;
import org.apache.iotdb.isession.ITableSession;
import org.apache.iotdb.isession.SessionDataSet;
import org.apache.iotdb.rpc.IoTDBConnectionException;
@@ -28,6 +33,13 @@ import org.apache.iotdb.rpc.UrlUtils;
import org.apache.iotdb.session.TableSessionBuilder;
import org.apache.iotdb.tool.common.Constants;
+import org.antlr.v4.runtime.BailErrorStrategy;
+import org.antlr.v4.runtime.BaseErrorListener;
+import org.antlr.v4.runtime.CharStreams;
+import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.RecognitionException;
+import org.antlr.v4.runtime.Recognizer;
+import org.antlr.v4.runtime.misc.ParseCancellationException;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.exception.write.WriteProcessException;
import org.apache.tsfile.external.commons.collections4.CollectionUtils;
@@ -51,6 +63,7 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
+import java.util.Locale;
import java.util.Map;
import java.util.stream.Collectors;
@@ -59,6 +72,7 @@ public class ExportDataTable extends AbstractExportData {
private static final IoTPrinter ioTPrinter = new IoTPrinter(System.out);
private static ITableSession tableSession;
private static List<String> tables = new ArrayList<>();
+ private static List<String> allTables = new ArrayList<>();
private static long processedRows;
private static long lastPrintTime;
@@ -85,17 +99,20 @@ public class ExportDataTable extends AbstractExportData {
System.exit(1);
}
sessionDataSet = tableSession.executeQueryStatement("show tables",
timeout);
+ allTables.clear();
while (sessionDataSet.hasNext()) {
- tables.add(sessionDataSet.next().getField(0).getStringValue());
+ allTables.add(sessionDataSet.next().getField(0).getStringValue());
}
- if (CollectionUtils.isEmpty(tables)
- || (ObjectUtils.isNotEmpty(table) && !tables.contains(table))) {
+ boolean filterByTable = StringUtils.isBlank(queryCommand) &&
ObjectUtils.isNotEmpty(table);
+ if (CollectionUtils.isEmpty(allTables) || (filterByTable &&
!allTables.contains(table))) {
ioTPrinter.println(String.format(Constants.TARGET_TABLE_NOT_EXIST_MSG,
table));
System.exit(1);
}
- if (ObjectUtils.isNotEmpty(table)) {
- tables.clear();
+ tables.clear();
+ if (filterByTable) {
tables.add(table);
+ } else {
+ tables.addAll(allTables);
}
if (ObjectUtils.isNotEmpty(sessionDataSet)) {
sessionDataSet.close();
@@ -103,9 +120,11 @@ public class ExportDataTable extends AbstractExportData {
}
@Override
- public void exportBySql(String sql, int index) {
+ public boolean exportBySql(String sql, int index) {
List<String> exportSql = new ArrayList<>();
- if (StringUtils.isNotBlank(sql)) {
+ boolean customQuery = StringUtils.isNotBlank(sql);
+ boolean success = true;
+ if (customQuery) {
if (Constants.SQL_SUFFIXS.equalsIgnoreCase(exportType)
|| Constants.TSFILE_SUFFIXS.equalsIgnoreCase(exportType)) {
legalCheck(sql);
@@ -114,7 +133,7 @@ public class ExportDataTable extends AbstractExportData {
} else {
StringBuilder sqlBuilder;
for (String table : tables) {
- sqlBuilder = new StringBuilder("select * from ").append(table);
+ sqlBuilder = new StringBuilder("select * from
").append(quoteIdentifier(table));
if (StringUtils.isNotBlank(startTime) ||
StringUtils.isNotBlank(endTime)) {
sqlBuilder.append(" where ");
if (StringUtils.isNotBlank(startTime)) {
@@ -131,31 +150,108 @@ public class ExportDataTable extends AbstractExportData {
}
}
for (int i = 0; i < exportSql.size(); i++) {
- String path = targetDirectory + targetFile + i;
- String table = tables.get(i);
- try (SessionDataSet sessionDataSet =
- tableSession.executeQueryStatement(exportSql.get(i), timeout)) {
- if (Constants.SQL_SUFFIXS.equalsIgnoreCase(exportType)) {
- exportToSqlFile(sessionDataSet, table, path);
- } else if (Constants.TSFILE_SUFFIXS.equalsIgnoreCase(exportType)) {
- long start = System.currentTimeMillis();
- boolean isComplete = exportToTsFile(sessionDataSet, path +
".tsfile", table);
- if (isComplete) {
- long end = System.currentTimeMillis();
- ioTPrinter.println("Export completely!cost: " + (end - start) + "
ms.");
+ String path = targetDirectory + targetFile + (customQuery ? index : i);
+ try {
+ String queryTable =
+ customQuery
+ ? (Constants.SQL_SUFFIXS.equalsIgnoreCase(exportType)
+ ||
Constants.TSFILE_SUFFIXS.equalsIgnoreCase(exportType)
+ ? getQueryTable(sql)
+ : null)
+ : tables.get(i);
+ try (SessionDataSet sessionDataSet =
+ tableSession.executeQueryStatement(exportSql.get(i), timeout)) {
+ if (Constants.SQL_SUFFIXS.equalsIgnoreCase(exportType)) {
+ exportToSqlFile(sessionDataSet, queryTable, path);
+ } else if (Constants.TSFILE_SUFFIXS.equalsIgnoreCase(exportType)) {
+ long start = System.currentTimeMillis();
+ boolean isComplete = exportToTsFile(sessionDataSet, path +
".tsfile", queryTable);
+ if (isComplete) {
+ long end = System.currentTimeMillis();
+ ioTPrinter.println("Export completely!cost: " + (end - start) +
" ms.");
+ }
+ } else {
+ exportToCsvFile(sessionDataSet, path);
}
- } else {
- exportToCsvFile(sessionDataSet, path);
+ sessionDataSet.closeOperationHandle();
+ ioTPrinter.println(Constants.EXPORT_COMPLETELY);
}
- sessionDataSet.closeOperationHandle();
- ioTPrinter.println(Constants.EXPORT_COMPLETELY);
} catch (StatementExecutionException
| IoTDBConnectionException
| IOException
| WriteProcessException e) {
ioTPrinter.println("Cannot dump result because: " + e.getMessage());
+ success = false;
}
}
+ return success;
+ }
+
+ private static String getQueryTable(String sql) throws
StatementExecutionException {
+ String message =
+ String.format(
+ CliMessages
+
.EXCEPTION_SQL_TSFILE_EXPORT_REQUIRES_A_QUERY_FROM_A_SINGLE_TABLE_IN_DATABASE_ARG_C9EB3768,
+ database);
+ RelationalSqlLexer lexer =
+ new RelationalSqlLexer(new
CaseInsensitiveStream(CharStreams.fromString(sql)));
+ RelationalSqlParser parser = new RelationalSqlParser(new
CommonTokenStream(lexer));
+ lexer.removeErrorListeners();
+ lexer.addErrorListener(
+ new BaseErrorListener() {
+ @Override
+ public void syntaxError(
+ Recognizer<?, ?> recognizer,
+ Object offendingSymbol,
+ int line,
+ int charPositionInLine,
+ String error,
+ RecognitionException cause) {
+ throw new ParseCancellationException(message);
+ }
+ });
+ parser.removeErrorListeners();
+ parser.setErrorHandler(new BailErrorStrategy());
+ List<RelationalSqlParser.QualifiedNameContext> sources = new ArrayList<>();
+ try {
+ new RelationalSqlBaseVisitor<Void>() {
+ @Override
+ public Void visitTableName(RelationalSqlParser.TableNameContext
context) {
+ sources.add(context.qualifiedName());
+ return null;
+ }
+ }.visit(parser.singleStatement());
+ } catch (ParseCancellationException e) {
+ throw new StatementExecutionException(message);
+ }
+ // Multiple source relations do not have an unambiguous destination table.
+ if (sources.size() != 1) {
+ throw new StatementExecutionException(message);
+ }
+ List<RelationalSqlParser.IdentifierContext> names =
sources.get(0).identifier();
+ if (names.size() > 2
+ || (names.size() == 2 &&
!identifierValue(names.get(0)).equalsIgnoreCase(database))) {
+ throw new StatementExecutionException(message);
+ }
+ String name = identifierValue(names.get(names.size() - 1));
+ if (!allTables.contains(name)) {
+ throw new StatementExecutionException(message);
+ }
+ return name;
+ }
+
+ private static String identifierValue(RelationalSqlParser.IdentifierContext
context) {
+ String value = context.getText();
+ if (value.startsWith("\"")) {
+ value = value.substring(1, value.length() - 1).replace("\"\"", "\"");
+ } else if (value.startsWith("`")) {
+ value = value.substring(1, value.length() - 1).replace("``", "`");
+ }
+ return value.toLowerCase(Locale.ROOT);
+ }
+
+ private static String quoteIdentifier(String name) {
+ return "\"" + name.replace("\"", "\"\"") + "\"";
}
private void exportToSqlFile(SessionDataSet sessionDataSet, String table,
String filePath)
@@ -164,7 +260,14 @@ public class ExportDataTable extends AbstractExportData {
lastPrintTime = 0;
StringBuilder sqlBuilder;
List<String> headers = sessionDataSet.getColumnNames();
- String prevSql = "insert into " + table + "(" + StringUtils.join(headers,
",") + ") values(";
+ String prevSql =
+ "insert into "
+ + quoteIdentifier(table)
+ + "("
+ + headers.stream()
+ .map(ExportDataTable::quoteIdentifier)
+ .collect(Collectors.joining(","))
+ + ") values(";
SessionDataSet.DataIterator iterator = sessionDataSet.iterator();
List<String> columnTypeList = iterator.getColumnTypeList();
int totalColumns = columnTypeList.size();
@@ -185,11 +288,7 @@ public class ExportDataTable extends AbstractExportData {
}
String columnType = columnTypeList.get(curColumnIndex);
String columnValue = iterator.getString(curColumnIndex + 1);
- if (columnType.equals("TEXT") || columnType.equals("STRING")) {
- sqlBuilder.append("'").append(columnValue).append("'");
- } else {
- sqlBuilder.append(columnValue);
- }
+ sqlBuilder.append(formatSqlValue(columnValue, columnType));
}
sqlBuilder.append(");\n");
writer.write(sqlBuilder.toString());
@@ -254,7 +353,7 @@ public class ExportDataTable extends AbstractExportData {
throws IOException, IoTDBConnectionException,
StatementExecutionException {
processedRows = 0;
lastPrintTime = 0;
- List<String> headers = sessionDataSet.getColumnNames();
+ List<String> headers = getCsvHeaders(sessionDataSet, false);
int fileIndex = 0;
SessionDataSet.DataIterator iterator = sessionDataSet.iterator();
List<String> columnTypeList = iterator.getColumnTypeList();
@@ -363,7 +462,8 @@ public class ExportDataTable extends AbstractExportData {
List<String> columnNames, Map<String, Integer> deviceColumnIndices,
String table)
throws IoTDBConnectionException, StatementExecutionException {
List<ColumnSchema> columnSchemas = new ArrayList<>();
- SessionDataSet sessionDataSet =
tableSession.executeQueryStatement("describe " + table);
+ SessionDataSet sessionDataSet =
+ tableSession.executeQueryStatement("describe " +
quoteIdentifier(table));
while (sessionDataSet.hasNext()) {
RowRecord rowRecord = sessionDataSet.next();
final String columnName = rowRecord.getField(0).getStringValue();
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTree.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTree.java
index fd8aeb0cbe3..bab4fa88b92 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTree.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTree.java
@@ -92,7 +92,7 @@ public class ExportDataTree extends AbstractExportData {
}
@Override
- public void exportBySql(String sql, int index) {
+ public boolean exportBySql(String sql, int index) {
if (Constants.SQL_SUFFIXS.equalsIgnoreCase(exportType)
|| Constants.TSFILE_SUFFIXS.equalsIgnoreCase(exportType)) {
legalCheck(sql);
@@ -109,29 +109,17 @@ public class ExportDataTree extends AbstractExportData {
ioTPrinter.println("Export completely!cost: " + (end - start) + "
ms.");
}
} else {
- List<String> headers = new ArrayList<>();
- List<String> names = sessionDataSet.getColumnNames();
- List<String> types = sessionDataSet.getColumnTypes();
- if (Boolean.TRUE.equals(needDataTypePrinted)) {
- for (int i = 0; i < names.size(); i++) {
- if (!"Time".equals(names.get(i)) &&
!"Device".equals(names.get(i))) {
- headers.add(String.format("%s(%s)", names.get(i), types.get(i)));
- } else {
- headers.add(names.get(i));
- }
- }
- } else {
- headers.addAll(names);
- }
exportToCsvFile(sessionDataSet, path);
}
sessionDataSet.closeOperationHandle();
ioTPrinter.println(Constants.EXPORT_COMPLETELY);
+ return true;
} catch (StatementExecutionException
| IoTDBConnectionException
| IOException
| WriteProcessException e) {
ioTPrinter.println("Cannot dump result because: " + e.getMessage());
+ return false;
}
}
@@ -140,7 +128,7 @@ public class ExportDataTree extends AbstractExportData {
processedRows = 0;
lastPrintTime = 0;
- List<String> measurementNames = sessionDataSet.getColumnNames();
+ List<String> measurementNames = new
ArrayList<>(sessionDataSet.getColumnNames());
if (CollectionUtils.isEmpty(measurementNames) || measurementNames.size()
<= 1) {
return;
} else {
@@ -171,12 +159,7 @@ public class ExportDataTree extends AbstractExportData {
}
List<String> values = new ArrayList<>();
for (int index = 2; index < totalColumns; index++) {
- if ("TEXT".equalsIgnoreCase(columnTypeList.get(index))
- || "STRING".equalsIgnoreCase(columnTypeList.get(index))) {
- values.add(String.format("\"%s\"", iterator.getString(index + 1)));
- } else {
- values.add(iterator.getString(index + 1));
- }
+ values.add(formatSqlValue(iterator.getString(index + 1),
columnTypeList.get(index)));
}
long timestamp = iterator.getLong(1);
writer.write(String.format(sqlPrefix, deviceName, timestamp,
String.join(",", values)));
@@ -205,73 +188,67 @@ public class ExportDataTree extends AbstractExportData {
throws IoTDBConnectionException, StatementExecutionException,
IOException {
processedRows = 0;
lastPrintTime = 0;
-
List<String> headers = sessionDataSet.getColumnNames();
- List<String> measurementNames = new ArrayList<>();
- String deviceName = null;
- if (CollectionUtils.isEmpty(headers) || headers.size() <= 1) {
- return;
- } else {
- headers.remove("Time");
- Path path = new Path(headers.get(0), true);
- deviceName = path.getDeviceString();
- for (String header : headers) {
- path = new Path(header, true);
- String meas = path.getMeasurement();
- if (path.getDeviceString().equals(deviceName)) {
- measurementNames.add(meas);
- }
+ Map<String, List<Integer>> deviceColumnIndices = new LinkedHashMap<>();
+ Map<Integer, String> measurements = new HashMap<>();
+ for (int i = 1; i < headers.size(); i++) {
+ Path path = new Path(headers.get(i), true);
+ String device = path.getDeviceString();
+ if (device.startsWith(SYSTEM_DATABASE + ".") ||
device.startsWith(AUDIT_DATABASE + ".")) {
+ continue;
}
+ deviceColumnIndices.computeIfAbsent(device, key -> new
ArrayList<>()).add(i);
+ measurements.put(i, path.getMeasurement());
}
- if (deviceName.startsWith(SYSTEM_DATABASE + ".")
- || deviceName.startsWith(AUDIT_DATABASE + ".")) {
- return;
- }
- String sqlPrefix =
- String.format(
- "INSERT INTO %s(TIMESTAMP,%s) VALUES (%s);\n",
- "%s", String.join(",", measurementNames), "%d,%s");
SessionDataSet.DataIterator iterator = sessionDataSet.iterator();
- List<String> columnTypeList = iterator.getColumnTypeList();
- int totalColumns = measurementNames.size();
+ List<String> columnTypes = iterator.getColumnTypeList();
int fileIndex = 0;
int currentLines = 0;
- String filePathTemplate = filePath + "_%d" + ".sql";
FileWriter writer = null;
- while (iterator.next()) {
- if (writer == null) {
- writer = new FileWriter(String.format(filePathTemplate, fileIndex));
- }
- List<String> values = new ArrayList<>();
- for (int index = 0; index < totalColumns; index++) {
- if ("TEXT".equalsIgnoreCase(columnTypeList.get(index + 1))
- || "STRING".equalsIgnoreCase(columnTypeList.get(index + 1))) {
- values.add(String.format("\"%s\"", iterator.getString(index + 2)));
- } else {
- values.add(iterator.getString(index + 2));
+ try {
+ while (iterator.next()) {
+ for (Map.Entry<String, List<Integer>> entry :
deviceColumnIndices.entrySet()) {
+ List<String> names = new ArrayList<>();
+ List<String> values = new ArrayList<>();
+ for (int index : entry.getValue()) {
+ if (!iterator.isNull(index + 1)) {
+ names.add(measurements.get(index));
+ values.add(formatSqlValue(iterator.getString(index + 1),
columnTypes.get(index)));
+ }
+ }
+ // A timestamp present in another device must not create an empty
record here.
+ if (values.isEmpty()) {
+ continue;
+ }
+ if (writer == null) {
+ writer = new FileWriter(filePath + "_" + fileIndex + ".sql");
+ }
+ writer.write(
+ String.format(
+ "INSERT INTO %s(TIMESTAMP,%s) VALUES (%d,%s);\n",
+ entry.getKey(),
+ String.join(",", names),
+ iterator.getLong(1),
+ String.join(",", values)));
+ if (++currentLines >= linesPerFile) {
+ writer.close();
+ writer = null;
+ fileIndex++;
+ currentLines = 0;
+ }
+ }
+ processedRows++;
+ if (System.currentTimeMillis() - lastPrintTime > updateTimeInterval) {
+ ioTPrinter.printf(Constants.PROCESSED_PROGRESS, processedRows);
+ lastPrintTime = System.currentTimeMillis();
}
}
- long timestamp = iterator.getLong(1);
- writer.write(String.format(sqlPrefix, deviceName, timestamp,
String.join(",", values)));
- processedRows += 1;
- currentLines += 1;
- if (System.currentTimeMillis() - lastPrintTime > updateTimeInterval) {
- ioTPrinter.printf(Constants.PROCESSED_PROGRESS, processedRows);
- lastPrintTime = System.currentTimeMillis();
- }
- if (currentLines >= linesPerFile) {
- writer.flush();
+ } finally {
+ if (writer != null) {
writer.close();
- fileIndex += 1;
- writer = null;
- currentLines = 0;
}
}
- if (writer != null) {
- writer.flush();
- writer.close();
- }
ioTPrinter.print("\n");
}
@@ -328,7 +305,7 @@ public class ExportDataTree extends AbstractExportData {
throws IOException, IoTDBConnectionException,
StatementExecutionException {
processedRows = 0;
lastPrintTime = 0;
- List<String> headers = sessionDataSet.getColumnNames();
+ List<String> headers = getCsvHeaders(sessionDataSet, true);
int fileIndex = 0;
SessionDataSet.DataIterator iterator = sessionDataSet.iterator();
List<String> columnTypeList = iterator.getColumnTypeList();
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportData.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportData.java
index 9525c171606..8686352cb5a 100644
--- a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportData.java
+++ b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportData.java
@@ -438,8 +438,7 @@ public class ImportData extends AbstractDataTool {
ioTPrinter.println(String.format("Source file or directory %s does not
exist", targetPath));
System.exit(Constants.CODE_ERROR);
}
- AbstractImportData.init(importData);
- return Constants.CODE_OK;
+ return AbstractImportData.init(importData) ? Constants.CODE_OK :
Constants.CODE_ERROR;
} catch (InterruptedException e) {
ioTPrinter.println(String.format("Import tsfile fail: %s",
e.getMessage()));
Thread.currentThread().interrupt();
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java
index f65b559a941..d866efbd242 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java
@@ -54,6 +54,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
+import java.util.Locale;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@@ -189,6 +190,7 @@ public class ImportDataTable extends AbstractImportData {
@SuppressWarnings("java:S2259")
protected void importFromSqlFile(File file) {
ArrayList<List<Object>> failedRecords = new ArrayList<>();
+ boolean fileFailed = false;
String failedFilePath;
if (failedFileDirectory == null) {
failedFilePath = file.getAbsolutePath() + ".failed";
@@ -210,18 +212,25 @@ public class ImportDataTable extends AbstractImportData {
"The extracted database '%s' in SQL statement does not
match the target database '%s'",
dbName, database));
failedRecords.add(Collections.singletonList(sql));
+ fileFailed = true;
continue;
}
session.executeNonQueryStatement(sql);
} catch (IoTDBConnectionException | StatementExecutionException e) {
+ fileFailed = true;
ioTPrinter.println(e.getMessage());
failedRecords.add(Collections.singletonList(sql));
}
}
- processSuccessFile();
} catch (IOException e) {
+ fileFailed = true;
ioTPrinter.println("SQL file read exception because: " + e.getMessage());
}
+ if (fileFailed) {
+ markFileFailed();
+ } else {
+ processSuccessFile();
+ }
if (!failedRecords.isEmpty()) {
FileWriter writer = null;
try {
@@ -230,6 +239,7 @@ public class ImportDataTable extends AbstractImportData {
writer.write(failedRecord.get(0).toString() + "\n");
}
} catch (IOException e) {
+ markImportFailed();
ioTPrinter.println("Cannot dump fail result because: " +
e.getMessage());
} finally {
if (ObjectUtils.isNotEmpty(writer)) {
@@ -257,15 +267,16 @@ public class ImportDataTable extends AbstractImportData {
protected void importFromCsvFile(File file) {
if (file.getName().endsWith(Constants.CSV_SUFFIXS)
|| file.getName().endsWith(Constants.TXT_SUFFIXS)) {
- try {
- CSVParser csvRecords = readCsvFile(file.getAbsolutePath());
+ try (CSVParser csvRecords = readCsvFile(file.getAbsolutePath())) {
List<String> headerNames = csvRecords.getHeaderNames();
Stream<CSVRecord> records = csvRecords.stream();
if (headerNames.isEmpty()) {
+ markFileFailed();
ioTPrinter.println("Empty file!");
return;
}
if
(!timeColumn.toLowerCase().equalsIgnoreCase(filterBomHeader(headerNames.get(0))))
{
+ markFileFailed();
ioTPrinter.println("The first field of header must be `time`!");
return;
}
@@ -275,21 +286,29 @@ public class ImportDataTable extends AbstractImportData {
} else {
failedFilePath = failedFileDirectory + file.getName() + ".failed";
}
- writeData(headerNames, records.collect(Collectors.toList()),
failedFilePath);
- processSuccessFile();
+ if (writeData(headerNames, records.collect(Collectors.toList()),
failedFilePath)) {
+ processSuccessFile();
+ }
} catch (IOException e) {
+ markFileFailed();
ioTPrinter.println("CSV file read exception because: " +
e.getMessage());
}
} else {
+ markFileFailed();
ioTPrinter.println("The file name must end with \"csv\" or \"txt\"!");
}
}
- protected void writeData(
+ protected boolean writeData(
List<String> headerNames, List<CSVRecord> records, String
failedFilePath) {
Map<String, TSDataType> headerTypeMap = new HashMap<>();
Map<String, String> headerNameMap = new HashMap<>();
- parseHeaders(headerNames, headerTypeMap, headerNameMap);
+ if (!parseHeaders(headerNames, headerTypeMap, headerNameMap)) {
+ markFileFailed();
+ ioTPrinter.println(
+
CliMessages.MESSAGE_MULTIPLE_CSV_HEADERS_RESOLVE_TO_THE_SAME_TARGET_COLUMN_99189563);
+ return false;
+ }
queryType(headerTypeMap);
List<List<Object>> failedRecords = new ArrayList<>();
List<IMeasurementSchema> schemas = new ArrayList<>();
@@ -367,41 +386,60 @@ public class ImportDataTable extends AbstractImportData {
if (!failedRecords.isEmpty()) {
writeFailedLinesFile(headerNames, failedFilePath, failedRecords);
}
+ return true;
}
private static void writeAndEmptyDataSet(Tablet tablet, int retryTime) {
- try (ITableSession session = sessionPool.getSession()) {
- session.insert(tablet);
- } catch (IoTDBConnectionException e) {
- if (retryTime > 0) {
- writeAndEmptyDataSet(tablet, --retryTime);
- }
- } catch (StatementExecutionException e) {
- ioTPrinter.println(Constants.INSERT_CSV_MEET_ERROR_MSG + e.getMessage());
- System.exit(1);
- }
+ executeWithRetry(
+ () -> {
+ try (ITableSession session = sessionPool.getSession()) {
+ session.insert(tablet);
+ }
+ },
+ retryTime);
}
- private void parseHeaders(
+ private boolean parseHeaders(
List<String> headerNames,
Map<String, TSDataType> headerTypeMap,
Map<String, String> headerNameMap) {
- String regex = "(?<=\\()\\S+(?=\\))";
+ String regex = "\\(([^()]+)\\)$";
Pattern pattern = Pattern.compile(regex);
+ List<String> unmatchedHeaders = new ArrayList<>();
for (String headerName : headerNames) {
- headerName = headerName.toLowerCase();
+ String normalizedHeaderName = headerName.toLowerCase(Locale.ROOT);
+ // A complete table column name takes precedence over a possible type
suffix.
+ if (dataTypes.containsKey(normalizedHeaderName)) {
+ if (headerNameMap.containsKey(normalizedHeaderName)) {
+ return false;
+ }
+ headerNameMap.put(normalizedHeaderName, headerName);
+ headerTypeMap.put(normalizedHeaderName,
dataTypes.get(normalizedHeaderName));
+ } else {
+ unmatchedHeaders.add(headerName);
+ }
+ }
+ for (String headerName : unmatchedHeaders) {
+ String normalizedHeaderName = headerName.toLowerCase(Locale.ROOT);
Matcher matcher = pattern.matcher(headerName);
String type;
String headerNameWithoutType;
if (matcher.find()) {
- type = matcher.group();
- headerNameWithoutType = headerName.replace("(" + type + ")",
"").replaceAll("\\s+", "");
+ type = matcher.group(1);
+ headerNameWithoutType = headerName.substring(0,
matcher.start()).toLowerCase(Locale.ROOT);
+ if (headerNameMap.containsKey(headerNameWithoutType)) {
+ return false;
+ }
headerNameMap.put(headerNameWithoutType, headerName);
- headerTypeMap.put(headerNameWithoutType, getType(type));
+ headerTypeMap.put(headerNameWithoutType,
getType(type.toUpperCase(Locale.ROOT)));
} else {
- headerNameMap.put(headerName, headerName);
+ if (headerNameMap.containsKey(normalizedHeaderName)) {
+ return false;
+ }
+ headerNameMap.put(normalizedHeaderName, headerName);
}
}
+ return true;
}
private void queryType(Map<String, TSDataType> dataType) {
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTree.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTree.java
index c976e847740..35641d3d501 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTree.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTree.java
@@ -103,6 +103,7 @@ public class ImportDataTree extends AbstractImportData {
@SuppressWarnings("java:S2259")
protected void importFromSqlFile(File file) {
ArrayList<List<Object>> failedRecords = new ArrayList<>();
+ boolean fileFailed = false;
String failedFilePath;
if (failedFileDirectory == null) {
failedFilePath = file.getAbsolutePath() + ".failed";
@@ -115,19 +116,26 @@ public class ImportDataTree extends AbstractImportData {
try {
sessionPool.executeNonQueryStatement(sql);
} catch (IoTDBConnectionException | StatementExecutionException e) {
+ fileFailed = true;
failedRecords.add(Collections.singletonList(sql));
}
}
- processSuccessFile(null);
} catch (IOException e) {
+ fileFailed = true;
ioTPrinter.println("SQL file read exception because: " + e.getMessage());
}
+ if (fileFailed) {
+ markFileFailed();
+ } else {
+ processSuccessFile(null);
+ }
if (!failedRecords.isEmpty()) {
try (FileWriter writer = new FileWriter(failedFilePath)) {
for (List<Object> failedRecord : failedRecords) {
writer.write(failedRecord.get(0).toString() + "\n");
}
} catch (IOException e) {
+ markImportFailed();
ioTPrinter.println("Cannot dump fail result because: " +
e.getMessage());
}
}
@@ -146,15 +154,16 @@ public class ImportDataTree extends AbstractImportData {
protected void importFromCsvFile(File file) {
if (file.getName().endsWith(Constants.CSV_SUFFIXS)
|| file.getName().endsWith(Constants.TXT_SUFFIXS)) {
- try {
- CSVParser csvRecords = readCsvFile(file.getAbsolutePath());
+ try (CSVParser csvRecords = readCsvFile(file.getAbsolutePath())) {
List<String> headerNames = csvRecords.getHeaderNames();
Stream<CSVRecord> records = csvRecords.stream();
if (headerNames.isEmpty()) {
+ markFileFailed();
ioTPrinter.println("Empty file!");
return;
}
if (!timeColumn.equalsIgnoreCase(filterBomHeader(headerNames.get(0))))
{
+ markFileFailed();
ioTPrinter.println("The first field of header must be `Time`!");
return;
}
@@ -171,9 +180,11 @@ public class ImportDataTree extends AbstractImportData {
}
processSuccessFile(null);
} catch (IOException | IllegalPathException e) {
+ markFileFailed();
ioTPrinter.println("CSV file read exception because: " +
e.getMessage());
}
} else {
+ markFileFailed();
ioTPrinter.println("The file name must end with \"csv\" or \"txt\"!");
}
}
@@ -401,26 +412,21 @@ public class ImportDataTree extends AbstractImportData {
List<List<Object>> valuesList,
List<List<String>> measurementsList,
int retryTime) {
- try {
- if (Boolean.FALSE.equals(aligned)) {
- sessionPool.insertRecordsOfOneDevice(
- device, times, measurementsList, typesList, valuesList);
- } else {
- sessionPool.insertAlignedRecordsOfOneDevice(
- device, times, measurementsList, typesList, valuesList);
- }
- } catch (IoTDBConnectionException e) {
- if (retryTime > 0) {
- writeAndEmptyDataSet(device, times, typesList, valuesList,
measurementsList, --retryTime);
- }
- } catch (StatementExecutionException e) {
- ioTPrinter.println(Constants.INSERT_CSV_MEET_ERROR_MSG + e.getMessage());
- } finally {
- times.clear();
- typesList.clear();
- valuesList.clear();
- measurementsList.clear();
- }
+ executeWithRetry(
+ () -> {
+ if (Boolean.FALSE.equals(aligned)) {
+ sessionPool.insertRecordsOfOneDevice(
+ device, times, measurementsList, typesList, valuesList);
+ } else {
+ sessionPool.insertAlignedRecordsOfOneDevice(
+ device, times, measurementsList, typesList, valuesList);
+ }
+ },
+ retryTime);
+ times.clear();
+ typesList.clear();
+ valuesList.clear();
+ measurementsList.clear();
}
private static void writeAndEmptyDataSet(
@@ -430,27 +436,21 @@ public class ImportDataTree extends AbstractImportData {
List<List<Object>> valuesList,
List<List<String>> measurementsList,
int retryTime) {
- try {
- if (Boolean.FALSE.equals(aligned)) {
- sessionPool.insertRecords(deviceIds, times, measurementsList,
typesList, valuesList);
- } else {
- sessionPool.insertAlignedRecords(deviceIds, times, measurementsList,
typesList, valuesList);
- }
- } catch (IoTDBConnectionException e) {
- if (retryTime > 0) {
- writeAndEmptyDataSet(
- deviceIds, times, typesList, valuesList, measurementsList,
--retryTime);
- }
- } catch (StatementExecutionException e) {
- ioTPrinter.println(Constants.INSERT_CSV_MEET_ERROR_MSG + e.getMessage());
- System.exit(1);
- } finally {
- deviceIds.clear();
- times.clear();
- typesList.clear();
- valuesList.clear();
- measurementsList.clear();
- }
+ executeWithRetry(
+ () -> {
+ if (Boolean.FALSE.equals(aligned)) {
+ sessionPool.insertRecords(deviceIds, times, measurementsList,
typesList, valuesList);
+ } else {
+ sessionPool.insertAlignedRecords(
+ deviceIds, times, measurementsList, typesList, valuesList);
+ }
+ },
+ retryTime);
+ deviceIds.clear();
+ times.clear();
+ typesList.clear();
+ valuesList.clear();
+ measurementsList.clear();
}
private static void queryType(
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFile.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFile.java
index 97ff8d164f3..3b9e36a7965 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFile.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFile.java
@@ -370,8 +370,7 @@ public class ImportTsFile extends AbstractTsFileTool {
ImportTsFileScanTool.traverseAndCollectFiles();
ImportTsFileScanTool.addNoResourceOrModsToQueue();
- asyncImportTsFiles();
- return CODE_OK;
+ return asyncImportTsFiles();
} catch (InterruptedException e) {
IOT_PRINTER.println(String.format("Import tsfile fail: %s",
e.getMessage()));
Thread.currentThread().interrupt();
@@ -413,7 +412,8 @@ public class ImportTsFile extends AbstractTsFileTool {
successDir, successOperation, failDir, failOperation);
}
- public static void asyncImportTsFiles() {
+ public static int asyncImportTsFiles() {
+ ImportTsFileBase.resetCounters();
final List<Thread> list = new ArrayList<>(threadNum);
for (int i = 0; i < threadNum; i++) {
final Thread thread =
@@ -433,5 +433,8 @@ public class ImportTsFile extends AbstractTsFileTool {
IOT_PRINTER.println("ImportTsFile thread join interrupted: " +
e.getMessage());
}
});
+ return Thread.currentThread().isInterrupted() ||
ImportTsFileBase.hasFailures()
+ ? CODE_ERROR
+ : CODE_OK;
}
}
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileBase.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileBase.java
index f1fd3bf2d30..634eda9741d 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileBase.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileBase.java
@@ -22,12 +22,9 @@ package org.apache.iotdb.tool.tsfile;
import org.apache.iotdb.cli.utils.IoTPrinter;
import org.apache.iotdb.tool.common.ImportTsFileOperation;
-import java.io.File;
-import java.nio.file.FileAlreadyExistsException;
-import java.nio.file.Files;
+import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
-import java.nio.file.StandardCopyOption;
import java.util.Objects;
import java.util.concurrent.atomic.LongAdder;
@@ -101,107 +98,26 @@ public abstract class ImportTsFileBase implements
Runnable {
}
}
- public static void processingFile(final String filePath, final boolean
isSuccess) {
- final String relativePath =
- filePath.substring(ImportTsFileScanTool.getSourceFullPathLength() + 1);
- final Path sourcePath = Paths.get(filePath);
-
- 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) {
- case DELETE:
- {
- try {
- Files.deleteIfExists(sourcePath);
- if (null != sourceResourcePath) {
- Files.deleteIfExists(sourceResourcePath);
- }
- if (null != sourceModsPath) {
- Files.deleteIfExists(sourceModsPath);
- }
- } catch (final Exception e) {
- IOT_PRINTER.println(String.format("Failed to delete file: %s",
e.getMessage()));
- }
- 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) {
- IOT_PRINTER.println(String.format("Failed to copy file: %s",
e.getMessage()));
- }
- break;
- }
- case HARDLINK:
- {
- try {
- Files.createLink(targetPath, sourcePath);
- } catch (FileAlreadyExistsException e) {
- IOT_PRINTER.println("Hardlink already exists: " + e.getMessage());
- } catch (final Exception e) {
- try {
- Files.copy(sourcePath, targetPath,
StandardCopyOption.REPLACE_EXISTING);
- } catch (final Exception copyException) {
- IOT_PRINTER.println(
- String.format("Failed to copy file: %s",
copyException.getMessage()));
- }
- }
-
- 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) {
- IOT_PRINTER.println(
- String.format("Failed to copy resource or mods file: %s",
e.getMessage()));
- }
- 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) {
- IOT_PRINTER.println(String.format("Failed to move file: %s",
e.getMessage()));
- }
- break;
- }
- default:
- break;
- }
+ public static void processingFile(final String filePath, final boolean
isSuccess)
+ throws IOException {
+ final Path target =
+ Paths.get(isSuccess ? successDir : failDir)
+ .resolve(ImportTsFileScanTool.getRelativePath(filePath));
+ final ImportTsFileOperation operation = isSuccess ? successOperation :
failOperation;
+ operation.processFileGroup(Paths.get(filePath), target);
+ }
+
+ protected static void resetCounters() {
+ loadFileSuccessfulNum.reset();
+ loadFileFailedNum.reset();
+ processingLoadSuccessfulFileSuccessfulNum.reset();
+ processingLoadFailedFileSuccessfulNum.reset();
+ }
+
+ protected static boolean hasFailures() {
+ return loadFileFailedNum.sum() > 0
+ || loadFileSuccessfulNum.sum() !=
processingLoadSuccessfulFileSuccessfulNum.sum()
+ || ImportTsFileScanTool.getTsFileQueueSize() > 0;
}
protected static void printResult(final long startTime) {
diff --git
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileScanTool.java
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileScanTool.java
index 4891f425486..f31e7bc72e9 100644
---
a/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileScanTool.java
+++
b/iotdb-client/cli/src/main/java/org/apache/iotdb/tool/tsfile/ImportTsFileScanTool.java
@@ -20,6 +20,8 @@
package org.apache.iotdb.tool.tsfile;
import java.io.File;
+import java.nio.file.Path;
+import java.nio.file.Paths;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
@@ -91,6 +93,12 @@ public class ImportTsFileScanTool {
: new File(ImportTsFileScanTool.sourceFullPath).getParent().length();
}
+ public static String getRelativePath(String file) {
+ Path source = Paths.get(sourceFullPath).toAbsolutePath().normalize();
+ Path directory = source.toFile().isDirectory() ? source :
source.getParent();
+ return
directory.relativize(Paths.get(file).toAbsolutePath().normalize()).toString();
+ }
+
public static int getTsFileQueueSize() {
return ImportTsFileScanTool.tsfileQueue.size();
}
diff --git
a/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/data/DataToolRegressionTest.java
b/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/data/DataToolRegressionTest.java
new file mode 100644
index 00000000000..ff4a59ada91
--- /dev/null
+++
b/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/data/DataToolRegressionTest.java
@@ -0,0 +1,1042 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.tool.data;
+
+import org.apache.iotdb.isession.ITableSession;
+import org.apache.iotdb.isession.SessionDataSet;
+import org.apache.iotdb.isession.pool.ITableSessionPool;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.session.Session;
+import org.apache.iotdb.session.pool.SessionPool;
+import org.apache.iotdb.tool.common.ImportTsFileOperation;
+import org.apache.iotdb.tool.tsfile.ImportTsFileBase;
+import org.apache.iotdb.tool.tsfile.ImportTsFileScanTool;
+
+import org.apache.commons.csv.CSVParser;
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.write.record.Tablet;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class DataToolRegressionTest extends AbstractDataTool {
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private final Map<Field, Object> originalFields = new HashMap<>();
+ private Path directory;
+
+ @Before
+ public void setUp() throws Exception {
+ for (Class<?> owner :
+ Arrays.asList(
+ AbstractDataTool.class,
+ ExportDataTree.class,
+ ExportDataTable.class,
+ ImportDataTree.class,
+ ImportDataTable.class,
+ ImportTsFileBase.class,
+ ImportTsFileScanTool.class)) {
+ for (Field field : owner.getDeclaredFields()) {
+ if (Modifier.isStatic(field.getModifiers()) &&
!Modifier.isFinal(field.getModifiers())) {
+ field.setAccessible(true);
+ originalFields.put(field, field.get(null));
+ }
+ }
+ }
+ directory = temporaryFolder.getRoot().toPath();
+ targetDirectory = directory + File.separator;
+ targetFile = "dump";
+ exportType = "sql";
+ fileType = "csv";
+ timeFormat = "timestamp";
+ timeColumn = "Time";
+ timestampPrecision = "ms";
+ database = "db";
+ table = "t";
+ aligned = false;
+ linesPerFile = 10000;
+ batchPointSize = 100;
+ threadNum = 1;
+ needDataTypePrinted = false;
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ for (Map.Entry<Field, Object> entry : originalFields.entrySet()) {
+ entry.getKey().set(null, entry.getValue());
+ }
+ }
+
+ @Test
+ public void testMultiDeviceSqlPreservesColumnMappingAndSplitsFiles() throws
Exception {
+ linesPerFile = 1;
+ mockTree(
+ dataSet(
+ List.of("Time", "root.a.d1.s1", "root.a.d2.s2", "root.a.d1.s3"),
+ List.of("INT64", "INT32", "INT32", "INT32"),
+ new Object[] {1L, 11, 22, 33},
+ new Object[] {2L, null, 44, null}));
+ assertTrue(new ExportDataTree().exportBySql("select * from root.a.*", 0));
+ assertEquals("INSERT INTO root.a.d1(TIMESTAMP,s1,s3) VALUES (1,11,33);\n",
read("dump0_0.sql"));
+ assertEquals("INSERT INTO root.a.d2(TIMESTAMP,s2) VALUES (1,22);\n",
read("dump0_1.sql"));
+ assertEquals("INSERT INTO root.a.d2(TIMESTAMP,s2) VALUES (2,44);\n",
read("dump0_2.sql"));
+ assertFalse(Files.exists(directory.resolve("dump0_3.sql")));
+ }
+
+ @Test
+ public void testCustomTableQueriesUseDistinctFilesAndActualTables() throws
Exception {
+ ITableSession session = mockTable("table_a", "table_b");
+ SessionDataSet first =
+ dataSet(List.of("time", "v"), List.of("TIMESTAMP", "INT32"), new
Object[] {1L, 11});
+ SessionDataSet second =
+ dataSet(List.of("time", "v"), List.of("TIMESTAMP", "INT32"), new
Object[] {2L, 22});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(first, second);
+ ExportDataTable tool = new ExportDataTable();
+ assertTrue(tool.exportBySql("select * from table_a", 0));
+ assertTrue(tool.exportBySql("SELECT * FROM \"db\".\"table_b\" WHERE tag =
'from table_a'", 1));
+ assertEquals("insert into \"table_a\"(\"time\",\"v\") values(1,11);\n",
read("dump0_0.sql"));
+ assertEquals("insert into \"table_b\"(\"time\",\"v\") values(2,22);\n",
read("dump1_0.sql"));
+ }
+
+ @Test
+ public void testCustomTableQueryOverridesTableFilter() throws Exception {
+ ITableSession session = mockTable("table_a", "table_b");
+ table = "table_a";
+ set(ExportDataTable.class, "tables", new ArrayList<>(List.of("table_a")));
+ SessionDataSet result =
+ dataSet(List.of("time", "v"), List.of("TIMESTAMP", "INT32"), new
Object[] {1L, 11});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+
+ assertTrue(new ExportDataTable().exportBySql("select * from table_b", 0));
+ assertEquals("insert into \"table_b\"(\"time\",\"v\") values(1,11);\n",
read("dump0_0.sql"));
+ }
+
+ @Test
+ public void testDefaultTableExportStillUsesTableFilter() throws Exception {
+ ITableSession session = mockTable("table_a", "table_b");
+ table = "table_a";
+ set(ExportDataTable.class, "tables", new ArrayList<>(List.of("table_a")));
+ SessionDataSet result =
+ dataSet(List.of("time", "v"), List.of("TIMESTAMP", "INT32"), new
Object[] {1L, 11});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+
+ assertTrue(new ExportDataTable().exportBySql(null, 0));
+ verify(session).executeQueryStatement("select * from \"table_a\"",
timeout);
+ assertEquals("insert into \"table_a\"(\"time\",\"v\") values(1,11);\n",
read("dump0_0.sql"));
+ assertFalse(Files.exists(directory.resolve("dump1_0.sql")));
+ }
+
+ @Test
+ public void testQuotedTableAndColumnNames() throws Exception {
+ ITableSession session = mockTable("odd\"name");
+ SessionDataSet result =
+ dataSet(List.of("time", "select"), List.of("TIMESTAMP", "INT32"), new
Object[] {1L, 11});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+ assertTrue(new ExportDataTable().exportBySql("select * from
\"ODD\"\"NAME\"", 0));
+ assertEquals(
+ "insert into \"odd\"\"name\"(\"time\",\"select\") values(1,11);\n",
read("dump0_0.sql"));
+ }
+
+ @Test
+ public void testAmbiguousAndCrossDatabaseTableQueriesFailBeforeWriting()
throws Exception {
+ ITableSession session = mockTable("table_a", "table_b");
+ ExportDataTable tool = new ExportDataTable();
+ assertFalse(
+ tool.exportBySql("select * from table_a join table_b on table_a.time =
table_b.time", 0));
+ assertFalse(tool.exportBySql("select * from other_db.table_b", 1));
+ assertFalse(tool.exportBySql("select * from", 2));
+ assertFalse(tool.exportBySql("select * from table_b #invalid_character",
3));
+ verify(session, never()).executeQueryStatement(anyString(), anyLong());
+ assertFalse(Files.exists(directory.resolve("dump0_0.sql")));
+ }
+
+ @Test
+ public void testTableSqlEscapesStringsAndPreservesNull() throws Exception {
+ ITableSession session = mockTable("t");
+ SessionDataSet result =
+ dataSet(
+ List.of("time", "tag", "value"),
+ List.of("TIMESTAMP", "STRING", "TEXT"),
+ new Object[] {1L, "O'Brien\\path", null});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+ assertTrue(new ExportDataTable().exportBySql("select * from t", 0));
+ assertEquals(
+ "insert into \"t\"(\"time\",\"tag\",\"value\")
values(1,'O''Brien\\path',NULL);\n",
+ read("dump0_0.sql"));
+ }
+
+ @Test
+ public void testAlignedTreeSqlEscapesStringsAndPreservesNull() throws
Exception {
+ mockTree(
+ dataSet(
+ List.of("Time", "Device", "s1", "s2"),
+ List.of("INT64", "TEXT", "TEXT", "STRING"),
+ new Object[] {1L, "root.a.d1", "a'b\"\\c", null}));
+ assertTrue(new ExportDataTree().exportBySql("select * from root.a.* align
by device", 0));
+ assertEquals(
+ "INSERT INTO root.a.d1(TIMESTAMP,s1,s2) ALIGNED VALUES
(1,'a''b\"\\c',NULL);\n",
+ read("dump0_0.sql"));
+ }
+
+ @Test
+ public void testTreeCsvTypeHeadersInEverySplitFile() throws Exception {
+ exportType = "csv";
+ linesPerFile = 1;
+ needDataTypePrinted = true;
+ mockTree(
+ dataSet(
+ List.of("Time", "Device", "s1"),
+ List.of("INT64", "TEXT", "DOUBLE"),
+ new Object[] {1L, "root.a.d1", 1.5},
+ new Object[] {2L, "root.a.d1", 2.5}));
+ assertTrue(new ExportDataTree().exportBySql("select * from root.a.* align
by device", 0));
+ assertEquals("Time,Device,s1(DOUBLE)", firstLine("dump0_0.csv"));
+ assertEquals("Time,Device,s1(DOUBLE)", firstLine("dump0_1.csv"));
+ needDataTypePrinted = false;
+ mockTree(
+ dataSet(
+ List.of("Time", "root.a.d1.s1"), List.of("INT64", "DOUBLE"), new
Object[] {1L, 1.5}));
+ assertTrue(new ExportDataTree().exportBySql("select * from root.a.d1", 1));
+ assertEquals("Time,root.a.d1.s1", firstLine("dump1_0.csv"));
+ }
+
+ @Test
+ public void testTableCsvTypeHeadersAndMultipleQueryFiles() throws Exception {
+ exportType = "csv";
+ needDataTypePrinted = true;
+ ITableSession session = mockTable("t");
+ SessionDataSet first =
+ dataSet(
+ List.of("time", "device", "v"),
+ List.of("TIMESTAMP", "STRING", "DOUBLE"),
+ new Object[] {1L, "sensor", 1.5});
+ SessionDataSet second =
+ dataSet(List.of("time", "v"), List.of("TIMESTAMP", "DOUBLE"), new
Object[] {2L, 2.5});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(first, second);
+ ExportDataTable tool = new ExportDataTable();
+ assertTrue(tool.exportBySql("select * from t", 0));
+ needDataTypePrinted = false;
+ assertTrue(tool.exportBySql("select time,v from t", 1));
+ assertEquals("time,device(STRING),v(DOUBLE)", firstLine("dump0_0.csv"));
+ assertEquals("time,v", firstLine("dump1_0.csv"));
+ }
+
+ @Test
+ public void testExportQueryFailureIsReported() throws Exception {
+ Session session = mock(Session.class);
+ when(session.executeQueryStatement(anyString(), anyLong()))
+ .thenThrow(new IoTDBConnectionException("disconnect"));
+ set(ExportDataTree.class, "session", session);
+ assertFalse(new ExportDataTree().exportBySql("select * from root.a.d1",
0));
+ }
+
+ @Test
+ public void testOneDeviceRetriesRetainFailedBatch() throws Exception {
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ List<Long> timestamps = new ArrayList<>(List.of(1L));
+ List<List<TSDataType>> types = new
ArrayList<>(List.of(List.of(TSDataType.INT32)));
+ List<List<Object>> values = new ArrayList<>(List.of(List.of(11)));
+ List<List<String>> measurements = new ArrayList<>(List.of(List.of("s1")));
+ IoTDBConnectionException failure = new
IoTDBConnectionException("disconnect");
+ doThrow(failure)
+ .when(pool)
+ .insertRecordsOfOneDevice("root.a.d1", timestamps, measurements,
types, values);
+ try {
+ invoke(
+ ImportDataTree.class,
+ "writeAndEmptyDataSet",
+ null,
+ new Class<?>[] {String.class, List.class, List.class, List.class,
List.class, int.class},
+ "root.a.d1",
+ timestamps,
+ types,
+ values,
+ measurements,
+ 3);
+ fail("Expected exhausted retries to fail");
+ } catch (IllegalStateException e) {
+ assertSame(failure, e.getCause());
+ }
+ verify(pool, times(4))
+ .insertRecordsOfOneDevice("root.a.d1", timestamps, measurements,
types, values);
+ assertEquals(List.of(1L), timestamps);
+ assertEquals(List.of(List.of(11)), values);
+ assertFalse(types.isEmpty());
+ assertFalse(measurements.isEmpty());
+ }
+
+ @Test
+ public void testOneDeviceSuccessfulRetryClearsBatchOnlyAfterInsert() throws
Exception {
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ List<Long> timestamps = new ArrayList<>(List.of(1L));
+ List<List<TSDataType>> types = new
ArrayList<>(List.of(List.of(TSDataType.INT32)));
+ List<List<Object>> values = new ArrayList<>(List.of(List.of(11)));
+ List<List<String>> measurements = new ArrayList<>(List.of(List.of("s1")));
+ AtomicInteger attempts = new AtomicInteger();
+ doAnswer(
+ invocation -> {
+ assertEquals(List.of(1L), timestamps);
+ assertEquals(List.of(List.of(11)), values);
+ if (attempts.incrementAndGet() < 3) {
+ throw new IoTDBConnectionException("disconnect");
+ }
+ return null;
+ })
+ .when(pool)
+ .insertRecordsOfOneDevice("root.a.d1", timestamps, measurements,
types, values);
+ invoke(
+ ImportDataTree.class,
+ "writeAndEmptyDataSet",
+ null,
+ new Class<?>[] {String.class, List.class, List.class, List.class,
List.class, int.class},
+ "root.a.d1",
+ timestamps,
+ types,
+ values,
+ measurements,
+ 3);
+ assertEquals(3, attempts.get());
+ assertTrue(timestamps.isEmpty());
+ assertTrue(values.isEmpty());
+ assertTrue(types.isEmpty());
+ assertTrue(measurements.isEmpty());
+ }
+
+ @Test
+ public void testAlignedMultiDeviceRetriesRetainFailedBatch() throws
Exception {
+ aligned = true;
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ List<String> devices = new ArrayList<>(List.of("root.a.d1"));
+ List<Long> timestamps = new ArrayList<>(List.of(1L));
+ List<List<TSDataType>> types = new
ArrayList<>(List.of(List.of(TSDataType.INT32)));
+ List<List<Object>> values = new ArrayList<>(List.of(List.of(11)));
+ List<List<String>> measurements = new ArrayList<>(List.of(List.of("s1")));
+ doThrow(new IoTDBConnectionException("disconnect"))
+ .when(pool)
+ .insertAlignedRecords(devices, timestamps, measurements, types,
values);
+ try {
+ invoke(
+ ImportDataTree.class,
+ "writeAndEmptyDataSet",
+ null,
+ new Class<?>[] {List.class, List.class, List.class, List.class,
List.class, int.class},
+ devices,
+ timestamps,
+ types,
+ values,
+ measurements,
+ 3);
+ fail("Expected exhausted retries to fail");
+ } catch (IllegalStateException expected) {
+ assertTrue(expected.getCause() instanceof IoTDBConnectionException);
+ }
+ verify(pool, times(4)).insertAlignedRecords(devices, timestamps,
measurements, types, values);
+ assertEquals(List.of("root.a.d1"), devices);
+ assertEquals(List.of(1L), timestamps);
+ assertEquals(List.of(List.of(11)), values);
+ }
+
+ @Test
+ public void testTableRetriesRetainTabletAndCloseSessions() throws Exception {
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mock(ITableSession.class);
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ Tablet tablet =
+ new Tablet(
+ "t", List.of("v"), List.of(TSDataType.DOUBLE),
List.of(ColumnCategory.FIELD), 100);
+ tablet.addTimestamp(0, 1L);
+ tablet.addValue("v", 0, 1.5);
+ doThrow(new
IoTDBConnectionException("disconnect")).when(session).insert(tablet);
+ try {
+ invoke(
+ ImportDataTable.class,
+ "writeAndEmptyDataSet",
+ null,
+ new Class<?>[] {Tablet.class, int.class},
+ tablet,
+ 3);
+ fail("Expected exhausted retries to fail");
+ } catch (IllegalStateException expected) {
+ assertTrue(expected.getCause() instanceof IoTDBConnectionException);
+ }
+ assertEquals(1, tablet.getRowSize());
+ verify(session, times(4)).insert(tablet);
+ verify(session, times(4)).close();
+ }
+
+ @Test
+ public void testWorkerFailurePreventsSuccessAndResetsForNextImport() throws
Exception {
+ Path csv = directory.resolve("input.csv");
+ Files.writeString(csv, "Time,root.a.d1.s1(INT32)\n1,11\n");
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ doThrow(new IoTDBConnectionException("disconnect"))
+ .when(pool)
+ .insertRecords(any(), any(), any(), any(), any());
+ long successes = loadFileSuccessfulNum.sum();
+ ImportDataScanTool.putToQueue(csv.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTree()));
+ assertEquals(successes, loadFileSuccessfulNum.sum());
+ doNothing().when(pool).insertRecords(any(), any(), any(), any(), any());
+ ImportDataScanTool.putToQueue(csv.toString());
+ assertTrue(AbstractImportData.init(new ImportDataTree()));
+ assertEquals(successes + 1, loadFileSuccessfulNum.sum());
+ }
+
+ @Test
+ public void testFailedCsvBatchDoesNotSkipNextFile() throws Exception {
+ Path bad = directory.resolve("a_bad.csv");
+ Path good = directory.resolve("b_good.csv");
+ Files.writeString(bad, "Time,root.a.d1.s1(INT32)\n1,11\n");
+ Files.writeString(good, "Time,root.a.d2.s1(INT32)\n1,22\n");
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ AtomicInteger goodInserted = new AtomicInteger();
+ doAnswer(
+ invocation -> {
+ if (((List<?>) invocation.getArgument(0)).contains("root.a.d1"))
{
+ throw new StatementExecutionException("type mismatch");
+ }
+ if (((List<?>) invocation.getArgument(0)).contains("root.a.d2"))
{
+ goodInserted.incrementAndGet();
+ }
+ return null;
+ })
+ .when(pool)
+ .insertRecords(any(), any(), any(), any(), any());
+ ImportDataScanTool.putToQueue(bad.toString());
+ ImportDataScanTool.putToQueue(good.toString());
+
+ assertFalse(AbstractImportData.init(new ImportDataTree()));
+
+ assertEquals(1, goodInserted.get());
+ assertNull(ImportDataScanTool.pollFromQueue());
+ }
+
+ @Test
+ public void testRejectedCsvAndFailedSqlAffectImportExitStatus() throws
Exception {
+ Path badCsv = directory.resolve("bad_header.csv");
+ Files.writeString(badCsv, "ts,root.a.d1.s1(INT32)\n1,11\n");
+ ImportDataScanTool.putToQueue(badCsv.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTree()));
+
+ Path badTableCsv = directory.resolve("bad_table_header.csv");
+ Files.writeString(badTableCsv, "ts,v\n1,11\n");
+ ImportDataScanTool.putToQueue(badTableCsv.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTable()));
+
+ Path failedTreeSql = directory.resolve("failed_tree.sql");
+ Files.writeString(failedTreeSql, "insert into root.a.d1(time,s1)
values(1,11)\n");
+ SessionPool treePool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", treePool);
+ doThrow(new StatementExecutionException("statement failed"))
+ .when(treePool)
+ .executeNonQueryStatement(anyString());
+ ImportDataScanTool.putToQueue(failedTreeSql.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTree()));
+
+ Path failedTableSql = directory.resolve("failed_table.sql");
+ Files.writeString(failedTableSql, "insert into db.t(time,v)
values(1,11)\n");
+ ITableSessionPool tablePool = mock(ITableSessionPool.class);
+ ITableSession tableSession = mock(ITableSession.class);
+ when(tablePool.getSession()).thenReturn(tableSession);
+ doThrow(new StatementExecutionException("statement failed"))
+ .when(tableSession)
+ .executeNonQueryStatement(anyString());
+ set(ImportDataTable.class, "sessionPool", tablePool);
+ ImportDataScanTool.putToQueue(failedTableSql.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTable()));
+ }
+
+ @Test
+ public void testRejectedFilesAreCountedOnceInBothModels() throws Exception {
+ Path empty = directory.resolve("empty.csv");
+ Path header = directory.resolve("invalid.csv");
+ Path extension = directory.resolve("input.other");
+ Files.writeString(empty, "");
+ Files.writeString(header, "ts,v\n1,11\n");
+ Files.writeString(extension, "Time,v\n1,11\n");
+ for (AbstractImportData importer : List.of(new ImportDataTree(), new
ImportDataTable())) {
+ for (Path source :
+ List.of(
+ empty,
+ header,
+ extension,
+ directory.resolve("missing.csv"),
+ directory.resolve("missing.sql"))) {
+ long failures = loadFileFailedNum.sum();
+ long successes = loadFileSuccessfulNum.sum();
+ ImportDataScanTool.putToQueue(source.toString());
+ assertFalse(AbstractImportData.init(importer));
+ assertEquals(failures + 1, loadFileFailedNum.sum());
+ assertEquals(successes, loadFileSuccessfulNum.sum());
+ }
+ }
+ }
+
+ @Test
+ public void testFailedDeviceCsvDoesNotSkipNextFile() throws Exception {
+ Path bad = directory.resolve("a_bad.csv");
+ Path good = directory.resolve("b_good.csv");
+ Files.writeString(bad, "Time,Device,s1(INT32)\n1,root.a.d1,11\n");
+ Files.writeString(good, "Time,Device,s1(INT32)\n1,root.a.d2,22\n");
+ SessionPool pool = mock(SessionPool.class);
+ set(ImportDataTree.class, "sessionPool", pool);
+ AtomicInteger goodInserted = new AtomicInteger();
+ doAnswer(
+ invocation -> {
+ if ("root.a.d1".equals(invocation.getArgument(0))) {
+ throw new StatementExecutionException("type mismatch");
+ }
+ goodInserted.incrementAndGet();
+ return null;
+ })
+ .when(pool)
+ .insertRecordsOfOneDevice(anyString(), any(), any(), any(), any());
+ ImportDataScanTool.putToQueue(bad.toString());
+ ImportDataScanTool.putToQueue(good.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTree()));
+ assertEquals(1, goodInserted.get());
+ assertNull(ImportDataScanTool.pollFromQueue());
+ }
+
+ @Test
+ public void testFailedTableCsvDoesNotSkipNextFile() throws Exception {
+ Path bad = directory.resolve("a_bad.csv");
+ Path good = directory.resolve("b_good.csv");
+ Files.writeString(bad, "Time,v(INT32)\n1,11\n");
+ Files.writeString(good, "Time,v(INT32)\n2,22\n");
+ set(ImportDataTable.class, "dataTypes", Map.of("v", TSDataType.INT32));
+ set(ImportDataTable.class, "columnCategory", Map.of("v",
ColumnCategory.FIELD));
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mock(ITableSession.class);
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ AtomicInteger goodInserted = new AtomicInteger();
+ doAnswer(
+ invocation -> {
+ Tablet tablet = invocation.getArgument(0);
+ if (tablet.getTimestamps()[0] == 1L) {
+ throw new StatementExecutionException("type mismatch");
+ }
+ goodInserted.incrementAndGet();
+ return null;
+ })
+ .when(session)
+ .insert(any(Tablet.class));
+ ImportDataScanTool.putToQueue(bad.toString());
+ ImportDataScanTool.putToQueue(good.toString());
+ assertFalse(AbstractImportData.init(new ImportDataTable()));
+ assertEquals(1, goodInserted.get());
+ assertNull(ImportDataScanTool.pollFromQueue());
+ }
+
+ @Test
+ public void testTypedTableCsvUsesOriginalHeaderSpelling() throws Exception {
+ Map<String, TSDataType> types = new LinkedHashMap<>();
+ types.put("id", TSDataType.STRING);
+ types.put("v", TSDataType.DOUBLE);
+ Map<String, ColumnCategory> categories = new LinkedHashMap<>();
+ categories.put("id", ColumnCategory.TAG);
+ categories.put("v", ColumnCategory.FIELD);
+ set(ImportDataTable.class, "dataTypes", types);
+ set(ImportDataTable.class, "columnCategory", categories);
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mock(ITableSession.class);
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ Path csv = directory.resolve("typed.csv");
+ Files.writeString(csv, "TIME,ID(STRING),v(double)\n1,sensor,1.5\n");
+ try (CSVParser records = readCsvFile(csv.toString())) {
+ new ImportDataTable()
+ .writeData(
+ records.getHeaderNames(),
+ records.getRecords(),
+ directory.resolve("typed.failed").toString());
+ }
+ verify(session).insert(any(Tablet.class));
+ assertFalse(Files.exists(directory.resolve("typed.failed")));
+ }
+
+ @Test
+ public void
testTableCsvRoundTripPreservesWhitespaceAndParenthesesInColumnNames()
+ throws Exception {
+ assertTableCsvRoundTrip("value(double) with space");
+ }
+
+ @Test
+ public void testTableCsvRoundTripPreservesTrailingParenthesesInColumnNames()
throws Exception {
+ for (String column : List.of("value(foo)", "value(double)",
"value(foo)(double)")) {
+ assertTableCsvRoundTrip(column);
+ }
+ }
+
+ @Test
+ public void testTableCsvPrefersCompleteColumnNameOverTypeSuffix() throws
Exception {
+ Map<String, TSDataType> types = new LinkedHashMap<>();
+ types.put("value", TSDataType.DOUBLE);
+ types.put("value(double)", TSDataType.DOUBLE);
+ Map<String, ColumnCategory> categories = new LinkedHashMap<>();
+ categories.put("value", ColumnCategory.FIELD);
+ categories.put("value(double)", ColumnCategory.FIELD);
+ set(ImportDataTable.class, "dataTypes", types);
+ set(ImportDataTable.class, "columnCategory", categories);
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mock(ITableSession.class);
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ doAnswer(
+ invocation -> {
+ Tablet tablet = invocation.getArgument(0);
+ assertEquals(2, tablet.getSchemas().size());
+ assertEquals("value(double)",
tablet.getSchemas().get(1).getMeasurementName());
+ assertEquals(1, tablet.getRowSize());
+ assertEquals(1L, tablet.getTimestamps()[0]);
+ assertEquals(1.5, ((double[]) tablet.getValues()[1])[0], 0);
+ assertEquals(0, ((double[]) tablet.getValues()[0])[0], 0);
+ return null;
+ })
+ .when(session)
+ .insert(any(Tablet.class));
+ Path csv = directory.resolve("parentheses.csv");
+ Files.writeString(csv, "TIME,VALUE(DOUBLE)\n1,1.5\n");
+ try (CSVParser records = readCsvFile(csv.toString())) {
+ new ImportDataTable()
+ .writeData(
+ records.getHeaderNames(),
+ records.getRecords(),
+ directory.resolve("parentheses.failed").toString());
+ }
+ verify(session).insert(any(Tablet.class));
+ assertFalse(Files.exists(directory.resolve("parentheses.failed")));
+ }
+
+ @Test
+ public void testTypedTableCsvRejectsAmbiguousTwoColumnRoundTrip() throws
Exception {
+ assertTableCsvTwoColumnRoundTrip(true);
+ }
+
+ @Test
+ public void testUntypedTableCsvPreservesTwoColumnRoundTrip() throws
Exception {
+ assertTableCsvTwoColumnRoundTrip(false);
+ }
+
+ private void assertTableCsvTwoColumnRoundTrip(boolean typed) throws
Exception {
+ Map<String, TSDataType> types = new LinkedHashMap<>();
+ types.put("value", TSDataType.DOUBLE);
+ types.put("value(double)", TSDataType.DOUBLE);
+ Map<String, ColumnCategory> categories = new LinkedHashMap<>();
+ categories.put("value", ColumnCategory.FIELD);
+ categories.put("value(double)", ColumnCategory.FIELD);
+ set(ImportDataTable.class, "dataTypes", types);
+ set(ImportDataTable.class, "columnCategory", categories);
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mockTable("t");
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ doAnswer(
+ invocation -> {
+ Tablet tablet = invocation.getArgument(0);
+ assertEquals(2, tablet.getSchemas().size());
+ assertEquals("value",
tablet.getSchemas().get(0).getMeasurementName());
+ assertEquals("value(double)",
tablet.getSchemas().get(1).getMeasurementName());
+ assertEquals(1, tablet.getRowSize());
+ assertEquals(1L, tablet.getTimestamps()[0]);
+ assertEquals(1.5, ((double[]) tablet.getValues()[0])[0], 0);
+ assertEquals(2.5, ((double[]) tablet.getValues()[1])[0], 0);
+ return null;
+ })
+ .when(session)
+ .insert(any(Tablet.class));
+ exportType = "csv";
+ needDataTypePrinted = typed;
+ for (int i = 0; i < 2; i++) {
+ List<String> columns =
+ i == 0
+ ? List.of("time", "value", "value(double)")
+ : List.of("time", "value(double)", "value");
+ Object[] row = i == 0 ? new Object[] {1L, 1.5, 2.5} : new Object[] {1L,
2.5, 1.5};
+ SessionDataSet result = dataSet(columns, List.of("TIMESTAMP", "DOUBLE",
"DOUBLE"), row);
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+ assertTrue(new ExportDataTable().exportBySql("select * from t", i));
+ long failedFiles = loadFileFailedNum.sum();
+ long successfulFiles = loadFileSuccessfulNum.sum();
+ ImportDataScanTool.putToQueue(directory.resolve("dump" + i +
"_0.csv").toString());
+ assertEquals(!typed, AbstractImportData.init(new ImportDataTable()));
+ assertEquals(failedFiles + (typed ? 1 : 0), loadFileFailedNum.sum());
+ assertEquals(successfulFiles + (typed ? 0 : 1),
loadFileSuccessfulNum.sum());
+ assertNull(ImportDataScanTool.pollFromQueue());
+ }
+ verify(session, times(typed ? 0 : 2)).insert(any(Tablet.class));
+ }
+
+ private void assertTableCsvRoundTrip(String column) throws Exception {
+ set(ImportDataTable.class, "dataTypes", Map.of(column, TSDataType.DOUBLE));
+ set(ImportDataTable.class, "columnCategory", Map.of(column,
ColumnCategory.FIELD));
+ ITableSessionPool pool = mock(ITableSessionPool.class);
+ ITableSession session = mockTable("t");
+ when(pool.getSession()).thenReturn(session);
+ set(ImportDataTable.class, "sessionPool", pool);
+ doAnswer(
+ invocation -> {
+ Tablet tablet = invocation.getArgument(0);
+ assertEquals(1, tablet.getSchemas().size());
+ assertEquals(column,
tablet.getSchemas().get(0).getMeasurementName());
+ assertEquals(1, tablet.getRowSize());
+ assertEquals(1L, tablet.getTimestamps()[0]);
+ assertEquals(1.5, ((double[]) tablet.getValues()[0])[0], 0);
+ return null;
+ })
+ .when(session)
+ .insert(any(Tablet.class));
+ exportType = "csv";
+ for (boolean typed : List.of(false, true)) {
+ needDataTypePrinted = typed;
+ SessionDataSet result =
+ dataSet(List.of("time", column), List.of("TIMESTAMP", "DOUBLE"), new
Object[] {1L, 1.5});
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(result);
+ int index = typed ? 1 : 0;
+ assertTrue(new ExportDataTable().exportBySql("select * from t", index));
+ String csv = "dump" + index + "_0.csv";
+ assertEquals("time," + column + (typed ? "(DOUBLE)" : ""),
firstLine(csv));
+ try (CSVParser records = readCsvFile(directory.resolve(csv).toString()))
{
+ new ImportDataTable()
+ .writeData(
+ records.getHeaderNames(),
+ records.getRecords(),
+ directory.resolve("roundtrip.failed").toString());
+ }
+ }
+ verify(session, times(2)).insert(any(Tablet.class));
+ assertFalse(Files.exists(directory.resolve("roundtrip.failed")));
+ }
+
+ @Test
+ public void testSuccessArchivalPreservesFilesAndSidecars() throws Exception {
+ for (ImportTsFileOperation operation :
+ List.of(
+ ImportTsFileOperation.CP, ImportTsFileOperation.HARDLINK,
ImportTsFileOperation.MV)) {
+ Path input =
Files.createDirectories(directory.resolve(operation.name()).resolve("input"));
+ Path nested = Files.createDirectories(input.resolve("nested"));
+ Path source = nested.resolve("one.tsfile");
+ Path archive =
+
Files.createDirectories(directory.resolve(operation.name()).resolve("success"));
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), "content" + suffix);
+ }
+ successDir = archive.toString();
+ successOperation = operation;
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+ new ImportDataTree().processingFile(source.toString(), true);
+ assertTrue(Files.isDirectory(archive));
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals(
+ "content" + suffix,
Files.readString(archive.resolve("nested/one.tsfile" + suffix)));
+ assertEquals(operation != ImportTsFileOperation.MV,
Files.exists(Path.of(source + suffix)));
+ }
+ Path other = input.resolve("other.tsfile");
+ Files.writeString(other, "other");
+ new ImportDataTree().processingFile(other.toString(), true);
+ assertEquals("other", Files.readString(archive.resolve("other.tsfile")));
+ assertEquals("content",
Files.readString(archive.resolve("nested/one.tsfile")));
+ }
+ }
+
+ @Test
+ public void testArchivalPreservesCollidingFlattenedNames() throws Exception {
+ for (boolean isSuccess : List.of(false, true)) {
+ for (ImportTsFileOperation operation :
+ List.of(
+ ImportTsFileOperation.CP, ImportTsFileOperation.HARDLINK,
ImportTsFileOperation.MV)) {
+ Path root = Files.createDirectories(directory.resolve(isSuccess + "-"
+ operation));
+ Path input = Files.createDirectories(root.resolve("input"));
+ Path archive = root.resolve("archive");
+ Files.createDirectories(input.resolve("a"));
+ successDir = failDir = archive.toString();
+ successOperation = failOperation = operation;
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+ for (String name : List.of("a/b.tsfile", "a_b.tsfile")) {
+ Path source = input.resolve(name);
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), name + suffix);
+ }
+ new ImportDataTree().processingFile(source.toString(), isSuccess);
+ }
+ for (String name : List.of("a/b.tsfile", "a_b.tsfile")) {
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals(name + suffix, Files.readString(archive.resolve(name
+ suffix)));
+ assertEquals(
+ operation != ImportTsFileOperation.MV,
Files.exists(input.resolve(name + suffix)));
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testExistingArchiveFilesAreNotOverwritten() throws Exception {
+ for (ImportTsFileOperation operation :
+ List.of(
+ ImportTsFileOperation.CP, ImportTsFileOperation.HARDLINK,
ImportTsFileOperation.MV)) {
+ for (String conflict : List.of("", ".resource", ".mods")) {
+ Path root = Files.createDirectories(directory.resolve(operation + "-"
+ conflict));
+ Path input = Files.createDirectories(root.resolve("input"));
+ Path archive = Files.createDirectories(root.resolve("archive"));
+ Path source = input.resolve("one.tsfile");
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), "source" + suffix);
+ }
+ Path occupied = archive.resolve("one.tsfile" + conflict);
+ Files.writeString(occupied, "existing");
+ successDir = archive.toString();
+ successOperation = operation;
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+ try {
+ new ImportDataTree().processingFile(source.toString(), true);
+ fail("Expected existing archive file to be rejected");
+ } catch (FileAlreadyExistsException expected) {
+ assertEquals("existing", Files.readString(occupied));
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals("source" + suffix, Files.readString(Path.of(source +
suffix)));
+ if (!suffix.equals(conflict)) {
+ assertFalse(Files.exists(archive.resolve("one.tsfile" +
suffix)));
+ }
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testArchivalErrorsArePropagated() throws Exception {
+ Path source = directory.resolve("one.tsfile");
+ Files.writeString(source, "source");
+ Path archive = Files.createDirectories(directory.resolve("success"));
+ Files.createDirectories(archive.resolve("one.tsfile"));
+ Files.writeString(archive.resolve("one.tsfile/child"), "occupied");
+ successDir = archive.toString();
+ successOperation = ImportTsFileOperation.CP;
+ ImportTsFileScanTool.setSourceFullPath(directory.toString());
+ try {
+ new ImportDataTree().processingFile(source.toString(), true);
+ fail("Expected archival error to propagate");
+ } catch (IOException expected) {
+ assertEquals("source", Files.readString(source));
+ }
+ }
+
+ @Test
+ public void testImportTsFileArchivalPreservesRelativePathsAndSidecars()
throws Exception {
+ for (boolean isSuccess : List.of(false, true)) {
+ for (ImportTsFileOperation operation :
+ List.of(
+ ImportTsFileOperation.CP, ImportTsFileOperation.HARDLINK,
ImportTsFileOperation.MV)) {
+ Path root = Files.createDirectories(directory.resolve("tsfile-" +
isSuccess + operation));
+ Path input = Files.createDirectories(root.resolve("input"));
+ Files.createDirectories(input.resolve("a"));
+ Path archive = Files.createDirectories(root.resolve("archive"));
+ ImportTsFileBase.setSuccessAndFailDirAndOperation(
+ archive.resolve("success").toString(),
+ operation,
+ archive.resolve("fail").toString(),
+ operation);
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+
+ for (String name : List.of("a/b.tsfile", "a_b.tsfile")) {
+ Path source = input.resolve(name);
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), name + suffix);
+ }
+ ImportTsFileBase.processingFile(source.toString(), isSuccess);
+ }
+
+ Path destination = archive.resolve(isSuccess ? "success" : "fail");
+ for (String name : List.of("a/b.tsfile", "a_b.tsfile")) {
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals(name + suffix,
Files.readString(destination.resolve(name + suffix)));
+ assertEquals(
+ operation != ImportTsFileOperation.MV,
Files.exists(input.resolve(name + suffix)));
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testImportTsFileArchiveConflictsKeepSourceFiles() throws
Exception {
+ for (boolean isSuccess : List.of(false, true)) {
+ for (String conflict : List.of("", ".resource", ".mods")) {
+ Path root = Files.createDirectories(directory.resolve("conflict-" +
isSuccess + conflict));
+ Path input = Files.createDirectories(root.resolve("input"));
+ Path archive = Files.createDirectories(root.resolve("archive"));
+ Path source = input.resolve("one.tsfile");
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), "source" + suffix);
+ }
+ Path occupied = archive.resolve("one.tsfile" + conflict);
+ Files.writeString(occupied, "existing");
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+ ImportTsFileBase.setSuccessAndFailDirAndOperation(
+ archive.toString(),
+ ImportTsFileOperation.MV,
+ archive.toString(),
+ ImportTsFileOperation.MV);
+ try {
+ ImportTsFileBase.processingFile(source.toString(), isSuccess);
+ fail("Expected existing archive file to be rejected");
+ } catch (FileAlreadyExistsException expected) {
+ assertEquals("existing", Files.readString(occupied));
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals("source" + suffix, Files.readString(Path.of(source +
suffix)));
+ if (!suffix.equals(conflict)) {
+ assertFalse(Files.exists(archive.resolve("one.tsfile" +
suffix)));
+ }
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testSingleFileRelativeSourceArchival() throws Exception {
+ Path workingDirectory = Path.of("").toAbsolutePath();
+ Path source = Files.createTempFile(workingDirectory, "one", ".tsfile");
+ try {
+ Files.writeString(source, "source");
+ Path archive = Files.createDirectories(directory.resolve("success"));
+ successDir = archive.toString();
+ successOperation = ImportTsFileOperation.CP;
+ ImportTsFileScanTool.setSourceFullPath(source.getFileName().toString());
+ new ImportDataTree().processingFile(source.toString(), true);
+ assertEquals("source",
Files.readString(archive.resolve(source.getFileName())));
+ } finally {
+ Files.deleteIfExists(source);
+ }
+ }
+
+ private String read(String name) throws IOException {
+ return Files.readString(directory.resolve(name));
+ }
+
+ private String firstLine(String name) throws IOException {
+ return Files.readAllLines(directory.resolve(name)).get(0);
+ }
+
+ private static void set(Class<?> owner, String name, Object value) throws
Exception {
+ Field field = owner.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(null, value);
+ }
+
+ private static Object invoke(
+ Class<?> owner, String name, Object target, Class<?>[] parameterTypes,
Object... arguments)
+ throws Exception {
+ Method method = owner.getDeclaredMethod(name, parameterTypes);
+ method.setAccessible(true);
+ try {
+ return method.invoke(target, arguments);
+ } catch (InvocationTargetException e) {
+ throw (Exception) e.getCause();
+ }
+ }
+
+ private static void mockTree(SessionDataSet dataSet) throws Exception {
+ Session session = mock(Session.class);
+ when(session.executeQueryStatement(anyString(),
anyLong())).thenReturn(dataSet);
+ set(ExportDataTree.class, "session", session);
+ }
+
+ private static ITableSession mockTable(String... names) throws Exception {
+ ITableSession session = mock(ITableSession.class);
+ set(ExportDataTable.class, "tableSession", session);
+ set(ExportDataTable.class, "tables", new
ArrayList<>(Arrays.asList(names)));
+ set(ExportDataTable.class, "allTables", new
ArrayList<>(Arrays.asList(names)));
+ return session;
+ }
+
+ private static SessionDataSet dataSet(List<String> names, List<String>
types, Object[]... rows)
+ throws Exception {
+ SessionDataSet dataSet = mock(SessionDataSet.class);
+ SessionDataSet.DataIterator iterator =
mock(SessionDataSet.DataIterator.class);
+ AtomicInteger row = new AtomicInteger(-1);
+ when(dataSet.getColumnNames()).thenAnswer(invocation -> new
ArrayList<>(names));
+ when(dataSet.getColumnTypes()).thenReturn(types);
+ when(dataSet.iterator()).thenReturn(iterator);
+ when(iterator.getColumnTypeList()).thenReturn(types);
+ when(iterator.next()).thenAnswer(invocation -> row.incrementAndGet() <
rows.length);
+ when(iterator.getString(anyInt()))
+ .thenAnswer(
+ invocation -> {
+ Object value = rows[row.get()][(Integer)
invocation.getArgument(0) - 1];
+ return value == null ? null : value.toString();
+ });
+ when(iterator.getLong(anyInt()))
+ .thenAnswer(
+ invocation ->
+ ((Number) rows[row.get()][(Integer) invocation.getArgument(0)
- 1]).longValue());
+ when(iterator.isNull(anyInt()))
+ .thenAnswer(invocation -> rows[row.get()][(Integer)
invocation.getArgument(0) - 1] == null);
+ return dataSet;
+ }
+}
diff --git
a/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/tsfile/ImportTsFileTest.java
b/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/tsfile/ImportTsFileTest.java
new file mode 100644
index 00000000000..481f678b52e
--- /dev/null
+++
b/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/tsfile/ImportTsFileTest.java
@@ -0,0 +1,205 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.tool.tsfile;
+
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.session.pool.SessionPool;
+import org.apache.iotdb.tool.common.ImportTsFileOperation;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+public class ImportTsFileTest extends AbstractTsFileTool {
+
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private final Map<Field, Object> originalFields = new HashMap<>();
+ private SessionPool sessionPool;
+ private Path input;
+ private Path archive;
+
+ @Before
+ public void setUp() throws Exception {
+ for (Class<?> owner :
+ List.of(
+ ImportTsFile.class,
+ ImportTsFileBase.class,
+ ImportTsFileLocally.class,
+ ImportTsFileRemotely.class,
+ ImportTsFileScanTool.class)) {
+ for (Field field : owner.getDeclaredFields()) {
+ if (Modifier.isStatic(field.getModifiers()) &&
!Modifier.isFinal(field.getModifiers())) {
+ field.setAccessible(true);
+ originalFields.put(field, field.get(null));
+ }
+ }
+ }
+ sessionPool = mock(SessionPool.class);
+ ImportTsFileLocally.setSessionPool(sessionPool);
+ setImportOption("isRemoteLoad", false);
+ setImportOption("threadNum", 2);
+ input =
Files.createDirectories(temporaryFolder.getRoot().toPath().resolve("input"));
+ archive =
Files.createDirectories(temporaryFolder.getRoot().toPath().resolve("archive"));
+ ImportTsFileScanTool.setSourceFullPath(input.toString());
+ ImportTsFileBase.setSuccessAndFailDirAndOperation(
+ archive.toString(), ImportTsFileOperation.CP, archive.toString(),
ImportTsFileOperation.CP);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ while (ImportTsFileScanTool.pollFromQueue() != null) {
+ // Drain any files left by a failing assertion before restoring the
shared state.
+ }
+ ImportTsFileBase.resetCounters();
+ for (Map.Entry<Field, Object> entry : originalFields.entrySet()) {
+ entry.getKey().set(null, entry.getValue());
+ }
+ }
+
+ @Test
+ public void testSuccessfulImportReturnsZero() throws Exception {
+ Path source = createTsFileGroup("one.tsfile");
+ ImportTsFileScanTool.putToQueue(source.toString());
+ assertEquals(CODE_OK, ImportTsFile.asyncImportTsFiles());
+ assertArchived("one.tsfile");
+ verify(sessionPool).executeNonQueryStatement(anyString());
+ assertNull(ImportTsFileScanTool.pollFromQueue());
+ }
+
+ @Test
+ public void testArchiveConflictsReturnErrorAndContinueImport() throws
Exception {
+ for (ImportTsFileOperation operation :
+ List.of(
+ ImportTsFileOperation.CP, ImportTsFileOperation.HARDLINK,
ImportTsFileOperation.MV)) {
+ ImportTsFileBase.setSuccessAndFailDirAndOperation(
+ archive.toString(), operation, archive.toString(), operation);
+ for (String conflict : List.of("", ".resource", ".mods")) {
+ String name = operation + conflict + ".tsfile";
+ Path source = createTsFileGroup(name);
+ Path occupied = archive.resolve(name + conflict);
+ Files.writeString(occupied, "existing");
+ String goodName = "good-" + name;
+ Path good = createTsFileGroup(goodName);
+ ImportTsFileScanTool.putToQueue(source.toString());
+ ImportTsFileScanTool.putToQueue(good.toString());
+
+ assertEquals(CODE_ERROR, ImportTsFile.asyncImportTsFiles());
+ assertEquals("existing", Files.readString(occupied));
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals("source" + suffix, Files.readString(Path.of(source +
suffix)));
+ if (!suffix.equals(conflict)) {
+ assertFalse(Files.exists(archive.resolve(name + suffix)));
+ }
+ }
+ assertArchived(goodName);
+ assertNull(ImportTsFileScanTool.pollFromQueue());
+ }
+ }
+ verify(sessionPool, times(18)).executeNonQueryStatement(anyString());
+ }
+
+ @Test
+ public void testArchiveIoFailureReturnsError() throws Exception {
+ Files.createDirectories(input.resolve("nested"));
+ Path source = createTsFileGroup("nested/one.tsfile");
+ Files.writeString(archive.resolve("nested"), "blocking file");
+ ImportTsFileScanTool.putToQueue(source.toString());
+ assertEquals(CODE_ERROR, ImportTsFile.asyncImportTsFiles());
+ assertEquals("source", Files.readString(source));
+ assertEquals("blocking file", Files.readString(archive.resolve("nested")));
+ }
+
+ @Test
+ public void testLoadFailureReturnsErrorWhenArchiveSucceeds() throws
Exception {
+ Path source = createTsFileGroup("one.tsfile");
+ doThrow(new StatementExecutionException("load failed"))
+ .when(sessionPool)
+ .executeNonQueryStatement(anyString());
+ ImportTsFileScanTool.putToQueue(source.toString());
+ assertEquals(CODE_ERROR, ImportTsFile.asyncImportTsFiles());
+ assertArchived("one.tsfile");
+ verify(sessionPool).executeNonQueryStatement(anyString());
+ }
+
+ @Test
+ public void testExitCodeResetsForNextImport() throws Exception {
+ Path source = createTsFileGroup("one.tsfile");
+ Path occupied = archive.resolve("one.tsfile");
+ Files.writeString(occupied, "existing");
+ ImportTsFileScanTool.putToQueue(source.toString());
+ assertEquals(CODE_ERROR, ImportTsFile.asyncImportTsFiles());
+ Files.delete(occupied);
+ ImportTsFileScanTool.putToQueue(source.toString());
+ assertEquals(CODE_OK, ImportTsFile.asyncImportTsFiles());
+ assertArchived("one.tsfile");
+ verify(sessionPool, times(2)).executeNonQueryStatement(anyString());
+ }
+
+ @Test
+ public void testImportFromTargetPathReturnsErrorForUnprocessedFiles() throws
Exception {
+ setImportOption("threadNum", 0);
+ // Exercise the entry point without starting workers or connecting to a
server.
+ setImportOption("isRemoteLoad", true);
+ Path source = createTsFileGroup("one.tsfile");
+ setImportOption("source", source.toString());
+ assertEquals(CODE_ERROR, ImportTsFile.importFromTargetPath());
+ assertEquals(source.toString(), ImportTsFileScanTool.pollFromQueue());
+ }
+
+ private Path createTsFileGroup(String name) throws Exception {
+ Path source = input.resolve(name);
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ Files.writeString(Path.of(source + suffix), "source" + suffix);
+ }
+ return source;
+ }
+
+ private void assertArchived(String name) throws Exception {
+ for (String suffix : List.of("", ".resource", ".mods")) {
+ assertEquals("source" + suffix, Files.readString(archive.resolve(name +
suffix)));
+ }
+ }
+
+ private static void setImportOption(String name, Object value) throws
Exception {
+ Field field = ImportTsFile.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(null, value);
+ }
+}
diff --git a/scripts/tools/windows/export-data.bat
b/scripts/tools/windows/export-data.bat
index f239b09bb48..06d489a2d30 100644
--- a/scripts/tools/windows/export-data.bat
+++ b/scripts/tools/windows/export-data.bat
@@ -57,6 +57,4 @@ pause
@REM
-----------------------------------------------------------------------------
:finally
-ENDLOCAL
-
-EXIT /B %ret_code%
+ENDLOCAL & EXIT /B %ret_code%
diff --git a/scripts/tools/windows/export-tsfile.bat
b/scripts/tools/windows/export-tsfile.bat
index 0b070b51fbc..95c9e5c29e9 100644
--- a/scripts/tools/windows/export-tsfile.bat
+++ b/scripts/tools/windows/export-tsfile.bat
@@ -57,6 +57,4 @@ pause
@REM
-----------------------------------------------------------------------------
:finally
-ENDLOCAL
-
-EXIT /B %ret_code%
\ No newline at end of file
+ENDLOCAL & EXIT /B %ret_code%
diff --git a/scripts/tools/windows/import-data.bat
b/scripts/tools/windows/import-data.bat
index 56518c78600..89d2feee0f6 100644
--- a/scripts/tools/windows/import-data.bat
+++ b/scripts/tools/windows/import-data.bat
@@ -58,6 +58,4 @@ pause
@REM
-----------------------------------------------------------------------------
:finally
-ENDLOCAL
-
-EXIT /B %ret_code%
+ENDLOCAL & EXIT /B %ret_code%
diff --git a/scripts/tools/windows/schema/export-schema.bat
b/scripts/tools/windows/schema/export-schema.bat
index 9a082ac05ff..c7f29781433 100644
--- a/scripts/tools/windows/schema/export-schema.bat
+++ b/scripts/tools/windows/schema/export-schema.bat
@@ -57,6 +57,4 @@ pause
@REM
-----------------------------------------------------------------------------
:finally
-ENDLOCAL
-
-EXIT /B %ret_code%
+ENDLOCAL & EXIT /B %ret_code%
diff --git a/scripts/tools/windows/schema/import-schema.bat
b/scripts/tools/windows/schema/import-schema.bat
index a4ddd93b6ff..930da21d1cc 100644
--- a/scripts/tools/windows/schema/import-schema.bat
+++ b/scripts/tools/windows/schema/import-schema.bat
@@ -58,6 +58,4 @@ pause
@REM
-----------------------------------------------------------------------------
:finally
-ENDLOCAL
-
-EXIT /B %ret_code%
+ENDLOCAL & EXIT /B %ret_code%
