From fc069b4811156553518920bf8a7c7986c9eb1b5d Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:35:17 +0800 Subject: [PATCH 1/3] Fix import-data and export-data correctness and error reporting --- .../apache/iotdb/cli/i18n/CliMessages.java | 7 + .../apache/iotdb/cli/i18n/CliMessages.java | 7 + .../iotdb/tool/data/AbstractExportData.java | 29 +- .../iotdb/tool/data/AbstractImportData.java | 179 +++--- .../apache/iotdb/tool/data/ExportData.java | 12 +- .../iotdb/tool/data/ExportDataTable.java | 151 ++++- .../iotdb/tool/data/ExportDataTree.java | 131 ++-- .../apache/iotdb/tool/data/ImportData.java | 3 +- .../iotdb/tool/data/ImportDataTable.java | 29 +- .../iotdb/tool/data/ImportDataTree.java | 71 +-- .../tool/tsfile/ImportTsFileScanTool.java | 8 + .../tool/data/DataToolRegressionTest.java | 570 ++++++++++++++++++ scripts/tools/windows/export-data.bat | 4 +- scripts/tools/windows/import-data.bat | 4 +- 14 files changed, 939 insertions(+), 266 deletions(-) create mode 100644 iotdb-client/cli/src/test/java/org/apache/iotdb/tool/data/DataToolRegressionTest.java 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..8edafac53db7c 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 @@ -77,6 +77,13 @@ public final class CliMessages { "Invalid args: Required values for option table not provided."; private CliMessages() {} + public static final String EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269 = + "Failed to insert CSV batch: %s"; + public static final String MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6 = + "[%s] - Unexpected error occurred: %s"; + public static final String + EXCEPTION_SQL_TSFILE_EXPORT_REQUIRES_A_QUERY_FROM_A_SINGLE_TABLE_IN_DATABASE_ARG_C9EB3768 = + "SQL/TsFile export requires a query from a single table in database %s."; // --------------------------------------------------------------------------- // Additional auto-collected messages // --------------------------------------------------------------------------- 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..155ba41b7dba6 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 @@ -77,6 +77,13 @@ public final class CliMessages { "参数无效：未提供 table 选项的必填值。"; private CliMessages() {} + public static final String EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269 = + "CSV 批次写入失败：%s"; + public static final String MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6 = + "[%s] - 发生意外错误：%s"; + public static final String + EXCEPTION_SQL_TSFILE_EXPORT_REQUIRES_A_QUERY_FROM_A_SINGLE_TABLE_IN_DATABASE_ARG_C9EB3768 = + "SQL/TsFile 导出要求查询仅来自数据库 %s 中的一张表。"; // --------------------------------------------------------------------------- // Additional auto-collected messages // --------------------------------------------------------------------------- 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..45ac89eec3c2c 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,13 +19,16 @@ package org.apache.iotdb.tool.data; +import org.apache.iotdb.cli.i18n.CliMessages; import org.apache.iotdb.cli.utils.IoTPrinter; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.iotdb.tool.common.Constants; +import org.apache.iotdb.tool.common.ImportTsFileOperation; import org.apache.iotdb.tool.tsfile.ImportTsFileScanTool; import java.io.File; +import java.io.IOException; import java.nio.file.FileAlreadyExistsException; import java.nio.file.Files; import java.nio.file.Path; @@ -34,10 +37,39 @@ import java.util.ArrayList; import java.util.List; import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; public abstract class AbstractImportData extends AbstractDataTool implements Runnable { private static final IoTPrinter ioTPrinter = new IoTPrinter(System.out); + private static final AtomicBoolean importFailed = new AtomicBoolean(); + + @FunctionalInterface + protected interface DataInsertion { + void insert() throws IoTDBConnectionException, StatementExecutionException; + } + + /** Retry connection failures without discarding the caller's batch when insertion fails. */ + protected static void executeWithRetry(DataInsertion insertion, int retryTime) { + for (int attempt = 0; ; attempt++) { + try { + insertion.insert(); + return; + } catch (IoTDBConnectionException e) { + if (attempt >= retryTime) { + throw new IllegalStateException( + String.format( + CliMessages.EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269, e.getMessage()), + e); + } + } catch (StatementExecutionException e) { + throw new IllegalStateException( + String.format( + CliMessages.EXCEPTION_FAILED_TO_INSERT_CSV_BATCH_ARG_159E0269, e.getMessage()), + e); + } + } + } public abstract void init() throws InterruptedException, IoTDBConnectionException, StatementExecutionException; @@ -65,15 +97,20 @@ public void run() { } } } catch (Exception e) { + importFailed.set(true); + loadFileFailedNum.increment(); ioTPrinter.println( - String.format("[%s] - Unexpected error occurred: %s", file.getName(), e.getMessage())); + String.format( + CliMessages.MESSAGE_ARG_UNEXPECTED_ERROR_OCCURRED_ARG_E7255EE6, + file == null ? filePath : file.getName(), + e.getMessage())); } } protected abstract Runnable getAsyncImportRunnable(); protected class ThreadManager { - public void asyncImportDataFiles() { + public boolean asyncImportDataFiles() { List list = new ArrayList<>(threadNum); for (int i = 0; i < threadNum; i++) { Thread thread = new Thread(getAsyncImportRunnable()); @@ -85,16 +122,21 @@ public void asyncImportDataFiles() { try { thread.join(); } catch (InterruptedException e) { + importFailed.set(true); Thread.currentThread().interrupt(); ioTPrinter.println("ImportData thread join interrupted: " + e.getMessage()); } }); - ioTPrinter.println(Constants.IMPORT_COMPLETELY); + if (!importFailed.get()) { + ioTPrinter.println(Constants.IMPORT_COMPLETELY); + } + return !importFailed.get(); } } - public static void init(AbstractImportData instance) { - instance.new ThreadManager().asyncImportDataFiles(); + public static boolean init(AbstractImportData instance) { + importFailed.set(false); + return instance.new ThreadManager().asyncImportDataFiles(); } protected abstract void importFromSqlFile(File file); @@ -111,6 +153,7 @@ protected void processSuccessFile(String file) { processingLoadSuccessfulFileSuccessfulNum.increment(); ioTPrinter.println("Processed success file [ " + file + " ] successfully!"); } catch (final Exception processSuccessException) { + importFailed.set(true); ioTPrinter.println( "Failed to process success file [ " + file @@ -120,103 +163,48 @@ protected void processSuccessFile(String file) { } } - protected void processingFile(final String file, final boolean isSuccess) { - final String relativePath = file.substring(ImportTsFileScanTool.getSourceFullPathLength() + 1); - final Path sourcePath = Paths.get(file); - + protected void processingFile(final String file, final boolean isSuccess) throws IOException { + final String relativePath = ImportTsFileScanTool.getRelativePath(file); final String target = - isSuccess - ? successDir - : failDir + File.separator + relativePath.replace(File.separator, "_"); - final Path targetPath = Paths.get(target); - - final String RESOURCE = ".resource"; - Path sourceResourcePath = Paths.get(sourcePath + RESOURCE); - sourceResourcePath = Files.exists(sourceResourcePath) ? sourceResourcePath : null; - final Path targetResourcePath = Paths.get(target + RESOURCE); - - final String MODS = ".mods"; - Path sourceModsPath = Paths.get(sourcePath + MODS); - sourceModsPath = Files.exists(sourceModsPath) ? sourceModsPath : null; - final Path targetModsPath = Paths.get(target + MODS); - - switch (isSuccess ? successOperation : failOperation) { + (isSuccess ? successDir : failDir) + + File.separator + + relativePath.replace(File.separator, "_"); + final ImportTsFileOperation operation = isSuccess ? successOperation : failOperation; + processFile(Paths.get(file), Paths.get(target), operation); + for (String suffix : new String[] {".resource", ".mods"}) { + Path sidecar = Paths.get(file + suffix); + if (Files.exists(sidecar)) { + processFile( + sidecar, + Paths.get(target + suffix), + operation == ImportTsFileOperation.HARDLINK ? ImportTsFileOperation.CP : operation); + } + } + } + + private static void processFile(Path source, Path target, ImportTsFileOperation operation) + throws IOException { + switch (operation) { case DELETE: - { - try { - Files.deleteIfExists(sourcePath); - if (null != sourceResourcePath) { - Files.deleteIfExists(sourceResourcePath); - } - if (null != sourceModsPath) { - Files.deleteIfExists(sourceModsPath); - } - } catch (final Exception e) { - ioTPrinter.println(String.format("Failed to delete file: %s", e.getMessage())); - } - break; - } + Files.deleteIfExists(source); + break; case CP: - { - try { - Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING); - if (null != sourceResourcePath) { - Files.copy( - sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING); - } - if (null != sourceModsPath) { - Files.copy(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING); - } - } catch (final Exception e) { - ioTPrinter.println(String.format("Failed to copy file: %s", e.getMessage())); - } - break; - } + Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING); + break; case HARDLINK: - { - try { - Files.createLink(targetPath, sourcePath); - } catch (FileAlreadyExistsException e) { - ioTPrinter.println("Hardlink already exists: " + e.getMessage()); - } catch (final Exception e) { - try { - Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING); - } catch (final Exception copyException) { - ioTPrinter.println( - String.format("Failed to copy file: %s", copyException.getMessage())); - } + try { + Files.createLink(target, source); + } catch (FileAlreadyExistsException e) { + if (!Files.isSameFile(source, target)) { + Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING); } - - try { - if (null != sourceResourcePath) { - Files.copy( - sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING); - } - if (null != sourceModsPath) { - Files.copy(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING); - } - } catch (final Exception e) { - ioTPrinter.println( - String.format("Failed to copy resource or mods file: %s", e.getMessage())); - } - break; + } catch (IOException | UnsupportedOperationException e) { + Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING); } + break; case MV: - { - try { - Files.move(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING); - if (null != sourceResourcePath) { - Files.move( - sourceResourcePath, targetResourcePath, StandardCopyOption.REPLACE_EXISTING); - } - if (null != sourceModsPath) { - Files.move(sourceModsPath, targetModsPath, StandardCopyOption.REPLACE_EXISTING); - } - } catch (final Exception e) { - ioTPrinter.println(String.format("Failed to move file: %s", e.getMessage())); - } - break; - } + Files.move(source, target, StandardCopyOption.REPLACE_EXISTING); + break; default: break; } @@ -232,6 +220,7 @@ protected void processFailFile(final String filePath, final Exception e) { } loadFileFailedNum.increment(); + importFailed.set(true); ioTPrinter.println("Failed to import [ " + filePath + " ] file: " + e.getMessage()); try { 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..7bed6a3292d27 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; @@ -85,6 +98,7 @@ public void init() throws IoTDBConnectionException, StatementExecutionException System.exit(1); } sessionDataSet = tableSession.executeQueryStatement("show tables", timeout); + tables.clear(); while (sessionDataSet.hasNext()) { tables.add(sessionDataSet.next().getField(0).getStringValue()); } @@ -103,9 +117,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 +130,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 +147,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 (!tables.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 +257,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 +285,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 +350,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 +459,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..25ad221a76a79 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))); + processedRows++; + if (++currentLines >= linesPerFile) { + writer.close(); + writer = null; + fileIndex++; + currentLines = 0; + } + } + 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..3b31c9613d586 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; @@ -370,16 +371,13 @@ protected void writeData( } 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( @@ -389,17 +387,20 @@ private void parseHeaders( String regex = "(?<=\\()\\S+(?=\\))"; Pattern pattern = Pattern.compile(regex); for (String headerName : headerNames) { - headerName = headerName.toLowerCase(); Matcher matcher = pattern.matcher(headerName); String type; String headerNameWithoutType; if (matcher.find()) { type = matcher.group(); - headerNameWithoutType = headerName.replace("(" + type + ")", "").replaceAll("\\s+", ""); + headerNameWithoutType = + headerName + .replace("(" + type + ")", "") + .replaceAll("\\s+", "") + .toLowerCase(Locale.ROOT); headerNameMap.put(headerNameWithoutType, headerName); - headerTypeMap.put(headerNameWithoutType, getType(type)); + headerTypeMap.put(headerNameWithoutType, getType(type.toUpperCase(Locale.ROOT))); } else { - headerNameMap.put(headerName, headerName); + headerNameMap.put(headerName.toLowerCase(Locale.ROOT), headerName); } } } 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..dae8e1b6ee9cb 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 @@ -401,26 +401,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 +425,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/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..f7a378e3e8199 --- /dev/null +++ b/iotdb-client/cli/src/test/java/org/apache/iotdb/tool/data/DataToolRegressionTest.java @@ -0,0 +1,570 @@ +/* + * 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.session.Session; +import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.tool.common.ImportTsFileOperation; +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.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.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, + 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 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 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 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))); + } + if (operation == ImportTsFileOperation.HARDLINK) { + new ImportDataTree().processingFile(source.toString(), true); + assertEquals("content", Files.readString(archive.resolve("nested_one.tsfile"))); + } + 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 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 testSingleFileRelativeSourceArchival() throws Exception { + Path source = directory.resolve("one.tsfile"); + Files.writeString(source, "source"); + Path archive = Files.createDirectories(directory.resolve("success")); + successDir = archive.toString(); + successOperation = ImportTsFileOperation.CP; + ImportTsFileScanTool.setSourceFullPath( + Path.of("").toAbsolutePath().relativize(source).toString()); + new ImportDataTree().processingFile(source.toString(), true); + assertEquals("source", Files.readString(archive.resolve("one.tsfile"))); + } + + 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))); + 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/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/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% From 126b3b6e03428e39576c8207d384f67b007dfd7e Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:12:40 +0800 Subject: [PATCH 2/3] Declare relational grammar dependency for CLI export --- iotdb-client/cli/pom.xml | 5 +++++ 1 file changed, 5 insertions(+) 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 From e5db31f1a860c19ba1a872f0b222e0612f5bb746 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:33:06 +0800 Subject: [PATCH 3/3] Fix Windows relative source archival test --- .../tool/data/DataToolRegressionTest.java | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) 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 index f7a378e3e8199..ce95f6100d757 100644 --- 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 @@ -493,15 +493,19 @@ public void testArchivalErrorsArePropagated() throws Exception { @Test public void testSingleFileRelativeSourceArchival() throws Exception { - Path source = directory.resolve("one.tsfile"); - Files.writeString(source, "source"); - Path archive = Files.createDirectories(directory.resolve("success")); - successDir = archive.toString(); - successOperation = ImportTsFileOperation.CP; - ImportTsFileScanTool.setSourceFullPath( - Path.of("").toAbsolutePath().relativize(source).toString()); - new ImportDataTree().processingFile(source.toString(), true); - assertEquals("source", Files.readString(archive.resolve("one.tsfile"))); + 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 {