diff --git a/iotdb-client/cli/pom.xml b/iotdb-client/cli/pom.xml index 5de247a084a09..a1e00508740bd 100644 --- a/iotdb-client/cli/pom.xml +++ b/iotdb-client/cli/pom.xml @@ -54,6 +54,11 @@ iotdb-antlr 2.0.11-SNAPSHOT + + org.apache.iotdb + iotdb-relational-grammar + 2.0.11-SNAPSHOT + org.apache.iotdb node-commons 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 0b0f30d970d95..ec00a7dc559c0 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 a260732ab554a..8321a28d08ef9 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 ed2d96dc8606b..b6999dea4277d 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 3769d855d161b..9934efbef6fc4 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 @@ 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 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 f1d506a72b49c..e3c33950d42d6 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 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 getCsvHeaders(SessionDataSet dataSet, boolean treeModel) { + List headers = new ArrayList<>(dataSet.getColumnNames()); + if (Boolean.TRUE.equals(needDataTypePrinted)) { + List 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 fa80c8feb2014..8aef206c61b5a 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 void init() @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 list = new ArrayList<>(threadNum); for (int i = 0; i < threadNum; i++) { Thread thread = new Thread(getAsyncImportRunnable()); @@ -85,16 +134,21 @@ public void asyncImportDataFiles() { try { thread.join(); } catch (InterruptedException e) { + importFailed.set(true); Thread.currentThread().interrupt(); ioTPrinter.println("ImportData thread join interrupted: " + e.getMessage()); } }); - ioTPrinter.println(Constants.IMPORT_COMPLETELY); + if (!importFailed.get()) { + ioTPrinter.println(Constants.IMPORT_COMPLETELY); + } + return !importFailed.get(); } } - public static void init(AbstractImportData instance) { - instance.new ThreadManager().asyncImportDataFiles(); + public static boolean init(AbstractImportData instance) { + importFailed.set(false); + return instance.new ThreadManager().asyncImportDataFiles(); } protected abstract void importFromSqlFile(File file); @@ -111,6 +165,7 @@ protected void processSuccessFile(String file) { processingLoadSuccessfulFileSuccessfulNum.increment(); ioTPrinter.println("Processed success file [ " + file + " ] successfully!"); } catch (final Exception processSuccessException) { + importFailed.set(true); ioTPrinter.println( "Failed to process success file [ " + file @@ -120,106 +175,11 @@ protected void processSuccessFile(String file) { } } - protected void processingFile(final String file, final boolean isSuccess) { - final String relativePath = file.substring(ImportTsFileScanTool.getSourceFullPathLength() + 1); - final Path sourcePath = Paths.get(file); - - 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 @@ protected void processFailFile(final String filePath, final Exception e) { } 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 322f10e059b66..63a2fdf45c723 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 static void main(String[] args) { ioTPrinter.println(sql); String[] values = sql.trim().split(";"); for (int i = 0; i < values.length; i++) { - exportData.exportBySql(values[i], i); + if (!exportData.exportBySql(values[i], i)) { + exitCode = Constants.CODE_ERROR; + } } } else { - exportData.exportBySql(null, 0); + if (!exportData.exportBySql(null, 0)) { + exitCode = Constants.CODE_ERROR; + } } } else { String[] values = queryCommand.trim().split(";"); for (int i = 0; i < values.length; i++) { - exportData.exportBySql(values[i], i); + if (!exportData.exportBySql(values[i], i)) { + exitCode = Constants.CODE_ERROR; + } } } } catch (IOException e) { 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 1694e86c9f09e..d60b5d559b62b 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.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.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 tables = new ArrayList<>(); + private static List allTables = new ArrayList<>(); private static long processedRows; private static long lastPrintTime; @@ -85,17 +99,20 @@ public void init() throws IoTDBConnectionException, StatementExecutionException 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 void init() throws IoTDBConnectionException, StatementExecutionException } @Override - public void exportBySql(String sql, int index) { + public boolean exportBySql(String sql, int index) { List 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 void exportBySql(String sql, int index) { } 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 void exportBySql(String sql, int index) { } } 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 sources = new ArrayList<>(); + try { + new RelationalSqlBaseVisitor() { + @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 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 @@ private void exportToSqlFile(SessionDataSet sessionDataSet, String table, String lastPrintTime = 0; StringBuilder sqlBuilder; List 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 columnTypeList = iterator.getColumnTypeList(); int totalColumns = columnTypeList.size(); @@ -185,11 +288,7 @@ private void exportToSqlFile(SessionDataSet sessionDataSet, String table, String } 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 @@ private void exportToCsvFile(SessionDataSet sessionDataSet, String filePath) throws IOException, IoTDBConnectionException, StatementExecutionException { processedRows = 0; lastPrintTime = 0; - List headers = sessionDataSet.getColumnNames(); + List headers = getCsvHeaders(sessionDataSet, false); int fileIndex = 0; SessionDataSet.DataIterator iterator = sessionDataSet.iterator(); List columnTypeList = iterator.getColumnTypeList(); @@ -363,7 +462,8 @@ private List collectSchemas( List columnNames, Map deviceColumnIndices, String table) throws IoTDBConnectionException, StatementExecutionException { List 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 fd8aeb0cbe352..bab4fa88b929c 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 void init() throws IoTDBConnectionException, StatementExecutionException, } @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 void exportBySql(String sql, int index) { ioTPrinter.println("Export completely!cost: " + (end - start) + " ms."); } } else { - List headers = new ArrayList<>(); - List names = sessionDataSet.getColumnNames(); - List 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 @@ private void exportToSqlFileWithAlignDevice(SessionDataSet sessionDataSet, Strin processedRows = 0; lastPrintTime = 0; - List measurementNames = sessionDataSet.getColumnNames(); + List measurementNames = new ArrayList<>(sessionDataSet.getColumnNames()); if (CollectionUtils.isEmpty(measurementNames) || measurementNames.size() <= 1) { return; } else { @@ -171,12 +159,7 @@ private void exportToSqlFileWithAlignDevice(SessionDataSet sessionDataSet, Strin } List 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 @@ private void exportToSqlFileWithoutAlign(SessionDataSet sessionDataSet, String f throws IoTDBConnectionException, StatementExecutionException, IOException { processedRows = 0; lastPrintTime = 0; - List headers = sessionDataSet.getColumnNames(); - List 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> deviceColumnIndices = new LinkedHashMap<>(); + Map 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 columnTypeList = iterator.getColumnTypeList(); - int totalColumns = measurementNames.size(); + List 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 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> entry : deviceColumnIndices.entrySet()) { + List names = new ArrayList<>(); + List 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 @@ private void exportToCsvFile(SessionDataSet sessionDataSet, String filePath) throws IOException, IoTDBConnectionException, StatementExecutionException { processedRows = 0; lastPrintTime = 0; - List headers = sessionDataSet.getColumnNames(); + List headers = getCsvHeaders(sessionDataSet, true); int fileIndex = 0; SessionDataSet.DataIterator iterator = sessionDataSet.iterator(); List 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 9525c17160687..8686352cb5a58 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 @@ private static int importFromTargetPathAsync() { 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 f65b559a9410f..d866efbd24252 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.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 @@ private static String extractDbFromSql(String sql) { @SuppressWarnings("java:S2259") protected void importFromSqlFile(File file) { ArrayList> failedRecords = new ArrayList<>(); + boolean fileFailed = false; String failedFilePath; if (failedFileDirectory == null) { failedFilePath = file.getAbsolutePath() + ".failed"; @@ -210,18 +212,25 @@ protected void importFromSqlFile(File file) { "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 @@ protected void importFromSqlFile(File file) { 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 @@ protected void importFromTsFile(File file) { 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 headerNames = csvRecords.getHeaderNames(); Stream 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 @@ protected void importFromCsvFile(File file) { } 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 headerNames, List records, String failedFilePath) { Map headerTypeMap = new HashMap<>(); Map 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> failedRecords = new ArrayList<>(); List schemas = new ArrayList<>(); @@ -367,41 +386,60 @@ protected void writeData( 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 headerNames, Map headerTypeMap, Map headerNameMap) { - String regex = "(?<=\\()\\S+(?=\\))"; + String regex = "\\(([^()]+)\\)$"; Pattern pattern = Pattern.compile(regex); + List 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 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 c976e847740d2..35641d3d50198 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 @@ protected Runnable getAsyncImportRunnable() { @SuppressWarnings("java:S2259") protected void importFromSqlFile(File file) { ArrayList> failedRecords = new ArrayList<>(); + boolean fileFailed = false; String failedFilePath; if (failedFileDirectory == null) { failedFilePath = file.getAbsolutePath() + ".failed"; @@ -115,19 +116,26 @@ protected void importFromSqlFile(File file) { 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 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 @@ protected void importFromTsFile(File file) { 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 headerNames = csvRecords.getHeaderNames(); Stream 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 @@ protected void importFromCsvFile(File file) { } 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 @@ private static void writeAndEmptyDataSet( List> valuesList, List> 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 @@ private static void writeAndEmptyDataSet( List> valuesList, List> 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 97ff8d164f359..3b9e36a79659b 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 static int importFromTargetPath() { 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 @@ private static void processSetParams() { successDir, successOperation, failDir, failOperation); } - public static void asyncImportTsFiles() { + public static int asyncImportTsFiles() { + ImportTsFileBase.resetCounters(); final List list = new ArrayList<>(threadNum); for (int i = 0; i < threadNum; i++) { final Thread thread = @@ -433,5 +433,8 @@ public static void asyncImportTsFiles() { 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 f1fd3bf2d30f5..634eda9741d20 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 @@ 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 @@ protected static void processSuccessFile(final String filePath) { } } - 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 4891f42548616..f31e7bc72e95b 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 static int getSourceFullPathLength() { : 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 0000000000000..ff4a59ada9119 --- /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 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 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 timestamps = new ArrayList<>(List.of(1L)); + List> types = new ArrayList<>(List.of(List.of(TSDataType.INT32))); + List> values = new ArrayList<>(List.of(List.of(11))); + List> 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 timestamps = new ArrayList<>(List.of(1L)); + List> types = new ArrayList<>(List.of(List.of(TSDataType.INT32))); + List> values = new ArrayList<>(List.of(List.of(11))); + List> 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 devices = new ArrayList<>(List.of("root.a.d1")); + List timestamps = new ArrayList<>(List.of(1L)); + List> types = new ArrayList<>(List.of(List.of(TSDataType.INT32))); + List> values = new ArrayList<>(List.of(List.of(11))); + List> 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 types = new LinkedHashMap<>(); + types.put("id", TSDataType.STRING); + types.put("v", TSDataType.DOUBLE); + Map 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 types = new LinkedHashMap<>(); + types.put("value", TSDataType.DOUBLE); + types.put("value(double)", TSDataType.DOUBLE); + Map 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 types = new LinkedHashMap<>(); + types.put("value", TSDataType.DOUBLE); + types.put("value(double)", TSDataType.DOUBLE); + Map 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 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 names, List 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 0000000000000..481f678b52e42 --- /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 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 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 f239b09bb489f..06d489a2d30c2 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 0b070b51fbc37..95c9e5c29e983 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 56518c78600a4..89d2feee0f6e2 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 9a082ac05ff6e..c7f2978143349 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 a4ddd93b6ffcf..930da21d1cc13 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%