From 0f2d66b5ca0fa88641b8047c91cf9fd3459e16e3 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Thu, 13 Aug 2026 16:36:54 +0800 Subject: [PATCH 1/4] ver1 --- .../iotdb/db/i18n/ImportWALMessages.java | 80 ++ .../iotdb/db/i18n/ImportWALMessages.java | 79 ++ .../dataregion/wal/io/WALReader.java | 9 + .../org/apache/iotdb/db/tools/ImportWAL.java | 977 ++++++++++++++++++ .../apache/iotdb/db/tools/ImportWALTest.java | 506 +++++++++ scripts/tools/import-wal.sh | 52 + scripts/tools/windows/import-wal.bat | 43 + 7 files changed, 1746 insertions(+) create mode 100644 iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java create mode 100644 iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java create mode 100644 scripts/tools/import-wal.sh create mode 100644 scripts/tools/windows/import-wal.bat diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java new file mode 100644 index 0000000000000..2ed53656ec0f0 --- /dev/null +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -0,0 +1,80 @@ +/* + * 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.db.i18n; + +/** Compile-time i18n constants for the WAL import tool (English). */ +public final class ImportWALMessages { + + public static final String MESSAGE_IMPORT_WAL_5E42804E = "import-wal"; + public static final String + MESSAGE_PATH_OF_A_WAL_FILE_OR_A_DIRECTORY_CONTAINING_WAL_FILES_473D0554 = + "Path of a WAL file or a directory containing WAL files."; + public static final String MESSAGE_TARGET_IOTDB_HOST_DEFAULT_127_0_0_1_3729156F = + "Target IoTDB host. Default: 127.0.0.1."; + public static final String MESSAGE_TARGET_IOTDB_RPC_PORT_DEFAULT_6667_FC0D345D = + "Target IoTDB RPC port. Default: 6667."; + public static final String MESSAGE_TARGET_IOTDB_USERNAME_DEFAULT_ROOT_EB91453B = + "Target IoTDB username. Default: root."; + public static final String + MESSAGE_TARGET_IOTDB_PASSWORD_PROMPTED_INTERACTIVELY_IF_OMITTED_29681961 = + "Target IoTDB password. Prompted interactively if omitted."; + public static final String MESSAGE_PASSWORD_PROMPT_F2D0E794 = "Password: "; + public static final String MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C = + "Target database for table-model WAL entries."; + public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = + "Print this help message."; + public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "Argument error: %s"; + public static final String MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA = "WAL import failed: %s"; + public static final String EXCEPTION_SOURCE_PATH_DOES_NOT_EXIST_ARG_7C806CA2 = + "Source path does not exist: %s"; + public static final String EXCEPTION_SOURCE_FILE_IS_NOT_A_WAL_FILE_ARG_14A43F76 = + "Source file is not a WAL file: %s"; + public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = + "No WAL files found under: %s"; + public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "Invalid port: %s"; + public static final String + MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = + "Replayed %d operations from %d WAL files; skipped %d entries."; + public static final String + MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F = + "Progress: %d/%d WAL files completed, %d/%d bytes (%.1f%%), elapsed %.1f s, rate %.1f MB/s."; + public static final String + MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D = + "Import duration: %.1f s; total size: %d bytes; average rate: %.1f MB/s."; + public static final String EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9 = + "Failed to replay WAL file %s at offset %d: %s"; + public static final String EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726 = + "Table-model WAL entries require -db/--database."; + public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = + "Unsupported WAL operation: %s"; + public static final String MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650 = + "Unsupported WAL operation: %s. Skip this entry? [y/N]: "; + public static final String EXCEPTION_INSERT_NODE_ARG_CONTAINS_NO_REPLAYABLE_DATA_5DA13453 = + "Insert node %s contains no replayable data."; + public static final String EXCEPTION_UNSUPPORTED_SNAPSHOT_DATA_TYPE_ARG_7A32D312 = + "Unsupported snapshot data type: %s"; + public static final String EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5 = + "The WAL file is truncated or corrupted."; + public static final String + EXCEPTION_PASSWORD_WAS_NOT_PROVIDED_AND_INTERACTIVE_INPUT_IS_UNAVAILABLE_40F42BCD = + "Password was not provided and interactive input is unavailable. Specify -pw/--password."; + + private ImportWALMessages() {} +} diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java new file mode 100644 index 0000000000000..c8dc2266dc8a1 --- /dev/null +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -0,0 +1,79 @@ +/* + * 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.db.i18n; + +/** WAL 导入工具的编译期国际化常量（中文）。 */ +public final class ImportWALMessages { + + public static final String MESSAGE_IMPORT_WAL_5E42804E = "import-wal"; + public static final String + MESSAGE_PATH_OF_A_WAL_FILE_OR_A_DIRECTORY_CONTAINING_WAL_FILES_473D0554 = + "WAL 文件或包含 WAL 文件的目录路径。"; + public static final String MESSAGE_TARGET_IOTDB_HOST_DEFAULT_127_0_0_1_3729156F = + "目标 IoTDB 主机。默认：127.0.0.1。"; + public static final String MESSAGE_TARGET_IOTDB_RPC_PORT_DEFAULT_6667_FC0D345D = + "目标 IoTDB RPC 端口。默认：6667。"; + public static final String MESSAGE_TARGET_IOTDB_USERNAME_DEFAULT_ROOT_EB91453B = + "目标 IoTDB 用户名。默认：root。"; + public static final String + MESSAGE_TARGET_IOTDB_PASSWORD_PROMPTED_INTERACTIVELY_IF_OMITTED_29681961 = + "目标 IoTDB 密码。未提供时将交互式询问。"; + public static final String MESSAGE_PASSWORD_PROMPT_F2D0E794 = "密码："; + public static final String MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C = + "表模型 WAL 条目的目标数据库。"; + public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "打印帮助信息。"; + public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "参数错误：%s"; + public static final String MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA = "WAL 导入失败：%s"; + public static final String EXCEPTION_SOURCE_PATH_DOES_NOT_EXIST_ARG_7C806CA2 = + "源路径不存在：%s"; + public static final String EXCEPTION_SOURCE_FILE_IS_NOT_A_WAL_FILE_ARG_14A43F76 = + "源文件不是 WAL 文件：%s"; + public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = + "路径下未找到 WAL 文件：%s"; + public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "无效端口：%s"; + public static final String + MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = + "已重放 %d 个操作（来自 %d 个 WAL 文件）；跳过 %d 个条目。"; + public static final String + MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F = + "进度：已完成 %d/%d 个 WAL 文件，已处理 %d/%d 字节（%.1f%%），耗时 %.1f 秒，速率 %.1f MB/s。"; + public static final String + MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D = + "导入耗时：%.1f 秒；文件总大小：%d 字节；平均速率：%.1f MB/s。"; + public static final String EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9 = + "重放 WAL 文件 %s 时失败，偏移量 %d：%s"; + public static final String EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726 = + "表模型 WAL 条目要求指定 -db/--database。"; + public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = + "不支持的 WAL 操作：%s"; + public static final String MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650 = + "不支持的 WAL 操作：%s。是否跳过此条目？[y/N]："; + public static final String EXCEPTION_INSERT_NODE_ARG_CONTAINS_NO_REPLAYABLE_DATA_5DA13453 = + "Insert node %s 不包含可重放数据。"; + public static final String EXCEPTION_UNSUPPORTED_SNAPSHOT_DATA_TYPE_ARG_7A32D312 = + "Unsupported snapshot data type: %s"; + public static final String EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5 = + "WAL 文件被截断或已损坏。"; + public static final String + EXCEPTION_PASSWORD_WAS_NOT_PROVIDED_AND_INTERACTIVE_INPUT_IS_UNAVAILABLE_40F42BCD = + "未提供密码且当前环境不支持交互式输入，请指定 -pw/--password。"; + + private ImportWALMessages() {} +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/io/WALReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/io/WALReader.java index befcf58c6323d..0a2c994eb549c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/io/WALReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/io/WALReader.java @@ -72,6 +72,11 @@ public boolean hasNext() { return false; } try { + // An active WAL has no end marker until its writer closes. Reaching EOF exactly between + // entries is therefore valid, while EOF during deserialization still marks a partial entry. + if (walInputStream.available() == 0) { + return false; + } nextEntry = WALEntry.deserialize(logStream); if (nextEntry.getType() == WALEntryType.WAL_FILE_INFO_END_MARKER) { nextEntry = null; @@ -98,6 +103,10 @@ public long getWALCurrentReadOffset() throws IOException { return walInputStream.getFileCurrentPos(); } + public boolean isFileCorrupted() { + return fileCorrupted; + } + /** * Like {@link Iterator#next()}. * diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java new file mode 100644 index 0000000000000..f8316499bdb5f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java @@ -0,0 +1,977 @@ +/* + * 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.db.tools; + +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.db.i18n.ImportWALMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ObjectNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalDeleteDataNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode; +import org.apache.iotdb.db.storageengine.dataregion.memtable.AlignedWritableMemChunk; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunk; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunkGroup; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryType; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALReader; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALFileUtils; +import org.apache.iotdb.db.subscription.broker.consensus.ConsensusLogToTabletConverter; +import org.apache.iotdb.db.subscription.columnfilter.ColumnFilterMatcher; +import org.apache.iotdb.db.utils.datastructure.AlignedTVList; +import org.apache.iotdb.db.utils.datastructure.TVList; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.session.Session; + +import org.apache.commons.cli.CommandLine; +import org.apache.commons.cli.DefaultParser; +import org.apache.commons.cli.HelpFormatter; +import org.apache.commons.cli.Option; +import org.apache.commons.cli.Options; +import org.apache.commons.cli.ParseException; +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.utils.BitMap; +import org.apache.tsfile.utils.DateUtils; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; + +import java.io.Console; +import java.io.IOException; +import java.io.PrintStream; +import java.io.PrintWriter; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +public class ImportWAL { + + private static final int CODE_OK = 0; + private static final int CODE_ERROR = 1; + private static final String DEFAULT_HOST = "127.0.0.1"; + private static final int DEFAULT_PORT = 6667; + private static final String DEFAULT_USER = "root"; + private static final int SNAPSHOT_TABLET_ROW_LIMIT = 1024; + + private ImportWAL() {} + + public static void main(final String[] args) { + System.exit(run(args, System.out, System.err)); + } + + static int run(final String[] args, final PrintStream out, final PrintStream err) { + final Options options = createOptions(); + if (containsHelpOption(args)) { + printHelp(options, out); + return CODE_OK; + } + + final CommandLine commandLine; + try { + commandLine = new DefaultParser().parse(options, args); + } catch (final ParseException e) { + err.printf(ImportWALMessages.MESSAGE_ARGUMENT_ERROR_ARG_A9767F62, e.getMessage()); + err.println(); + printHelp(options, err); + return CODE_ERROR; + } + + try { + final Path source = Paths.get(commandLine.getOptionValue("file")); + final List walFiles = collectWALFiles(source); + final String database = commandLine.getOptionValue("database"); + final String password = getPassword(commandLine); + final Session treeSession = + createSession( + commandLine.getOptionValue("host", DEFAULT_HOST), + parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), + commandLine.getOptionValue("username", DEFAULT_USER), + password, + null); + final Session tableSession = + database == null + ? null + : createSession( + commandLine.getOptionValue("host", DEFAULT_HOST), + parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), + commandLine.getOptionValue("username", DEFAULT_USER), + password, + database); + try { + treeSession.open(false); + if (tableSession != null) { + tableSession.open(false); + } + final ReplayStatistics statistics = + replayWALFiles(walFiles, new WALReplayer(treeSession, tableSession, database), out); + out.printf( + ImportWALMessages + .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, + statistics.replayedOperationCount, + walFiles.size(), + statistics.skippedEntryCount); + out.println(); + out.printf( + ImportWALMessages + .MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D, + statistics.getElapsedSeconds(), + statistics.getTotalBytes(), + statistics.getAverageRateMbPerSecond()); + out.println(); + } finally { + closeSession(tableSession); + closeSession(treeSession); + } + return CODE_OK; + } catch (final Exception e) { + err.printf(ImportWALMessages.MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA, e.getMessage()); + err.println(); + return CODE_ERROR; + } + } + + private static Options createOptions() { + final Options options = new Options(); + options.addOption( + Option.builder("f") + .longOpt("file") + .hasArg() + .required() + .desc( + ImportWALMessages + .MESSAGE_PATH_OF_A_WAL_FILE_OR_A_DIRECTORY_CONTAINING_WAL_FILES_473D0554) + .build()); + options.addOption( + Option.builder("h") + .longOpt("host") + .hasArg() + .desc(ImportWALMessages.MESSAGE_TARGET_IOTDB_HOST_DEFAULT_127_0_0_1_3729156F) + .build()); + options.addOption( + Option.builder("p") + .longOpt("port") + .hasArg() + .desc(ImportWALMessages.MESSAGE_TARGET_IOTDB_RPC_PORT_DEFAULT_6667_FC0D345D) + .build()); + options.addOption( + Option.builder("u") + .longOpt("username") + .hasArg() + .desc(ImportWALMessages.MESSAGE_TARGET_IOTDB_USERNAME_DEFAULT_ROOT_EB91453B) + .build()); + options.addOption( + Option.builder("pw") + .longOpt("password") + .hasArg() + .desc( + ImportWALMessages + .MESSAGE_TARGET_IOTDB_PASSWORD_PROMPTED_INTERACTIVELY_IF_OMITTED_29681961) + .build()); + options.addOption( + Option.builder("db") + .longOpt("database") + .hasArg() + .desc(ImportWALMessages.MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C) + .build()); + options.addOption( + Option.builder() + .longOpt("help") + .desc(ImportWALMessages.MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A) + .build()); + return options; + } + + private static String getPassword(final CommandLine commandLine) { + return getPassword(commandLine, System.console()); + } + + static String getPassword(final CommandLine commandLine, final Console console) { + if (commandLine.hasOption("password")) { + return commandLine.getOptionValue("password"); + } + if (console == null) { + throw new IllegalArgumentException( + ImportWALMessages + .EXCEPTION_PASSWORD_WAS_NOT_PROVIDED_AND_INTERACTIVE_INPUT_IS_UNAVAILABLE_40F42BCD); + } + final char[] password = + console.readPassword(ImportWALMessages.MESSAGE_PASSWORD_PROMPT_F2D0E794); + if (password == null) { + throw new IllegalArgumentException( + ImportWALMessages + .EXCEPTION_PASSWORD_WAS_NOT_PROVIDED_AND_INTERACTIVE_INPUT_IS_UNAVAILABLE_40F42BCD); + } + return new String(password); + } + + private static boolean containsHelpOption(final String[] args) { + if (args == null) { + return false; + } + for (final String arg : args) { + if ("--help".equals(arg) || "-help".equals(arg)) { + return true; + } + } + return false; + } + + private static void printHelp(final Options options, final PrintStream stream) { + final HelpFormatter formatter = new HelpFormatter(); + formatter.setWidth(120); + formatter.printHelp( + new PrintWriter(stream, true), + 120, + ImportWALMessages.MESSAGE_IMPORT_WAL_5E42804E, + null, + options, + 2, + 2, + null, + true); + } + + private static int parsePort(final String port) { + try { + final int value = Integer.parseInt(port); + if (value <= 0 || value > 65535) { + throw new NumberFormatException(port); + } + return value; + } catch (final NumberFormatException e) { + throw new IllegalArgumentException( + String.format(ImportWALMessages.EXCEPTION_INVALID_PORT_ARG_A7CDD5AC, port), e); + } + } + + private static Session createSession( + final String host, + final int port, + final String username, + final String password, + final String database) { + final Session.Builder builder = + new Session.Builder().host(host).port(port).username(username).password(password); + if (database != null) { + builder.sqlDialect("table").database(database); + } + return builder.build(); + } + + private static void closeSession(final Session session) { + if (session == null) { + return; + } + try { + session.close(); + } catch (final IoTDBConnectionException ignored) { + // The import result has already been determined; closing failure must not hide it. + } + } + + static List collectWALFiles(final Path source) throws IOException { + if (!Files.exists(source)) { + throw new IOException( + String.format( + ImportWALMessages.EXCEPTION_SOURCE_PATH_DOES_NOT_EXIST_ARG_7C806CA2, source)); + } + if (Files.isRegularFile(source)) { + if (!isWALFile(source)) { + throw new IOException( + String.format( + ImportWALMessages.EXCEPTION_SOURCE_FILE_IS_NOT_A_WAL_FILE_ARG_14A43F76, source)); + } + return List.of(source.toAbsolutePath().normalize()); + } + if (!Files.isDirectory(source)) { + throw new IOException( + String.format( + ImportWALMessages.EXCEPTION_SOURCE_PATH_DOES_NOT_EXIST_ARG_7C806CA2, source)); + } + + final List walFiles; + try (Stream stream = Files.walk(source)) { + walFiles = + stream + .filter(Files::isRegularFile) + .filter(ImportWAL::isWALFile) + .map(path -> path.toAbsolutePath().normalize()) + .sorted(WAL_FILE_COMPARATOR) + .collect(Collectors.toList()); + } + if (walFiles.isEmpty()) { + throw new IOException( + String.format(ImportWALMessages.EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22, source)); + } + return walFiles; + } + + private static boolean isWALFile(final Path path) { + return path.getFileName().toString().toLowerCase(Locale.ROOT).endsWith(".wal"); + } + + private static final Comparator WAL_FILE_COMPARATOR = + Comparator.comparing((Path path) -> Objects.toString(path.getParent(), "")) + .thenComparingLong(ImportWAL::getWALVersion) + .thenComparing(path -> path.getFileName().toString()); + + private static long getWALVersion(final Path path) { + final String filename = path.getFileName().toString(); + return WALFileUtils.WAL_FILE_NAME_PATTERN.matcher(filename).find() + ? WALFileUtils.parseVersionId(filename) + : Long.MAX_VALUE; + } + + static ReplayStatistics replayWALFiles(final List walFiles, final WALReplayer replayer) + throws IOException { + return replayWALFiles(walFiles, replayer, null); + } + + static ReplayStatistics replayWALFiles( + final List walFiles, final WALReplayer replayer, final PrintStream progressStream) + throws IOException { + final ReplayStatistics statistics = new ReplayStatistics(); + final long startNanos = System.nanoTime(); + for (final Path walFile : walFiles) { + statistics.totalBytes += Files.size(walFile); + } + for (final Path walFile : walFiles) { + try (WALReader reader = new WALReader(walFile.toFile())) { + long offset = reader.getWALCurrentReadOffset(); + while (reader.hasNext()) { + final WALEntry entry = reader.next(); + try { + if (replayer.replay(entry)) { + statistics.replayedOperationCount++; + } else { + statistics.skippedEntryCount++; + } + } catch (final IoTDBConnectionException | StatementExecutionException e) { + throw new WALReplayException( + String.format( + ImportWALMessages + .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, + walFile, + offset, + e.getMessage()), + e); + } + offset = reader.getWALCurrentReadOffset(); + } + if (reader.isFileCorrupted()) { + throw new WALReplayException( + String.format( + ImportWALMessages + .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, + walFile, + reader.getWALCurrentReadOffset(), + ImportWALMessages.EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5), + null); + } + statistics.completedFileCount++; + statistics.processedBytes += Files.size(walFile); + statistics.elapsedNanos = System.nanoTime() - startNanos; + if (progressStream != null) { + progressStream.printf( + ImportWALMessages + .MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F, + statistics.completedFileCount, + walFiles.size(), + statistics.processedBytes, + statistics.totalBytes, + statistics.getProgressPercent(), + statistics.getElapsedSeconds(), + statistics.getAverageRateMbPerSecond()); + progressStream.println(); + } + } catch (final IOException e) { + if (e instanceof WALReplayException walReplayException) { + throw walReplayException; + } + throw new WALReplayException( + String.format( + ImportWALMessages + .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, + walFile, + 0, + e.getMessage()), + e); + } + } + return statistics; + } + + static class WALReplayer { + + private final Session treeSession; + private final Session tableSession; + private final ConsensusLogToTabletConverter converter; + private final UnsupportedEntryPrompt unsupportedEntryPrompt; + private final Map> tableTagSchemas = new HashMap<>(); + + WALReplayer( + final Session treeSession, final Session tableSession, final String tableDatabaseName) { + this( + treeSession, + tableSession, + tableDatabaseName, + createUnsupportedEntryPrompt(System.console())); + } + + WALReplayer( + final Session treeSession, + final Session tableSession, + final String tableDatabaseName, + final UnsupportedEntryPrompt unsupportedEntryPrompt) { + this.treeSession = treeSession; + this.tableSession = tableSession; + this.unsupportedEntryPrompt = unsupportedEntryPrompt; + converter = + new ConsensusLogToTabletConverter( + null, null, ColumnFilterMatcher.matchAll(), tableDatabaseName); + } + + boolean replay(final WALEntry entry) + throws IoTDBConnectionException, StatementExecutionException { + if (entry.getType() == WALEntryType.MEMORY_TABLE_SNAPSHOT + || entry.getType() == WALEntryType.OLD_MEMORY_TABLE_SNAPSHOT) { + return replayMemTableSnapshot((IMemTable) entry.getValue()); + } + if (entry.getValue() instanceof InsertNode insertNode) { + replayInsert(insertNode); + return true; + } + if (entry.getValue() instanceof DeleteDataNode deleteDataNode) { + replayTreeDelete(deleteDataNode); + return true; + } + if (entry.getValue() instanceof RelationalDeleteDataNode + || entry.getValue() instanceof ObjectNode) { + // A null prompt means no interactive console is available, so preserve fail-fast behavior. + if (unsupportedEntryPrompt != null && unsupportedEntryPrompt.shouldSkip(entry)) { + return false; + } + throw unsupportedOperation(entry); + } + return false; + } + + private static UnsupportedEntryPrompt createUnsupportedEntryPrompt(final Console console) { + if (console == null) { + return null; + } + return entry -> { + final String answer = + console.readLine( + ImportWALMessages + .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, + entry.getType()); + return isSkipConfirmation(answer); + }; + } + + static boolean isSkipConfirmation(final String answer) { + return answer != null + && ("y".equalsIgnoreCase(answer.trim()) || "yes".equalsIgnoreCase(answer.trim())); + } + + @FunctionalInterface + interface UnsupportedEntryPrompt { + + boolean shouldSkip(WALEntry entry); + } + + private static StatementExecutionException unsupportedOperation(final WALEntry entry) { + return new StatementExecutionException( + String.format( + ImportWALMessages.EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0, entry.getType())); + } + + private void replayInsert(final InsertNode node) + throws IoTDBConnectionException, StatementExecutionException { + if (node instanceof InsertRowsNode insertRowsNode + && !(node instanceof RelationalInsertRowsNode)) { + for (final InsertNode rowNode : insertRowsNode.getInsertRowNodeList()) { + replayInsert(rowNode); + } + return; + } + final List tablets = converter.convert(node); + if (tablets.isEmpty()) { + throw new StatementExecutionException( + String.format( + ImportWALMessages.EXCEPTION_INSERT_NODE_ARG_CONTAINS_NO_REPLAYABLE_DATA_5DA13453, + node.getType())); + } + final boolean tableModel = isTableModelInsert(node); + if (tableModel && tableSession == null) { + throw new StatementExecutionException( + ImportWALMessages.EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726); + } + for (final Tablet tablet : tablets) { + if (tableModel) { + tableSession.insertRelationalTablet(tablet); + } else if (node.isAligned()) { + treeSession.insertAlignedTablet(tablet); + } else { + treeSession.insertTablet(tablet); + } + } + } + + private static boolean isTableModelInsert(final InsertNode node) { + return node instanceof RelationalInsertRowNode + || node instanceof RelationalInsertRowsNode + || node instanceof RelationalInsertTabletNode; + } + + private void replayTreeDelete(final DeleteDataNode node) + throws IoTDBConnectionException, StatementExecutionException { + final List paths = new ArrayList<>(node.getPathList().size()); + for (final MeasurementPath path : node.getPathList()) { + paths.add(path.getFullPath()); + } + treeSession.deleteData(paths, node.getDeleteStartTime(), node.getDeleteEndTime()); + } + + private boolean replayMemTableSnapshot(final IMemTable memTable) + throws IoTDBConnectionException, StatementExecutionException { + if (memTable == null || memTable.isSignalMemTable()) { + return false; + } + boolean replayed = false; + for (Map.Entry deviceEntry : + memTable.getMemTableMap().entrySet()) { + final IDeviceID deviceId = deviceEntry.getKey(); + final IWritableMemChunkGroup group = deviceEntry.getValue(); + for (IWritableMemChunk chunk : group.getMemChunkMap().values()) { + if (chunk == null || chunk.isEmpty()) { + continue; + } + if (chunk instanceof AlignedWritableMemChunk) { + replayed |= replayAlignedMemChunk(deviceId, (AlignedWritableMemChunk) chunk); + } else { + replayed |= replayNonAlignedMemChunk(deviceId, chunk); + } + } + } + return replayed; + } + + private boolean replayNonAlignedMemChunk( + final IDeviceID deviceId, final IWritableMemChunk chunk) + throws IoTDBConnectionException, StatementExecutionException { + final List schemas = Collections.singletonList(chunk.getSchema()); + final boolean tableModel = deviceId.isTableModel(); + requireTableSessionIfNeeded(tableModel); + final TableTabletSchema tabletSchema = createTableTabletSchema(deviceId, schemas); + final List lists = new ArrayList<>(); + lists.addAll(chunk.getSortedList()); + lists.add(chunk.getWorkingTVList()); + boolean replayed = false; + for (TVList list : lists) { + if (list == null || list.rowCount() == 0) { + continue; + } + if (!list.isSorted()) { + list.sort(); + } + for (int start = 0; start < list.rowCount(); start += SNAPSHOT_TABLET_ROW_LIMIT) { + final int end = Math.min(start + SNAPSHOT_TABLET_ROW_LIMIT, list.rowCount()); + final Tablet tablet = + buildNonAlignedTablet(deviceId, tabletSchema, schemas, list, start, end); + sendTablet(tablet, tableModel, false); + replayed = true; + } + } + return replayed; + } + + private boolean replayAlignedMemChunk( + final IDeviceID deviceId, final AlignedWritableMemChunk chunk) + throws IoTDBConnectionException, StatementExecutionException { + final boolean tableModel = deviceId.isTableModel(); + requireTableSessionIfNeeded(tableModel); + final List schemas = chunk.getSchemaList(); + final TableTabletSchema tabletSchema = createTableTabletSchema(deviceId, schemas); + final List lists = new ArrayList<>(); + lists.addAll(chunk.getSortedList()); + lists.add(chunk.getWorkingTVList()); + boolean replayed = false; + for (AlignedTVList list : lists) { + if (list == null || list.rowCount() == 0) { + continue; + } + if (!list.isSorted()) { + list.sort(); + } + final List replayableRows = getReplayableAlignedRows(list); + for (int start = 0; start < replayableRows.size(); start += SNAPSHOT_TABLET_ROW_LIMIT) { + final int end = Math.min(start + SNAPSHOT_TABLET_ROW_LIMIT, replayableRows.size()); + final Tablet tablet = + buildAlignedTablet(deviceId, tabletSchema, schemas, list, replayableRows, start, end); + sendTablet(tablet, tableModel, true); + replayed = true; + } + } + return replayed; + } + + private static List getReplayableAlignedRows(final AlignedTVList list) { + final List replayableRows = new ArrayList<>(list.rowCount()); + for (int row = 0; row < list.rowCount(); row++) { + if (!list.isTimeDeleted(row)) { + replayableRows.add(row); + } + } + return replayableRows; + } + + private void requireTableSessionIfNeeded(final boolean tableModel) + throws StatementExecutionException { + if (tableModel && tableSession == null) { + throw new StatementExecutionException( + ImportWALMessages.EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726); + } + } + + private TableTabletSchema createTableTabletSchema( + final IDeviceID deviceId, final List fieldSchemas) + throws IoTDBConnectionException, StatementExecutionException { + if (!deviceId.isTableModel()) { + return new TableTabletSchema(fieldSchemas, null, 0); + } + final String tableName = deviceId.getTableName(); + List tagSchemas = tableTagSchemas.get(tableName); + if (tagSchemas == null) { + tagSchemas = new ArrayList<>(); + try (SessionDataSet dataSet = + tableSession.executeQueryStatement("DESCRIBE " + quoteIdentifier(tableName))) { + final SessionDataSet.DataIterator iterator = dataSet.iterator(); + while (iterator.next()) { + final String category = iterator.getString(3); + if ("TAG".equalsIgnoreCase(category)) { + tagSchemas.add( + new MeasurementSchema( + iterator.getString(1), TSDataType.valueOf(iterator.getString(2)))); + } + } + } + tableTagSchemas.put(tableName, tagSchemas); + } + final List schemas = + new ArrayList<>(tagSchemas.size() + fieldSchemas.size()); + schemas.addAll(tagSchemas); + schemas.addAll(fieldSchemas); + final List categories = new ArrayList<>(schemas.size()); + categories.addAll(Collections.nCopies(tagSchemas.size(), ColumnCategory.TAG)); + categories.addAll(Collections.nCopies(fieldSchemas.size(), ColumnCategory.FIELD)); + return new TableTabletSchema(schemas, categories, tagSchemas.size()); + } + + static String quoteIdentifier(final String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } + + private void sendTablet(final Tablet tablet, final boolean tableModel, final boolean aligned) + throws IoTDBConnectionException, StatementExecutionException { + if (tableModel) { + tableSession.insertRelationalTablet(tablet); + } else if (aligned) { + treeSession.insertAlignedTablet(tablet); + } else { + treeSession.insertTablet(tablet); + } + } + + private static Tablet buildNonAlignedTablet( + final IDeviceID deviceId, + final TableTabletSchema tabletSchema, + final List sourceSchemas, + final TVList list, + final int start, + final int end) { + final int rowCount = end - start; + final long[] times = new long[rowCount]; + final Object[] values = createValueArrays(tabletSchema.schemas, rowCount); + final BitMap[] bitMaps = new BitMap[tabletSchema.schemas.size()]; + final int fieldColumnIndex = tabletSchema.tagCount; + final TSDataType type = sourceSchemas.get(0).getType(); + for (int i = 0; i < rowCount; i++) { + final int scanIndex = start + i; + final int valueIndex = list.getValueIndex(scanIndex); + times[i] = list.getTime(scanIndex); + if (list.isNullValue(valueIndex)) { + if (bitMaps[fieldColumnIndex] == null) { + bitMaps[fieldColumnIndex] = new BitMap(rowCount); + } + bitMaps[fieldColumnIndex].mark(i); + } else { + putValue(values[fieldColumnIndex], i, type, list, scanIndex); + } + } + populateTableTags(deviceId, tabletSchema.categories, values, bitMaps, rowCount); + return tabletSchema.categories == null + ? new Tablet(deviceId.toString(), tabletSchema.schemas, times, values, bitMaps, rowCount) + : new Tablet( + deviceId.getTableName(), + tabletSchema.schemas, + tabletSchema.categories, + times, + values, + bitMaps, + rowCount); + } + + private static Tablet buildAlignedTablet( + final IDeviceID deviceId, + final TableTabletSchema tabletSchema, + final List sourceSchemas, + final AlignedTVList list, + final List replayableRows, + final int start, + final int end) { + final int rowCount = end - start; + final long[] times = new long[rowCount]; + final Object[] values = createValueArrays(tabletSchema.schemas, rowCount); + final BitMap[] bitMaps = new BitMap[tabletSchema.schemas.size()]; + final List types = list.getTsDataTypes(); + for (int i = 0; i < rowCount; i++) { + final int scanIndex = replayableRows.get(start + i); + final int valueIndex = list.getValueIndex(scanIndex); + times[i] = list.getTime(scanIndex); + for (int c = 0; c < sourceSchemas.size(); c++) { + final int targetColumnIndex = tabletSchema.tagCount + c; + if (c >= types.size() || list.isNullValue(valueIndex, c)) { + if (bitMaps[targetColumnIndex] == null) { + bitMaps[targetColumnIndex] = new BitMap(rowCount); + } + bitMaps[targetColumnIndex].mark(i); + } else { + putValue(values[targetColumnIndex], i, types.get(c), list, valueIndex, c); + } + } + } + populateTableTags(deviceId, tabletSchema.categories, values, bitMaps, rowCount); + return tabletSchema.categories == null + ? new Tablet(deviceId.toString(), tabletSchema.schemas, times, values, bitMaps, rowCount) + : new Tablet( + deviceId.getTableName(), + tabletSchema.schemas, + tabletSchema.categories, + times, + values, + bitMaps, + rowCount); + } + + private static void populateTableTags( + final IDeviceID deviceId, + final List columnCategories, + final Object[] values, + final BitMap[] bitMaps, + final int rowCount) { + if (columnCategories == null) { + return; + } + int tagSegmentIndex = 1; + for (int columnIndex = 0; columnIndex < columnCategories.size(); columnIndex++) { + if (columnCategories.get(columnIndex) != ColumnCategory.TAG) { + continue; + } + final Object segment = + tagSegmentIndex < deviceId.segmentNum() ? deviceId.segment(tagSegmentIndex) : null; + tagSegmentIndex++; + final Binary tagValue = + segment == null ? null : new Binary(segment.toString(), TSFileConfig.STRING_CHARSET); + for (int row = 0; row < rowCount; row++) { + if (tagValue == null) { + if (bitMaps[columnIndex] == null) { + bitMaps[columnIndex] = new BitMap(rowCount); + } + bitMaps[columnIndex].mark(row); + } else { + ((Binary[]) values[columnIndex])[row] = tagValue; + if (bitMaps[columnIndex] != null) { + bitMaps[columnIndex].unmark(row); + } + } + } + } + } + + private static Object createValueArray(final TSDataType type, final int rowCount) { + return switch (type) { + case BOOLEAN -> new boolean[rowCount]; + case INT32 -> new int[rowCount]; + case DATE -> new LocalDate[rowCount]; + case INT64, TIMESTAMP -> new long[rowCount]; + case FLOAT -> new float[rowCount]; + case DOUBLE -> new double[rowCount]; + case TEXT, STRING, BLOB, OBJECT -> new Binary[rowCount]; + case VECTOR, UNKNOWN -> throw unsupportedSnapshotDataType(type); + }; + } + + private static Object[] createValueArrays( + final List schemas, final int rowCount) { + final Object[] values = new Object[schemas.size()]; + for (int column = 0; column < schemas.size(); column++) { + values[column] = createValueArray(schemas.get(column).getType(), rowCount); + } + return values; + } + + private static void putValue( + final Object target, + final int targetIndex, + final TSDataType type, + final TVList list, + final int sourceIndex) { + switch (type) { + case BOOLEAN -> ((boolean[]) target)[targetIndex] = list.getBoolean(sourceIndex); + case INT32 -> ((int[]) target)[targetIndex] = list.getInt(sourceIndex); + case DATE -> + ((LocalDate[]) target)[targetIndex] = + DateUtils.parseIntToLocalDate(list.getInt(sourceIndex)); + case INT64, TIMESTAMP -> ((long[]) target)[targetIndex] = list.getLong(sourceIndex); + case FLOAT -> ((float[]) target)[targetIndex] = list.getFloat(sourceIndex); + case DOUBLE -> ((double[]) target)[targetIndex] = list.getDouble(sourceIndex); + case TEXT, STRING, BLOB, OBJECT -> + ((Binary[]) target)[targetIndex] = list.getBinary(sourceIndex); + case VECTOR, UNKNOWN -> throw unsupportedSnapshotDataType(type); + } + } + + private static void putValue( + final Object target, + final int targetIndex, + final TSDataType type, + final AlignedTVList list, + final int sourceIndex, + final int columnIndex) { + switch (type) { + case BOOLEAN -> + ((boolean[]) target)[targetIndex] = + list.getBooleanByValueIndex(sourceIndex, columnIndex); + case INT32 -> + ((int[]) target)[targetIndex] = list.getIntByValueIndex(sourceIndex, columnIndex); + case DATE -> + ((LocalDate[]) target)[targetIndex] = + DateUtils.parseIntToLocalDate(list.getIntByValueIndex(sourceIndex, columnIndex)); + case INT64, TIMESTAMP -> + ((long[]) target)[targetIndex] = list.getLongByValueIndex(sourceIndex, columnIndex); + case FLOAT -> + ((float[]) target)[targetIndex] = list.getFloatByValueIndex(sourceIndex, columnIndex); + case DOUBLE -> + ((double[]) target)[targetIndex] = list.getDoubleByValueIndex(sourceIndex, columnIndex); + case TEXT, STRING, BLOB, OBJECT -> + ((Binary[]) target)[targetIndex] = list.getBinaryByValueIndex(sourceIndex, columnIndex); + case VECTOR, UNKNOWN -> throw unsupportedSnapshotDataType(type); + } + } + + private static IllegalArgumentException unsupportedSnapshotDataType(final TSDataType type) { + return new IllegalArgumentException( + String.format( + ImportWALMessages.EXCEPTION_UNSUPPORTED_SNAPSHOT_DATA_TYPE_ARG_7A32D312, type)); + } + + private static class TableTabletSchema { + private final List schemas; + private final List categories; + private final int tagCount; + + private TableTabletSchema( + final List schemas, + final List categories, + final int tagCount) { + this.schemas = schemas; + this.categories = categories; + this.tagCount = tagCount; + } + } + } + + static class ReplayStatistics { + private long replayedOperationCount; + private long skippedEntryCount; + private long totalBytes; + private long processedBytes; + private long completedFileCount; + private long elapsedNanos; + + long getReplayedOperationCount() { + return replayedOperationCount; + } + + long getSkippedEntryCount() { + return skippedEntryCount; + } + + long getTotalBytes() { + return totalBytes; + } + + long getCompletedFileCount() { + return completedFileCount; + } + + double getElapsedSeconds() { + return elapsedNanos / 1_000_000_000.0; + } + + double getProgressPercent() { + return totalBytes == 0 ? 100.0 : processedBytes * 100.0 / totalBytes; + } + + double getAverageRateMbPerSecond() { + final double elapsedSeconds = getElapsedSeconds(); + return elapsedSeconds <= 0 ? 0.0 : processedBytes / elapsedSeconds / (1024.0 * 1024.0); + } + } + + private static class WALReplayException extends IOException { + private WALReplayException(final String message, final Throwable cause) { + super(message, cause); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java new file mode 100644 index 0000000000000..787e06e197511 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java @@ -0,0 +1,506 @@ +/* + * 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.db.tools; + +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalDeleteDataNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; +import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable; +import org.apache.iotdb.db.storageengine.dataregion.wal.WALTestUtils; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryType; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALInfoEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALSignalEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.ILogWriter; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALFileTest; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALWriter; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALByteBufferForTest; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALFileStatus; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALFileUtils; +import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.session.Session; + +import org.apache.commons.cli.CommandLine; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.ArgumentCaptor; + +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.IOException; +import java.io.PrintStream; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +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 ImportWALTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + /** + * Covers recursive directory discovery with WAL versions 2 and 10 in separate node folders. The + * result must ignore non-WAL files and preserve parent-folder then numeric-version order. + */ + @Test + public void testCollectWALFilesRecursivelyAndSortByVersion() throws IOException { + final Path source = temporaryFolder.newFolder("wal-root").toPath(); + final Path nodeA = Files.createDirectory(source.resolve("node-a")); + final Path nodeB = Files.createDirectory(source.resolve("node-b")); + final Path a10 = createWALFile(nodeA, 10); + final Path a2 = createWALFile(nodeA, 2); + final Path b1 = createWALFile(nodeB, 1); + Files.createFile(nodeA.resolve("ignore.txt")); + + final List files = ImportWAL.collectWALFiles(source); + + assertEquals( + Arrays.asList( + a2.toAbsolutePath().normalize(), + a10.toAbsolutePath().normalize(), + b1.toAbsolutePath().normalize()), + files); + } + + /** + * Covers a real WAL file containing one tree insert and one internal signal. The insert must be + * sent once as a Tablet, while the signal is counted as skipped and no corruption is reported. + */ + @Test + public void testReplayWALFileReplaysInsertAndSkipsInternalEntry() throws Exception { + final File walFile = createWALFile(0); + final InsertRowNode rowNode = WALTestUtils.getInsertRowNode("root.sg.d1", 100); + writeWAL(walFile, new WALInfoEntry(1, rowNode), new WALSignalEntry(WALEntryType.CLOSE_SIGNAL)); + final Session treeSession = mock(Session.class); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALFiles( + Collections.singletonList(walFile.toPath()), + new ImportWAL.WALReplayer(treeSession, null, null)); + + assertEquals(1, statistics.getReplayedOperationCount()); + assertEquals(1, statistics.getSkippedEntryCount()); + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(treeSession).insertTablet(tabletCaptor.capture()); + assertEquals("root.sg.d1", tabletCaptor.getValue().getDeviceId()); + assertEquals(100, tabletCaptor.getValue().getTimestamp(0)); + } + + @Test + public void testReplayReportsProgressAndFileStatistics() throws Exception { + final File walFile = createWALFile(0); + writeWAL( + walFile, + new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.progress", 100)), + new WALSignalEntry(WALEntryType.CLOSE_SIGNAL)); + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALFiles( + Collections.singletonList(walFile.toPath()), + new ImportWAL.WALReplayer(mock(Session.class), null, null), + new PrintStream(output)); + + assertEquals(1, statistics.getCompletedFileCount()); + assertEquals(Files.size(walFile.toPath()), statistics.getTotalBytes()); + assertTrue(statistics.getElapsedSeconds() >= 0); + assertTrue(output.toString().contains("1/1")); + } + + /** + * Covers an active WAL containing one complete entry but no end marker or metadata. Replay must + * treat EOF at the entry boundary as clean and import the entry without waiting for writer close. + */ + @Test + public void testReplayActiveWALFileAtEntryBoundary() throws Exception { + final File walFile = createWALFile(0); + final InsertRowNode rowNode = WALTestUtils.getInsertRowNode("root.sg.active", 102); + final Session treeSession = mock(Session.class); + + try (WALWriter writer = new WALWriter(walFile)) { + writer.write(serializeWAL(new WALInfoEntry(1, rowNode))); + writer.force(); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALFiles( + Collections.singletonList(walFile.toPath()), + new ImportWAL.WALReplayer(treeSession, null, null)); + + assertEquals(1, statistics.getReplayedOperationCount()); + assertEquals(0, statistics.getSkippedEntryCount()); + verify(treeSession).insertTablet(any(Tablet.class)); + } + } + + /** + * Covers an aligned tree row. The replay must use the aligned Session API and must not fall back + * to the non-aligned Tablet API. + */ + @Test + public void testReplayAlignedTreeInsertUsesAlignedSessionAPI() throws Exception { + final InsertRowNode rowNode = WALTestUtils.getInsertRowNode("root.sg.aligned", 101); + rowNode.setAligned(true); + final Session treeSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, rowNode)); + + verify(treeSession).insertAlignedTablet(any(Tablet.class)); + verify(treeSession, never()).insertTablet(any(Tablet.class)); + } + + /** + * Covers a table-model WAL tablet with an explicit target database. Replay must use the table + * Session and retain the source table name in the converted Tablet. + */ + @Test + public void testReplayTableInsertUsesRelationalSessionAPI() throws Exception { + final RelationalInsertTabletNode node = WALFileTest.getRelationalInsertTabletNode("table1"); + final Session treeSession = mock(Session.class); + final Session tableSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, tableSession, "db").replay(new WALInfoEntry(1, node)); + + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(tableSession).insertRelationalTablet(tabletCaptor.capture()); + assertEquals("table1", tabletCaptor.getValue().getTableName()); + verify(treeSession, never()).insertTablet(any(Tablet.class)); + } + + /** + * Covers a table-model WAL insert without a target database. Replay must fail before issuing any + * write because WAL insert entries do not carry their source database name. + */ + @Test + public void testReplayTableInsertRequiresDatabase() throws Exception { + final RelationalInsertTabletNode node = WALFileTest.getRelationalInsertTabletNode("table1"); + final Session treeSession = mock(Session.class); + + assertThrows( + StatementExecutionException.class, + () -> new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, node))); + + verify(treeSession, never()).insertTablet(any(Tablet.class)); + } + + /** + * Covers a tree deletion with multiple paths and a bounded time range. Replay must pass the + * original paths and inclusive time bounds to Session.deleteData. + */ + @Test + public void testReplayTreeDeletePreservesPathsAndTimeRange() throws Exception { + final DeleteDataNode deleteNode = + new DeleteDataNode( + new PlanNodeId(""), + Arrays.asList( + new MeasurementPath("root.sg.d1.s1"), new MeasurementPath("root.sg.d2.*")), + 10, + 20); + final Session treeSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, deleteNode)); + + verify(treeSession) + .deleteData(eq(Arrays.asList("root.sg.d1.s1", "root.sg.d2.*")), eq(10L), eq(20L)); + } + + /** Covers an unsupported entry when the interactive user explicitly chooses to skip it. */ + @Test + public void testReplayUnsupportedEntrySkipsAfterConfirmation() throws Exception { + final WALEntry entry = mockUnsupportedEntry(); + + final boolean replayed = + new ImportWAL.WALReplayer(mock(Session.class), null, null, ignored -> true).replay(entry); + + assertFalse(replayed); + } + + /** Covers an unsupported entry when the interactive user declines the skip prompt. */ + @Test + public void testReplayUnsupportedEntryFailsAfterDecliningSkip() { + final WALEntry entry = mockUnsupportedEntry(); + + assertThrows( + StatementExecutionException.class, + () -> + new ImportWAL.WALReplayer(mock(Session.class), null, null, ignored -> false) + .replay(entry)); + } + + /** Covers non-interactive execution, which must retain the original fail-fast behavior. */ + @Test + public void testReplayUnsupportedEntryFailsWithoutInteractiveInput() { + final WALEntry entry = mockUnsupportedEntry(); + + assertThrows( + StatementExecutionException.class, + () -> new ImportWAL.WALReplayer(mock(Session.class), null, null, null).replay(entry)); + } + + /** Covers accepted confirmations and the safe default for all other prompt answers. */ + @Test + public void testUnsupportedEntrySkipConfirmationParsing() { + assertTrue(ImportWAL.WALReplayer.isSkipConfirmation("y")); + assertTrue(ImportWAL.WALReplayer.isSkipConfirmation(" YES ")); + assertFalse(ImportWAL.WALReplayer.isSkipConfirmation("n")); + assertFalse(ImportWAL.WALReplayer.isSkipConfirmation("")); + assertFalse(ImportWAL.WALReplayer.isSkipConfirmation(null)); + } + + /** Covers a non-aligned snapshot whose measurements have independent time axes. */ + @Test + public void testReplayNonAlignedMemTableSnapshotAsTablets() throws Exception { + final PrimitiveMemTable memTable = new PrimitiveMemTable("root.sg", "0"); + final List schemas = + Arrays.asList( + new MeasurementSchema("s1", TSDataType.INT32), + new MeasurementSchema("s2", TSDataType.INT64)); + final StringArrayDeviceID deviceId = new StringArrayDeviceID("root.sg.d1"); + memTable.write(deviceId, schemas, 3, new Object[] {30, 300L}); + memTable.write(deviceId, schemas, 1, new Object[] {10, null}); + final Session treeSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, memTable)); + + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(treeSession, times(2)).insertTablet(tabletCaptor.capture()); + final Tablet s1Tablet = + tabletCaptor.getAllValues().stream() + .filter(tablet -> "s1".equals(tablet.getSchemas().get(0).getMeasurementName())) + .findFirst() + .orElseThrow(AssertionError::new); + assertEquals(2, s1Tablet.getRowSize()); + assertEquals(1, s1Tablet.getTimestamp(0)); + assertEquals(3, s1Tablet.getTimestamp(1)); + assertArrayEquals(new int[] {10, 30}, (int[]) s1Tablet.getValues()[0]); + } + + /** Covers an aligned snapshot with nulls and verifies the aligned Session API is used. */ + @Test + public void testReplayAlignedMemTableSnapshotPreservesNulls() throws Exception { + final PrimitiveMemTable memTable = new PrimitiveMemTable("root.sg", "0"); + final List schemas = + Arrays.asList( + new MeasurementSchema("s1", TSDataType.INT32), + new MeasurementSchema("s2", TSDataType.INT64)); + final StringArrayDeviceID deviceId = new StringArrayDeviceID("root.sg.d1"); + memTable.writeAlignedRow(deviceId, schemas, 2, new Object[] {20, null}); + memTable.writeAlignedRow(deviceId, schemas, 1, new Object[] {10, 100L}); + final Session treeSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, memTable)); + + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(treeSession).insertAlignedTablet(tabletCaptor.capture()); + verify(treeSession, never()).insertTablet(any(Tablet.class)); + final Tablet tablet = tabletCaptor.getValue(); + assertEquals(2, tablet.getRowSize()); + assertEquals(1, tablet.getTimestamp(0)); + assertEquals(2, tablet.getTimestamp(1)); + assertArrayEquals(new int[] {10, 20}, (int[]) tablet.getValues()[0]); + assertTrue(tablet.getBitMaps()[1].isMarked(1)); + } + + /** Covers snapshot serialization and deserialization through a real WAL file. */ + @Test + public void testReplaySerializedMemTableSnapshot() throws Exception { + final PrimitiveMemTable memTable = new PrimitiveMemTable("root.sg", "0"); + memTable.write( + new StringArrayDeviceID("root.sg.d1"), + Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT32)), + 7, + new Object[] {70}); + final File walFile = createWALFile(0); + writeWAL(walFile, new WALInfoEntry(1, memTable)); + final Session treeSession = mock(Session.class); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALFiles( + Collections.singletonList(walFile.toPath()), + new ImportWAL.WALReplayer(treeSession, null, null)); + + assertEquals(1, statistics.getReplayedOperationCount()); + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(treeSession).insertTablet(tabletCaptor.capture()); + assertEquals(7, tabletCaptor.getValue().getTimestamp(0)); + assertEquals(70, ((int[]) tabletCaptor.getValue().getValues()[0])[0]); + } + + /** Covers a snapshot larger than the replay batch limit. */ + @Test + public void testReplayMemTableSnapshotSplitsLargeChunk() throws Exception { + final PrimitiveMemTable memTable = new PrimitiveMemTable("root.sg", "0"); + final List schemas = + Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT32)); + final StringArrayDeviceID deviceId = new StringArrayDeviceID("root.sg.d1"); + for (int i = 0; i < 1025; i++) { + memTable.write(deviceId, schemas, i, new Object[] {i}); + } + final Session treeSession = mock(Session.class); + + new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, memTable)); + + final ArgumentCaptor tabletCaptor = ArgumentCaptor.forClass(Tablet.class); + verify(treeSession, times(2)).insertTablet(tabletCaptor.capture()); + assertEquals(1024, tabletCaptor.getAllValues().get(0).getRowSize()); + assertEquals(1, tabletCaptor.getAllValues().get(1).getRowSize()); + } + + /** Covers a signal snapshot, which carries no user data and must be skipped. */ + @Test + public void testReplaySignalMemTableSnapshotIsSkipped() throws Exception { + final IMemTable signalMemTable = mock(IMemTable.class); + when(signalMemTable.isSignalMemTable()).thenReturn(true); + final Session treeSession = mock(Session.class); + + final boolean replayed = + new ImportWAL.WALReplayer(treeSession, null, null) + .replay(new WALInfoEntry(1, signalMemTable)); + + assertFalse(replayed); + verify(treeSession, never()).insertTablet(any(Tablet.class)); + verify(treeSession, never()).insertAlignedTablet(any(Tablet.class)); + } + + /** Covers a table-model snapshot without a target database. */ + @Test + public void testReplayTableMemTableSnapshotRequiresDatabase() throws Exception { + final PrimitiveMemTable memTable = new PrimitiveMemTable("db", "0"); + memTable.writeAlignedRow( + new StringArrayDeviceID("table1", "device1"), + Collections.singletonList(new MeasurementSchema("temperature", TSDataType.FLOAT)), + 1, + new Object[] {1.0F}); + final Session treeSession = mock(Session.class); + + assertThrows( + StatementExecutionException.class, + () -> + new ImportWAL.WALReplayer(treeSession, null, null) + .replay(new WALInfoEntry(1, memTable))); + + verify(treeSession, never()).insertAlignedTablet(any(Tablet.class)); + } + + /** Covers table-model identifier quoting, including an embedded double quote. */ + @Test + public void testQuoteTableIdentifierForDescribe() { + assertEquals("\"table\"", ImportWAL.WALReplayer.quoteIdentifier("table")); + assertEquals("\"table\"\"name\"", ImportWAL.WALReplayer.quoteIdentifier("table\"name")); + } + + /** + * Covers a truncated WAL that cannot yield a complete entry. The file-level replay must fail so + * callers cannot mistake a partial replay for success. + */ + @Test + public void testReplayWALFileFailsOnCorruption() throws Exception { + final File walFile = createWALFile(0); + Files.write(walFile.toPath(), new byte[] {WALEntryType.INSERT_ROW_NODE.getCode()}); + + final IOException exception = + assertThrows( + IOException.class, + () -> + ImportWAL.replayWALFiles( + Collections.singletonList(walFile.toPath()), + new ImportWAL.WALReplayer(mock(Session.class), null, null))); + + assertTrue(exception.getMessage().contains(walFile.getName())); + } + + @Test + public void testPasswordIsRequiredWhenInteractiveInputIsUnavailable() { + final CommandLine commandLine = mock(CommandLine.class); + when(commandLine.hasOption("password")).thenReturn(false); + + assertThrows(IllegalArgumentException.class, () -> ImportWAL.getPassword(commandLine, null)); + } + + @Test + public void testExplicitPasswordTakesPrecedence() { + final CommandLine commandLine = mock(CommandLine.class); + when(commandLine.hasOption("password")).thenReturn(true); + when(commandLine.getOptionValue("password")).thenReturn("secret"); + + assertEquals("secret", ImportWAL.getPassword(commandLine, null)); + } + + private Path createWALFile(final Path parent, final long version) throws IOException { + return Files.createFile( + parent.resolve( + WALFileUtils.getLogFileName(version, 0, WALFileStatus.CONTAINS_SEARCH_INDEX))); + } + + private File createWALFile(final long version) throws IOException { + return temporaryFolder.newFile( + WALFileUtils.getLogFileName(version, 0, WALFileStatus.CONTAINS_SEARCH_INDEX)); + } + + private static void writeWAL(final File walFile, final WALEntry... entries) throws IOException { + try (ILogWriter writer = new WALWriter(walFile)) { + writer.write(serializeWAL(entries)); + } + } + + private static ByteBuffer serializeWAL(final WALEntry... entries) { + int serializedSize = 0; + for (final WALEntry entry : entries) { + serializedSize += entry.serializedSize(); + } + final WALByteBufferForTest buffer = + new WALByteBufferForTest(ByteBuffer.allocate(serializedSize)); + for (final WALEntry entry : entries) { + entry.serialize(buffer); + } + return buffer.getBuffer(); + } + + private static WALEntry mockUnsupportedEntry() { + final WALEntry entry = mock(WALEntry.class); + when(entry.getType()).thenReturn(WALEntryType.RELATIONAL_DELETE_DATA_NODE); + when(entry.getValue()).thenReturn(mock(RelationalDeleteDataNode.class)); + return entry; + } +} diff --git a/scripts/tools/import-wal.sh b/scripts/tools/import-wal.sh new file mode 100644 index 0000000000000..58bfe63c9eefd --- /dev/null +++ b/scripts/tools/import-wal.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# +# 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. +# + +if [ -z "${IOTDB_INCLUDE}" ]; then + : +elif [ -r "$IOTDB_INCLUDE" ]; then + . "$IOTDB_INCLUDE" +fi + +if [ -z "${IOTDB_HOME}" ]; then + export IOTDB_HOME="$(cd "$(dirname "$0")"/..; pwd)" +fi + +if [ -n "$JAVA_HOME" ]; then + for java in "$JAVA_HOME"/bin/amd64/java "$JAVA_HOME"/bin/java; do + if [ -x "$java" ]; then + JAVA="$java" + break + fi + done +else + JAVA=java +fi + +if [ -z "$JAVA" ]; then + echo "Unable to find java executable. Check JAVA_HOME and PATH environment variables." > /dev/stderr + exit 1 +fi + +JVM_OPTS="-Dsun.jnu.encoding=UTF-8 -Dfile.encoding=UTF-8" +CLASSPATH="${IOTDB_HOME}/lib/*" +MAIN_CLASS=org.apache.iotdb.db.tools.ImportWAL + +"$JAVA" $JVM_OPTS -DIOTDB_HOME="${IOTDB_HOME}" -cp "$CLASSPATH" "$MAIN_CLASS" "$@" +exit $? diff --git a/scripts/tools/windows/import-wal.bat b/scripts/tools/windows/import-wal.bat new file mode 100644 index 0000000000000..dc69c440d816e --- /dev/null +++ b/scripts/tools/windows/import-wal.bat @@ -0,0 +1,43 @@ +@REM +@REM Licensed to the Apache Software Foundation (ASF) under one +@REM or more contributor license agreements. See the NOTICE file +@REM distributed with this work for additional information +@REM regarding copyright ownership. The ASF licenses this file +@REM to you under the Apache License, Version 2.0 (the +@REM "License"); you may not use this file except in compliance +@REM with the License. You may obtain a copy of the License at +@REM +@REM http://www.apache.org/licenses/LICENSE-2.0 +@REM +@REM Unless required by applicable law or agreed to in writing, +@REM software distributed under the License is distributed on an +@REM "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +@REM KIND, either express or implied. See the License for the +@REM specific language governing permissions and limitations +@REM under the License. +@REM + +@echo off +if "%OS%" == "Windows_NT" setlocal + +pushd %~dp0..\.. +if NOT DEFINED IOTDB_HOME set IOTDB_HOME=%CD% +popd + +if NOT DEFINED MAIN_CLASS set MAIN_CLASS=org.apache.iotdb.db.tools.ImportWAL +if NOT DEFINED JAVA_HOME goto :err + +set JAVA_OPTS=-ea^ + -DIOTDB_HOME="%IOTDB_HOME%" -Dsun.jnu.encoding=UTF-8 -Dfile.encoding=UTF-8 +set CLASSPATH="%IOTDB_HOME%\lib\*" + +"%JAVA_HOME%\bin\java" %JAVA_OPTS% -cp %CLASSPATH% %MAIN_CLASS% %* +set ret_code=%ERRORLEVEL% +goto finally + +:err +echo JAVA_HOME environment variable must be set! +set ret_code=1 + +:finally +ENDLOCAL & EXIT /B %ret_code% From a25258ac70aff80e3af284f99208c29db1194390 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Thu, 13 Aug 2026 17:05:40 +0800 Subject: [PATCH 2/4] add on_success --- .../iotdb/db/i18n/ImportWALMessages.java | 10 +++ .../iotdb/db/i18n/ImportWALMessages.java | 10 +++ .../org/apache/iotdb/db/tools/ImportWAL.java | 61 +++++++++++++++++- .../apache/iotdb/db/tools/ImportWALTest.java | 63 +++++++++++++++++++ 4 files changed, 143 insertions(+), 1 deletion(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java index 2ed53656ec0f0..0dd487bfe8b80 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -38,6 +38,9 @@ public final class ImportWALMessages { public static final String MESSAGE_PASSWORD_PROMPT_F2D0E794 = "Password: "; public static final String MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C = "Target database for table-model WAL entries."; + public static final String + MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = + "When all WAL files are replayed successfully, do operation on source WAL files. Optional parameters are none (default) and delete."; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "Print this help message."; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "Argument error: %s"; @@ -58,8 +61,15 @@ public final class ImportWALMessages { public static final String MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D = "Import duration: %.1f s; total size: %d bytes; average rate: %.1f MB/s."; + public static final String MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B = + "Deleted %d source WAL files."; public static final String EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9 = "Failed to replay WAL file %s at offset %d: %s"; + public static final String EXCEPTION_FAILED_TO_DELETE_SOURCE_WAL_FILE_ARG_ARG_236AF580 = + "Failed to delete source WAL file %s: %s"; + public static final String + EXCEPTION_UNSUPPORTED_ON_SUCCESS_VALUE_ARG_EXPECTED_NONE_OR_DELETE_F1C8EACE = + "Unsupported on_success value: %s. Expected none or delete."; public static final String EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726 = "Table-model WAL entries require -db/--database."; public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java index c8dc2266dc8a1..57052048e6332 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -38,6 +38,9 @@ public final class ImportWALMessages { public static final String MESSAGE_PASSWORD_PROMPT_F2D0E794 = "密码："; public static final String MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C = "表模型 WAL 条目的目标数据库。"; + public static final String + MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = + "所有 WAL 文件成功重放后，对源 WAL 文件执行操作。可选参数为 none（默认）和 delete。"; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "打印帮助信息。"; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "参数错误：%s"; public static final String MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA = "WAL 导入失败：%s"; @@ -57,8 +60,15 @@ public final class ImportWALMessages { public static final String MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D = "导入耗时：%.1f 秒；文件总大小：%d 字节；平均速率：%.1f MB/s。"; + public static final String MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B = + "已删除 %d 个源 WAL 文件。"; public static final String EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9 = "重放 WAL 文件 %s 时失败，偏移量 %d：%s"; + public static final String EXCEPTION_FAILED_TO_DELETE_SOURCE_WAL_FILE_ARG_ARG_236AF580 = + "删除源 WAL 文件 %s 失败：%s"; + public static final String + EXCEPTION_UNSUPPORTED_ON_SUCCESS_VALUE_ARG_EXPECTED_NONE_OR_DELETE_F1C8EACE = + "不支持的 on_success 值：%s。应为 none 或 delete。"; public static final String EXCEPTION_TABLE_MODEL_WAL_ENTRIES_REQUIRE_DB_DATABASE_F7597726 = "表模型 WAL 条目要求指定 -db/--database。"; public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java index f8316499bdb5f..c1281c00ee301 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java @@ -118,6 +118,8 @@ static int run(final String[] args, final PrintStream out, final PrintStream err final Path source = Paths.get(commandLine.getOptionValue("file")); final List walFiles = collectWALFiles(source); final String database = commandLine.getOptionValue("database"); + final boolean deleteSource = + shouldDeleteSource(commandLine.getOptionValue("on_success", "none")); final String password = getPassword(commandLine); final Session treeSession = createSession( @@ -141,7 +143,8 @@ static int run(final String[] args, final PrintStream out, final PrintStream err tableSession.open(false); } final ReplayStatistics statistics = - replayWALFiles(walFiles, new WALReplayer(treeSession, tableSession, database), out); + replayWALFiles( + walFiles, new WALReplayer(treeSession, tableSession, database), out, deleteSource); out.printf( ImportWALMessages .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, @@ -156,6 +159,11 @@ static int run(final String[] args, final PrintStream out, final PrintStream err statistics.getTotalBytes(), statistics.getAverageRateMbPerSecond()); out.println(); + if (deleteSource) { + out.printf( + ImportWALMessages.MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B, walFiles.size()); + out.println(); + } } finally { closeSession(tableSession); closeSession(treeSession); @@ -211,6 +219,15 @@ private static Options createOptions() { .hasArg() .desc(ImportWALMessages.MESSAGE_TARGET_DATABASE_FOR_TABLE_MODEL_WAL_ENTRIES_27BACD1C) .build()); + options.addOption( + Option.builder("os") + .longOpt("on_success") + .argName("on_success") + .hasArg() + .desc( + ImportWALMessages + .MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66) + .build()); options.addOption( Option.builder() .longOpt("help") @@ -242,6 +259,21 @@ static String getPassword(final CommandLine commandLine, final Console console) return new String(password); } + static boolean shouldDeleteSource(final String onSuccess) { + final String normalizedOnSuccess = onSuccess.trim(); + if ("none".equalsIgnoreCase(normalizedOnSuccess)) { + return false; + } + if ("delete".equalsIgnoreCase(normalizedOnSuccess)) { + return true; + } + throw new IllegalArgumentException( + String.format( + ImportWALMessages + .EXCEPTION_UNSUPPORTED_ON_SUCCESS_VALUE_ARG_EXPECTED_NONE_OR_DELETE_F1C8EACE, + onSuccess)); + } + private static boolean containsHelpOption(final String[] args) { if (args == null) { return false; @@ -368,6 +400,15 @@ static ReplayStatistics replayWALFiles(final List walFiles, final WALRepla static ReplayStatistics replayWALFiles( final List walFiles, final WALReplayer replayer, final PrintStream progressStream) throws IOException { + return replayWALFiles(walFiles, replayer, progressStream, false); + } + + static ReplayStatistics replayWALFiles( + final List walFiles, + final WALReplayer replayer, + final PrintStream progressStream, + final boolean deleteSource) + throws IOException { final ReplayStatistics statistics = new ReplayStatistics(); final long startNanos = System.nanoTime(); for (final Path walFile : walFiles) { @@ -436,9 +477,27 @@ static ReplayStatistics replayWALFiles( e); } } + if (deleteSource) { + deleteSourceWALFiles(walFiles); + } return statistics; } + private static void deleteSourceWALFiles(final List walFiles) throws IOException { + for (final Path walFile : walFiles) { + try { + Files.delete(walFile); + } catch (final IOException e) { + throw new IOException( + String.format( + ImportWALMessages.EXCEPTION_FAILED_TO_DELETE_SOURCE_WAL_FILE_ARG_ARG_236AF580, + walFile, + e.getMessage()), + e); + } + } + } + static class WALReplayer { private final Session treeSession; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java index 787e06e197511..c8efb24ef9fcd 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java @@ -104,6 +104,26 @@ public void testCollectWALFilesRecursivelyAndSortByVersion() throws IOException files); } + /** Covers CLI discovery of the source-file operation without opening a Session. */ + @Test + public void testHelpDescribesDeleteSourceOption() { + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + + final int exitCode = + ImportWAL.run(new String[] {"--help"}, new PrintStream(output), new PrintStream(output)); + + assertEquals(0, exitCode); + assertTrue(output.toString().contains("--on_success")); + } + + /** Covers the default retention value, deletion value, normalization, and invalid input. */ + @Test + public void testParseOnSuccessOption() { + assertFalse(ImportWAL.shouldDeleteSource("none")); + assertTrue(ImportWAL.shouldDeleteSource(" DELETE ")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.shouldDeleteSource("unsupported")); + } + /** * Covers a real WAL file containing one tree insert and one internal signal. The insert must be * sent once as a Tablet, while the signal is counted as skipped and no corruption is reported. @@ -126,6 +146,49 @@ public void testReplayWALFileReplaysInsertAndSkipsInternalEntry() throws Excepti verify(treeSession).insertTablet(tabletCaptor.capture()); assertEquals("root.sg.d1", tabletCaptor.getValue().getDeviceId()); assertEquals(100, tabletCaptor.getValue().getTimestamp(0)); + assertTrue(walFile.exists()); + } + + /** Covers opt-in deletion after every source WAL file completes successfully. */ + @Test + public void testReplayDeletesSourceFilesAfterAllFilesSucceed() throws Exception { + final File firstWALFile = createWALFile(0); + final File secondWALFile = createWALFile(1); + writeWAL( + firstWALFile, new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.delete.d1", 1))); + writeWAL( + secondWALFile, new WALInfoEntry(2, WALTestUtils.getInsertRowNode("root.sg.delete.d2", 2))); + + ImportWAL.replayWALFiles( + Arrays.asList(firstWALFile.toPath(), secondWALFile.toPath()), + new ImportWAL.WALReplayer(mock(Session.class), null, null), + null, + true); + + assertFalse(firstWALFile.exists()); + assertFalse(secondWALFile.exists()); + } + + /** Covers all-or-nothing replay gating: a later failure must retain every source WAL file. */ + @Test + public void testReplayRetainsAllSourceFilesWhenAnyFileFails() throws Exception { + final File validWALFile = createWALFile(0); + final File corruptedWALFile = createWALFile(1); + writeWAL( + validWALFile, new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.retain.d1", 1))); + Files.write(corruptedWALFile.toPath(), new byte[] {WALEntryType.INSERT_ROW_NODE.getCode()}); + + assertThrows( + IOException.class, + () -> + ImportWAL.replayWALFiles( + Arrays.asList(validWALFile.toPath(), corruptedWALFile.toPath()), + new ImportWAL.WALReplayer(mock(Session.class), null, null), + null, + true)); + + assertTrue(validWALFile.exists()); + assertTrue(corruptedWALFile.exists()); } @Test From 2cd41b697d462a1cb9b48e2dfe7150f0d7a14904 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Mon, 24 Aug 2026 11:18:24 +0800 Subject: [PATCH 3/4] add thread_num --- .../iotdb/db/i18n/ImportWALMessages.java | 8 + .../iotdb/db/i18n/ImportWALMessages.java | 8 + .../org/apache/iotdb/db/tools/ImportWAL.java | 437 +++++++++++++----- .../apache/iotdb/db/tools/ImportWALTest.java | 112 ++++- 4 files changed, 459 insertions(+), 106 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java index 0dd487bfe8b80..3061a1845be8f 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -41,6 +41,9 @@ public final class ImportWALMessages { public static final String MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = "When all WAL files are replayed successfully, do operation on source WAL files. Optional parameters are none (default) and delete."; + public static final String + MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50 = + "Number of threads used to replay WAL directories in parallel. Default: 1."; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "Print this help message."; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "Argument error: %s"; @@ -52,6 +55,11 @@ public final class ImportWALMessages { public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = "No WAL files found under: %s"; public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "Invalid port: %s"; + public static final String + EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD = + "Invalid thread count: %s. Expected a positive integer."; + public static final String EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD = + "WAL replay was interrupted."; public static final String MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = "Replayed %d operations from %d WAL files; skipped %d entries."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java index 57052048e6332..120c32659996b 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -41,6 +41,9 @@ public final class ImportWALMessages { public static final String MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = "所有 WAL 文件成功重放后，对源 WAL 文件执行操作。可选参数为 none（默认）和 delete。"; + public static final String + MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50 = + "并行重放 WAL 目录所用的线程数。默认：1。"; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "打印帮助信息。"; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "参数错误：%s"; public static final String MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA = "WAL 导入失败：%s"; @@ -51,6 +54,11 @@ public final class ImportWALMessages { public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = "路径下未找到 WAL 文件：%s"; public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "无效端口：%s"; + public static final String + EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD = + "无效线程数：%s。应为正整数。"; + public static final String EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD = + "WAL 重放被中断。"; public static final String MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = "已重放 %d 个操作（来自 %d 个 WAL 文件）；跳过 %d 个条目。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java index c1281c00ee301..2214b69db9210 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java @@ -75,10 +75,17 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -89,7 +96,9 @@ public class ImportWAL { private static final String DEFAULT_HOST = "127.0.0.1"; private static final int DEFAULT_PORT = 6667; private static final String DEFAULT_USER = "root"; + private static final int DEFAULT_THREAD_NUM = 1; private static final int SNAPSHOT_TABLET_ROW_LIMIT = 1024; + private static final Object CONSOLE_PROMPT_LOCK = new Object(); private ImportWAL() {} @@ -120,53 +129,38 @@ static int run(final String[] args, final PrintStream out, final PrintStream err final String database = commandLine.getOptionValue("database"); final boolean deleteSource = shouldDeleteSource(commandLine.getOptionValue("on_success", "none")); + final int threadNum = + parseThreadNum( + commandLine.getOptionValue("thread_num", String.valueOf(DEFAULT_THREAD_NUM))); final String password = getPassword(commandLine); - final Session treeSession = - createSession( - commandLine.getOptionValue("host", DEFAULT_HOST), - parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), - commandLine.getOptionValue("username", DEFAULT_USER), - password, - null); - final Session tableSession = - database == null - ? null - : createSession( - commandLine.getOptionValue("host", DEFAULT_HOST), - parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), - commandLine.getOptionValue("username", DEFAULT_USER), - password, - database); - try { - treeSession.open(false); - if (tableSession != null) { - tableSession.open(false); - } - final ReplayStatistics statistics = - replayWALFiles( - walFiles, new WALReplayer(treeSession, tableSession, database), out, deleteSource); + final String host = commandLine.getOptionValue("host", DEFAULT_HOST); + final int port = parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))); + final String username = commandLine.getOptionValue("username", DEFAULT_USER); + final ReplayStatistics statistics = + replayWALDirectories( + walFiles, + threadNum, + () -> createWALReplayWorker(host, port, username, password, database), + out, + deleteSource); + out.printf( + ImportWALMessages + .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, + statistics.replayedOperationCount, + walFiles.size(), + statistics.skippedEntryCount); + out.println(); + out.printf( + ImportWALMessages + .MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D, + statistics.getElapsedSeconds(), + statistics.getTotalBytes(), + statistics.getAverageRateMbPerSecond()); + out.println(); + if (deleteSource) { out.printf( - ImportWALMessages - .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, - statistics.replayedOperationCount, - walFiles.size(), - statistics.skippedEntryCount); + ImportWALMessages.MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B, walFiles.size()); out.println(); - out.printf( - ImportWALMessages - .MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D, - statistics.getElapsedSeconds(), - statistics.getTotalBytes(), - statistics.getAverageRateMbPerSecond()); - out.println(); - if (deleteSource) { - out.printf( - ImportWALMessages.MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B, walFiles.size()); - out.println(); - } - } finally { - closeSession(tableSession); - closeSession(treeSession); } return CODE_OK; } catch (final Exception e) { @@ -228,6 +222,15 @@ private static Options createOptions() { ImportWALMessages .MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66) .build()); + options.addOption( + Option.builder("tn") + .longOpt("thread_num") + .argName("thread_num") + .hasArg() + .desc( + ImportWALMessages + .MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50) + .build()); options.addOption( Option.builder() .longOpt("help") @@ -314,6 +317,23 @@ private static int parsePort(final String port) { } } + static int parseThreadNum(final String threadNum) { + try { + final int value = Integer.parseInt(threadNum); + if (value <= 0) { + throw new NumberFormatException(threadNum); + } + return value; + } catch (final NumberFormatException e) { + throw new IllegalArgumentException( + String.format( + ImportWALMessages + .EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD, + threadNum), + e); + } + } + private static Session createSession( final String host, final int port, @@ -328,6 +348,29 @@ private static Session createSession( return builder.build(); } + private static WALReplayWorker createWALReplayWorker( + final String host, + final int port, + final String username, + final String password, + final String database) + throws IOException { + final Session treeSession = createSession(host, port, username, password, null); + final Session tableSession = + database == null ? null : createSession(host, port, username, password, database); + try { + treeSession.open(false); + if (tableSession != null) { + tableSession.open(false); + } + return new SessionWALReplayer(treeSession, tableSession, database); + } catch (final IoTDBConnectionException e) { + closeSession(tableSession); + closeSession(treeSession); + throw new IOException(e.getMessage(), e); + } + } + private static void closeSession(final Session session) { if (session == null) { return; @@ -409,78 +452,219 @@ static ReplayStatistics replayWALFiles( final PrintStream progressStream, final boolean deleteSource) throws IOException { - final ReplayStatistics statistics = new ReplayStatistics(); final long startNanos = System.nanoTime(); + final ReplayStatistics statistics = createReplayStatistics(walFiles); for (final Path walFile : walFiles) { - statistics.totalBytes += Files.size(walFile); + recordCompletedFile( + statistics, + replayWALFile(walFile, replayer), + walFiles.size(), + startNanos, + progressStream); + } + if (deleteSource) { + deleteSourceWALFiles(walFiles); + } + return statistics; + } + + static ReplayStatistics replayWALDirectories( + final List walFiles, + final int threadNum, + final WALReplayWorkerFactory workerFactory, + final PrintStream progressStream, + final boolean deleteSource) + throws IOException { + if (threadNum <= 0) { + throw new IllegalArgumentException( + String.format( + ImportWALMessages + .EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD, + threadNum)); + } + final long startNanos = System.nanoTime(); + final ReplayStatistics statistics = createReplayStatistics(walFiles); + final List> walDirectories = groupWALFilesByDirectory(walFiles); + if (walDirectories.isEmpty()) { + statistics.elapsedNanos = System.nanoTime() - startNanos; + return statistics; + } + + final int workerCount = Math.min(threadNum, walDirectories.size()); + final AtomicInteger nextDirectoryIndex = new AtomicInteger(); + final AtomicReference replayFailure = new AtomicReference<>(); + final ExecutorService executor = Executors.newFixedThreadPool(workerCount); + final List> futures = new ArrayList<>(workerCount); + for (int i = 0; i < workerCount; i++) { + futures.add( + executor.submit( + () -> { + try (WALReplayWorker worker = workerFactory.create()) { + int directoryIndex; + // A worker owns one directory at a time so WAL files from that directory remain + // ordered, while independent directories can make progress concurrently. + while (replayFailure.get() == null + && (directoryIndex = nextDirectoryIndex.getAndIncrement()) + < walDirectories.size()) { + for (final Path walFile : walDirectories.get(directoryIndex)) { + recordCompletedFile( + statistics, + replayWALFile(walFile, worker), + walFiles.size(), + startNanos, + progressStream); + } + } + } catch (final Exception e) { + replayFailure.compareAndSet(null, e); + } + })); + } + executor.shutdown(); + + boolean interrupted = false; + for (final Future future : futures) { + boolean completed = false; + // Wait for every worker even after a failure so sessions are closed before source deletion or + // the failure is reported to the caller. + while (!completed) { + try { + future.get(); + completed = true; + } catch (final InterruptedException e) { + interrupted = true; + replayFailure.compareAndSet( + null, + new IOException(ImportWALMessages.EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD, e)); + } catch (final ExecutionException e) { + replayFailure.compareAndSet(null, e.getCause()); + completed = true; + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + statistics.elapsedNanos = System.nanoTime() - startNanos; + rethrowReplayFailure(replayFailure.get()); + if (deleteSource) { + deleteSourceWALFiles(walFiles); + } + return statistics; + } + + static List> groupWALFilesByDirectory(final List walFiles) { + final List sortedWALFiles = new ArrayList<>(walFiles); + sortedWALFiles.sort(WAL_FILE_COMPARATOR); + final Map> filesByDirectory = new LinkedHashMap<>(); + for (final Path walFile : sortedWALFiles) { + filesByDirectory + .computeIfAbsent(walFile.getParent(), ignored -> new ArrayList<>()) + .add(walFile); } + return new ArrayList<>(filesByDirectory.values()); + } + + private static ReplayStatistics createReplayStatistics(final List walFiles) + throws IOException { + final ReplayStatistics statistics = new ReplayStatistics(); for (final Path walFile : walFiles) { - try (WALReader reader = new WALReader(walFile.toFile())) { - long offset = reader.getWALCurrentReadOffset(); - while (reader.hasNext()) { - final WALEntry entry = reader.next(); - try { - if (replayer.replay(entry)) { - statistics.replayedOperationCount++; - } else { - statistics.skippedEntryCount++; - } - } catch (final IoTDBConnectionException | StatementExecutionException e) { - throw new WALReplayException( - String.format( - ImportWALMessages - .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, - walFile, - offset, - e.getMessage()), - e); + statistics.totalBytes += Files.size(walFile); + } + return statistics; + } + + private static ReplayStatistics replayWALFile(final Path walFile, final WALReplayWorker replayer) + throws IOException { + final ReplayStatistics statistics = new ReplayStatistics(); + try (WALReader reader = new WALReader(walFile.toFile())) { + long offset = reader.getWALCurrentReadOffset(); + while (reader.hasNext()) { + final WALEntry entry = reader.next(); + try { + if (replayer.replay(entry)) { + statistics.replayedOperationCount++; + } else { + statistics.skippedEntryCount++; } - offset = reader.getWALCurrentReadOffset(); - } - if (reader.isFileCorrupted()) { + } catch (final IoTDBConnectionException | StatementExecutionException e) { throw new WALReplayException( String.format( ImportWALMessages .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, walFile, - reader.getWALCurrentReadOffset(), - ImportWALMessages.EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5), - null); - } - statistics.completedFileCount++; - statistics.processedBytes += Files.size(walFile); - statistics.elapsedNanos = System.nanoTime() - startNanos; - if (progressStream != null) { - progressStream.printf( - ImportWALMessages - .MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F, - statistics.completedFileCount, - walFiles.size(), - statistics.processedBytes, - statistics.totalBytes, - statistics.getProgressPercent(), - statistics.getElapsedSeconds(), - statistics.getAverageRateMbPerSecond()); - progressStream.println(); - } - } catch (final IOException e) { - if (e instanceof WALReplayException walReplayException) { - throw walReplayException; + offset, + e.getMessage()), + e); } + offset = reader.getWALCurrentReadOffset(); + } + if (reader.isFileCorrupted()) { throw new WALReplayException( String.format( ImportWALMessages .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, walFile, - 0, - e.getMessage()), - e); + reader.getWALCurrentReadOffset(), + ImportWALMessages.EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5), + null); } + statistics.completedFileCount++; + statistics.processedBytes += Files.size(walFile); + return statistics; + } catch (final IOException e) { + if (e instanceof WALReplayException walReplayException) { + throw walReplayException; + } + throw new WALReplayException( + String.format( + ImportWALMessages.EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, + walFile, + 0, + e.getMessage()), + e); } - if (deleteSource) { - deleteSourceWALFiles(walFiles); + } + + private static void recordCompletedFile( + final ReplayStatistics statistics, + final ReplayStatistics completedFileStatistics, + final int totalFileCount, + final long startNanos, + final PrintStream progressStream) { + synchronized (statistics) { + statistics.add(completedFileStatistics); + statistics.elapsedNanos = System.nanoTime() - startNanos; + if (progressStream != null) { + progressStream.printf( + ImportWALMessages + .MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F, + statistics.completedFileCount, + totalFileCount, + statistics.processedBytes, + statistics.totalBytes, + statistics.getProgressPercent(), + statistics.getElapsedSeconds(), + statistics.getAverageRateMbPerSecond()); + progressStream.println(); + } } - return statistics; + } + + private static void rethrowReplayFailure(final Throwable failure) throws IOException { + if (failure == null) { + return; + } + if (failure instanceof IOException ioException) { + throw ioException; + } + if (failure instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (failure instanceof Error error) { + throw error; + } + throw new IOException(failure); } private static void deleteSourceWALFiles(final List walFiles) throws IOException { @@ -498,7 +682,40 @@ private static void deleteSourceWALFiles(final List walFiles) throws IOExc } } - static class WALReplayer { + @FunctionalInterface + interface WALReplayWorkerFactory { + + WALReplayWorker create() throws Exception; + } + + interface WALReplayWorker extends AutoCloseable { + + boolean replay(WALEntry entry) throws IoTDBConnectionException, StatementExecutionException; + + @Override + default void close() {} + } + + private static class SessionWALReplayer extends WALReplayer { + + private final Session treeSession; + private final Session tableSession; + + private SessionWALReplayer( + final Session treeSession, final Session tableSession, final String tableDatabaseName) { + super(treeSession, tableSession, tableDatabaseName); + this.treeSession = treeSession; + this.tableSession = tableSession; + } + + @Override + public void close() { + closeSession(tableSession); + closeSession(treeSession); + } + } + + static class WALReplayer implements WALReplayWorker { private final Session treeSession; private final Session tableSession; @@ -528,7 +745,8 @@ static class WALReplayer { null, null, ColumnFilterMatcher.matchAll(), tableDatabaseName); } - boolean replay(final WALEntry entry) + @Override + public boolean replay(final WALEntry entry) throws IoTDBConnectionException, StatementExecutionException { if (entry.getType() == WALEntryType.MEMORY_TABLE_SNAPSHOT || entry.getType() == WALEntryType.OLD_MEMORY_TABLE_SNAPSHOT) { @@ -558,12 +776,14 @@ private static UnsupportedEntryPrompt createUnsupportedEntryPrompt(final Console return null; } return entry -> { - final String answer = - console.readLine( - ImportWALMessages - .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, - entry.getType()); - return isSkipConfirmation(answer); + synchronized (CONSOLE_PROMPT_LOCK) { + final String answer = + console.readLine( + ImportWALMessages + .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, + entry.getType()); + return isSkipConfirmation(answer); + } }; } @@ -998,6 +1218,13 @@ static class ReplayStatistics { private long completedFileCount; private long elapsedNanos; + private void add(final ReplayStatistics statistics) { + replayedOperationCount += statistics.replayedOperationCount; + skippedEntryCount += statistics.skippedEntryCount; + processedBytes += statistics.processedBytes; + completedFileCount += statistics.completedFileCount; + } + long getReplayedOperationCount() { return replayedOperationCount; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java index c8efb24ef9fcd..a40f020c18c5a 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java @@ -62,6 +62,12 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -70,6 +76,7 @@ import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -104,7 +111,7 @@ public void testCollectWALFilesRecursivelyAndSortByVersion() throws IOException files); } - /** Covers CLI discovery of the source-file operation without opening a Session. */ + /** Covers CLI discovery of source-file and parallel replay options without opening a Session. */ @Test public void testHelpDescribesDeleteSourceOption() { final ByteArrayOutputStream output = new ByteArrayOutputStream(); @@ -114,6 +121,79 @@ public void testHelpDescribesDeleteSourceOption() { assertEquals(0, exitCode); assertTrue(output.toString().contains("--on_success")); + assertTrue(output.toString().contains("--thread_num")); + } + + /** Covers valid thread counts and rejects zero, negative, and non-numeric values. */ + @Test + public void testParseThreadNumOption() { + assertEquals(1, ImportWAL.parseThreadNum("1")); + assertEquals(4, ImportWAL.parseThreadNum("4")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("0")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("-1")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("invalid")); + } + + /** + * Covers directory-level parallel replay with two WAL files per directory. Different directories + * must overlap, while versions within each directory must retain ascending replay order. + */ + @Test + public void testReplayWALDirectoriesInParallelAndPreserveDirectoryOrder() throws Exception { + final Path source = temporaryFolder.newFolder("parallel-wal-root").toPath(); + final Path nodeA = Files.createDirectory(source.resolve("node-a")); + final Path nodeB = Files.createDirectory(source.resolve("node-b")); + final Path a1 = createWALFile(nodeA, 1); + final Path a2 = createWALFile(nodeA, 2); + final Path b1 = createWALFile(nodeB, 1); + final Path b2 = createWALFile(nodeB, 2); + writeWAL(a1.toFile(), new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.a", 1))); + writeWAL(a2.toFile(), new WALInfoEntry(2, WALTestUtils.getInsertRowNode("root.sg.a", 2))); + writeWAL(b1.toFile(), new WALInfoEntry(3, WALTestUtils.getInsertRowNode("root.sg.b", 1))); + writeWAL(b2.toFile(), new WALInfoEntry(4, WALTestUtils.getInsertRowNode("root.sg.b", 2))); + final List walFiles = ImportWAL.collectWALFiles(source); + final CyclicBarrier replayBarrier = new CyclicBarrier(2); + final AtomicInteger activeReplays = new AtomicInteger(); + final AtomicInteger maxActiveReplays = new AtomicInteger(); + final AtomicInteger createdWorkers = new AtomicInteger(); + final Map> replayedTimestamps = new ConcurrentHashMap<>(); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALDirectories( + walFiles, + 2, + () -> { + createdWorkers.incrementAndGet(); + final Session session = mock(Session.class); + doAnswer( + invocation -> { + final Tablet tablet = invocation.getArgument(0); + final int active = activeReplays.incrementAndGet(); + maxActiveReplays.accumulateAndGet(active, Math::max); + try { + replayBarrier.await(5, TimeUnit.SECONDS); + replayedTimestamps + .computeIfAbsent( + tablet.getDeviceId(), ignored -> new CopyOnWriteArrayList<>()) + .add(tablet.getTimestamp(0)); + } finally { + activeReplays.decrementAndGet(); + } + return null; + }) + .when(session) + .insertTablet(any(Tablet.class)); + return new ImportWAL.WALReplayer(session, null, null); + }, + null, + false); + + assertEquals(2, createdWorkers.get()); + assertTrue(maxActiveReplays.get() >= 2); + assertEquals(Arrays.asList(1L, 2L), replayedTimestamps.get("root.sg.a")); + assertEquals(Arrays.asList(1L, 2L), replayedTimestamps.get("root.sg.b")); + assertEquals(4, statistics.getReplayedOperationCount()); + assertEquals(4, statistics.getCompletedFileCount()); } /** Covers the default retention value, deletion value, normalization, and invalid input. */ @@ -191,6 +271,36 @@ public void testReplayRetainsAllSourceFilesWhenAnyFileFails() throws Exception { assertTrue(corruptedWALFile.exists()); } + /** + * Covers all-or-nothing deletion during directory-level parallel replay. A corrupted directory + * must retain WAL files from both the failed directory and another concurrently replayed one. + */ + @Test + public void testParallelReplayRetainsAllSourceFilesWhenAnyDirectoryFails() throws Exception { + final Path source = temporaryFolder.newFolder("parallel-retain-wal-root").toPath(); + final Path validDirectory = Files.createDirectory(source.resolve("valid")); + final Path corruptedDirectory = Files.createDirectory(source.resolve("corrupted")); + final Path validWALFile = createWALFile(validDirectory, 0); + final Path corruptedWALFile = createWALFile(corruptedDirectory, 0); + writeWAL( + validWALFile.toFile(), + new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.parallel.retain", 1))); + Files.write(corruptedWALFile, new byte[] {WALEntryType.INSERT_ROW_NODE.getCode()}); + + assertThrows( + IOException.class, + () -> + ImportWAL.replayWALDirectories( + ImportWAL.collectWALFiles(source), + 2, + () -> new ImportWAL.WALReplayer(mock(Session.class), null, null), + null, + true)); + + assertTrue(Files.exists(validWALFile)); + assertTrue(Files.exists(corruptedWALFile)); + } + @Test public void testReplayReportsProgressAndFileStatistics() throws Exception { final File walFile = createWALFile(0); From fa7f2dd1a5699885c086c1fb71748d189693ec4a Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Mon, 24 Aug 2026 18:48:31 +0800 Subject: [PATCH 4/4] add detailed prompt --- .../iotdb/db/i18n/ImportWALMessages.java | 10 +- .../iotdb/db/i18n/ImportWALMessages.java | 10 +- .../org/apache/iotdb/db/tools/ImportWAL.java | 137 +++++++++++---- .../apache/iotdb/db/tools/ImportWALTest.java | 158 ++++++++++++++++-- 4 files changed, 268 insertions(+), 47 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java index 3061a1845be8f..1be122abbfe31 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -82,8 +82,14 @@ public final class ImportWALMessages { "Table-model WAL entries require -db/--database."; public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = "Unsupported WAL operation: %s"; - public static final String MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650 = - "Unsupported WAL operation: %s. Skip this entry? [y/N]: "; + public static final String + MESSAGE_TREE_MODEL_DELETE_OPERATION_DETECTED_ARG_CHOOSE_E_EXECUTE_S_SKIP_A_EXECUTE_ALL_L_SKIP_ALL_Q_QUIT_11E39FD7 = + "Tree-model delete operation detected: %s. Choose e=execute, s=skip, a=execute all, l=skip all, q=quit: "; + public static final String + MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_CHOOSE_S_SKIP_L_SKIP_ALL_Q_QUIT_0A734E52 = + "Unsupported WAL operation: %s. Choose s=skip, l=skip all, q=quit: "; + public static final String EXCEPTION_WAL_REPLAY_WAS_TERMINATED_BY_THE_USER_E0BD6197 = + "WAL replay was terminated by the user."; public static final String EXCEPTION_INSERT_NODE_ARG_CONTAINS_NO_REPLAYABLE_DATA_5DA13453 = "Insert node %s contains no replayable data."; public static final String EXCEPTION_UNSUPPORTED_SNAPSHOT_DATA_TYPE_ARG_7A32D312 = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java index 120c32659996b..912ab66ca268f 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -81,8 +81,14 @@ public final class ImportWALMessages { "表模型 WAL 条目要求指定 -db/--database。"; public static final String EXCEPTION_UNSUPPORTED_WAL_OPERATION_ARG_ABD227A0 = "不支持的 WAL 操作：%s"; - public static final String MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650 = - "不支持的 WAL 操作：%s。是否跳过此条目？[y/N]："; + public static final String + MESSAGE_TREE_MODEL_DELETE_OPERATION_DETECTED_ARG_CHOOSE_E_EXECUTE_S_SKIP_A_EXECUTE_ALL_L_SKIP_ALL_Q_QUIT_11E39FD7 = + "检测到树模型删除操作：%s。请选择 e=执行、s=跳过、a=全部执行、l=全部跳过、q=终止重放："; + public static final String + MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_CHOOSE_S_SKIP_L_SKIP_ALL_Q_QUIT_0A734E52 = + "不支持的 WAL 操作：%s。请选择 s=跳过、l=全部跳过、q=终止重放："; + public static final String EXCEPTION_WAL_REPLAY_WAS_TERMINATED_BY_THE_USER_E0BD6197 = + "用户终止了 WAL 重放。"; public static final String EXCEPTION_INSERT_NODE_ARG_CONTAINS_NO_REPLAYABLE_DATA_5DA13453 = "Insert node %s 不包含可重放数据。"; public static final String EXCEPTION_UNSUPPORTED_SNAPSHOT_DATA_TYPE_ARG_7A32D312 = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java index 2214b69db9210..dd1aca285b40e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java @@ -98,7 +98,6 @@ public class ImportWAL { private static final String DEFAULT_USER = "root"; private static final int DEFAULT_THREAD_NUM = 1; private static final int SNAPSHOT_TABLET_ROW_LIMIT = 1024; - private static final Object CONSOLE_PROMPT_LOCK = new Object(); private ImportWAL() {} @@ -136,11 +135,15 @@ static int run(final String[] args, final PrintStream out, final PrintStream err final String host = commandLine.getOptionValue("host", DEFAULT_HOST); final int port = parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))); final String username = commandLine.getOptionValue("username", DEFAULT_USER); + final WALReplayer.ReplayDecisionController replayDecisionController = + new WALReplayer.ReplayDecisionController(System.console()); final ReplayStatistics statistics = replayWALDirectories( walFiles, threadNum, - () -> createWALReplayWorker(host, port, username, password, database), + () -> + createWALReplayWorker( + host, port, username, password, database, replayDecisionController), out, deleteSource); out.printf( @@ -353,7 +356,8 @@ private static WALReplayWorker createWALReplayWorker( final int port, final String username, final String password, - final String database) + final String database, + final WALReplayer.ReplayDecisionController replayDecisionController) throws IOException { final Session treeSession = createSession(host, port, username, password, null); final Session tableSession = @@ -363,7 +367,7 @@ private static WALReplayWorker createWALReplayWorker( if (tableSession != null) { tableSession.open(false); } - return new SessionWALReplayer(treeSession, tableSession, database); + return new SessionWALReplayer(treeSession, tableSession, database, replayDecisionController); } catch (final IoTDBConnectionException e) { closeSession(tableSession); closeSession(treeSession); @@ -702,8 +706,11 @@ private static class SessionWALReplayer extends WALReplayer { private final Session tableSession; private SessionWALReplayer( - final Session treeSession, final Session tableSession, final String tableDatabaseName) { - super(treeSession, tableSession, tableDatabaseName); + final Session treeSession, + final Session tableSession, + final String tableDatabaseName, + final WALReplayer.ReplayDecisionController replayDecisionController) { + super(treeSession, tableSession, tableDatabaseName, replayDecisionController); this.treeSession = treeSession; this.tableSession = tableSession; } @@ -720,7 +727,7 @@ static class WALReplayer implements WALReplayWorker { private final Session treeSession; private final Session tableSession; private final ConsensusLogToTabletConverter converter; - private final UnsupportedEntryPrompt unsupportedEntryPrompt; + private final ReplayDecisionPrompt replayDecisionPrompt; private final Map> tableTagSchemas = new HashMap<>(); WALReplayer( @@ -729,17 +736,17 @@ static class WALReplayer implements WALReplayWorker { treeSession, tableSession, tableDatabaseName, - createUnsupportedEntryPrompt(System.console())); + new ReplayDecisionController(System.console())); } WALReplayer( final Session treeSession, final Session tableSession, final String tableDatabaseName, - final UnsupportedEntryPrompt unsupportedEntryPrompt) { + final ReplayDecisionPrompt replayDecisionPrompt) { this.treeSession = treeSession; this.tableSession = tableSession; - this.unsupportedEntryPrompt = unsupportedEntryPrompt; + this.replayDecisionPrompt = replayDecisionPrompt; converter = new ConsensusLogToTabletConverter( null, null, ColumnFilterMatcher.matchAll(), tableDatabaseName); @@ -757,45 +764,111 @@ public boolean replay(final WALEntry entry) return true; } if (entry.getValue() instanceof DeleteDataNode deleteDataNode) { + final ReplayDecision decision = replayDecisionPrompt.decide(entry, true); + if (decision == ReplayDecision.SKIP || decision == ReplayDecision.SKIP_ALL) { + return false; + } + if (decision == ReplayDecision.TERMINATE) { + throw replayTerminatedByUser(); + } replayTreeDelete(deleteDataNode); return true; } if (entry.getValue() instanceof RelationalDeleteDataNode || entry.getValue() instanceof ObjectNode) { - // A null prompt means no interactive console is available, so preserve fail-fast behavior. - if (unsupportedEntryPrompt != null && unsupportedEntryPrompt.shouldSkip(entry)) { + final ReplayDecision decision = replayDecisionPrompt.decide(entry, false); + if (decision == ReplayDecision.SKIP || decision == ReplayDecision.SKIP_ALL) { return false; } + if (decision == ReplayDecision.TERMINATE) { + throw replayTerminatedByUser(); + } throw unsupportedOperation(entry); } return false; } - private static UnsupportedEntryPrompt createUnsupportedEntryPrompt(final Console console) { - if (console == null) { - return null; - } - return entry -> { - synchronized (CONSOLE_PROMPT_LOCK) { - final String answer = - console.readLine( - ImportWALMessages - .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, - entry.getType()); - return isSkipConfirmation(answer); - } - }; + enum ReplayDecision { + EXECUTE, + SKIP, + EXECUTE_ALL, + SKIP_ALL, + TERMINATE } - static boolean isSkipConfirmation(final String answer) { - return answer != null - && ("y".equalsIgnoreCase(answer.trim()) || "yes".equalsIgnoreCase(answer.trim())); + @FunctionalInterface + interface ReplayDecisionPrompt { + + ReplayDecision decide(WALEntry entry, boolean treeDelete); } - @FunctionalInterface - interface UnsupportedEntryPrompt { + static class ReplayDecisionController implements ReplayDecisionPrompt { + + private final Console console; + private ReplayDecision treeDeleteDecision; + private boolean skipAllUnsupportedEntries; - boolean shouldSkip(WALEntry entry); + ReplayDecisionController(final Console console) { + this.console = console; + } + + // The controller is shared by parallel workers so an "all" choice applies to the whole + // import rather than only to the WAL files assigned to one worker. + @Override + public synchronized ReplayDecision decide(final WALEntry entry, final boolean treeDelete) { + if (treeDelete && treeDeleteDecision != null) { + return treeDeleteDecision == ReplayDecision.EXECUTE_ALL + ? ReplayDecision.EXECUTE + : ReplayDecision.SKIP; + } + if (!treeDelete && skipAllUnsupportedEntries) { + return ReplayDecision.SKIP; + } + if (console == null) { + return ReplayDecision.TERMINATE; + } + final String answer = + console.readLine( + treeDelete + ? ImportWALMessages + .MESSAGE_TREE_MODEL_DELETE_OPERATION_DETECTED_ARG_CHOOSE_E_EXECUTE_S_SKIP_A_EXECUTE_ALL_L_SKIP_ALL_Q_QUIT_11E39FD7 + : ImportWALMessages + .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_CHOOSE_S_SKIP_L_SKIP_ALL_Q_QUIT_0A734E52, + entry.getType()); + final ReplayDecision decision = parseDecision(answer, treeDelete); + rememberAllDecision(decision, treeDelete); + return decision; + } + + private void rememberAllDecision(final ReplayDecision decision, final boolean treeDelete) { + if (treeDelete + && (decision == ReplayDecision.EXECUTE_ALL || decision == ReplayDecision.SKIP_ALL)) { + treeDeleteDecision = decision; + } else if (!treeDelete && decision == ReplayDecision.SKIP_ALL) { + skipAllUnsupportedEntries = true; + } + } + + static ReplayDecision parseDecision(final String answer, final boolean treeDelete) { + if (answer == null) { + return ReplayDecision.TERMINATE; + } + return switch (answer.trim().toLowerCase(Locale.ROOT)) { + case "e", "execute", "yes", "y" -> + treeDelete ? ReplayDecision.EXECUTE : ReplayDecision.TERMINATE; + case "s", "skip", "no", "n" -> ReplayDecision.SKIP; + case "a", "all", "execute_all" -> + treeDelete ? ReplayDecision.EXECUTE_ALL : ReplayDecision.TERMINATE; + case "l", "skip_all" -> ReplayDecision.SKIP_ALL; + case "q", "quit", "terminate", "t" -> ReplayDecision.TERMINATE; + default -> ReplayDecision.TERMINATE; + }; + } + } + + private static StatementExecutionException replayTerminatedByUser() { + return new StatementExecutionException( + ImportWALMessages.EXCEPTION_WAL_REPLAY_WAS_TERMINATED_BY_THE_USER_E0BD6197); } private static StatementExecutionException unsupportedOperation(final WALEntry entry) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java index a40f020c18c5a..6cef00311bff6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java @@ -23,6 +23,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ObjectNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalDeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode; import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; @@ -75,6 +76,7 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; @@ -412,23 +414,132 @@ public void testReplayTreeDeletePreservesPathsAndTimeRange() throws Exception { 20); final Session treeSession = mock(Session.class); - new ImportWAL.WALReplayer(treeSession, null, null).replay(new WALInfoEntry(1, deleteNode)); + new ImportWAL.WALReplayer( + treeSession, + null, + null, + (entry, treeDelete) -> ImportWAL.WALReplayer.ReplayDecision.EXECUTE) + .replay(new WALInfoEntry(1, deleteNode)); verify(treeSession) .deleteData(eq(Arrays.asList("root.sg.d1.s1", "root.sg.d2.*")), eq(10L), eq(20L)); } + @Test + public void testReplayTreeDeleteSkipsAfterConfirmation() throws Exception { + final DeleteDataNode deleteNode = + new DeleteDataNode( + new PlanNodeId(""), List.of(new MeasurementPath("root.sg.d1.s1")), 10, 20); + final Session treeSession = mock(Session.class); + + final boolean replayed = + new ImportWAL.WALReplayer( + treeSession, + null, + null, + (entry, treeDelete) -> ImportWAL.WALReplayer.ReplayDecision.SKIP) + .replay(new WALInfoEntry(1, deleteNode)); + + assertFalse(replayed); + verify(treeSession, never()).deleteData(any(), anyLong(), anyLong()); + } + + @Test + public void testReplayTreeDeleteExecuteAllAndSkipAllDecisions() throws Exception { + final DeleteDataNode deleteNode = + new DeleteDataNode( + new PlanNodeId(""), List.of(new MeasurementPath("root.sg.d1.s1")), 10, 20); + final Session treeSession = mock(Session.class); + final AtomicInteger executeAllPromptCount = new AtomicInteger(); + final ImportWAL.WALReplayer.ReplayDecisionPrompt executeAllPrompt = + (entry, treeDelete) -> { + executeAllPromptCount.incrementAndGet(); + return ImportWAL.WALReplayer.ReplayDecision.EXECUTE_ALL; + }; + final ImportWAL.WALReplayer firstExecuteAllReplayer = + new ImportWAL.WALReplayer(treeSession, null, null, executeAllPrompt); + final ImportWAL.WALReplayer secondExecuteAllReplayer = + new ImportWAL.WALReplayer(treeSession, null, null, executeAllPrompt); + + assertTrue(firstExecuteAllReplayer.replay(new WALInfoEntry(1, deleteNode))); + assertTrue(secondExecuteAllReplayer.replay(new WALInfoEntry(2, deleteNode))); + assertEquals(2, executeAllPromptCount.get()); + verify(treeSession, times(2)).deleteData(any(), eq(10L), eq(20L)); + + final Session skippedTreeSession = mock(Session.class); + final AtomicInteger skipAllPromptCount = new AtomicInteger(); + final ImportWAL.WALReplayer.ReplayDecisionPrompt skipAllPrompt = + (entry, treeDelete) -> { + skipAllPromptCount.incrementAndGet(); + return ImportWAL.WALReplayer.ReplayDecision.SKIP_ALL; + }; + final ImportWAL.WALReplayer firstSkipAllReplayer = + new ImportWAL.WALReplayer(skippedTreeSession, null, null, skipAllPrompt); + final ImportWAL.WALReplayer secondSkipAllReplayer = + new ImportWAL.WALReplayer(skippedTreeSession, null, null, skipAllPrompt); + + assertFalse(firstSkipAllReplayer.replay(new WALInfoEntry(1, deleteNode))); + assertFalse(secondSkipAllReplayer.replay(new WALInfoEntry(2, deleteNode))); + assertEquals(2, skipAllPromptCount.get()); + verify(skippedTreeSession, never()).deleteData(any(), anyLong(), anyLong()); + } + + @Test + public void testReplayTreeDeleteTerminatesAfterConfirmation() throws Exception { + final DeleteDataNode deleteNode = + new DeleteDataNode( + new PlanNodeId(""), List.of(new MeasurementPath("root.sg.d1.s1")), 10, 20); + + assertThrows( + StatementExecutionException.class, + () -> + new ImportWAL.WALReplayer( + mock(Session.class), + null, + null, + (entry, treeDelete) -> ImportWAL.WALReplayer.ReplayDecision.TERMINATE) + .replay(new WALInfoEntry(1, deleteNode))); + } + /** Covers an unsupported entry when the interactive user explicitly chooses to skip it. */ @Test public void testReplayUnsupportedEntrySkipsAfterConfirmation() throws Exception { final WALEntry entry = mockUnsupportedEntry(); final boolean replayed = - new ImportWAL.WALReplayer(mock(Session.class), null, null, ignored -> true).replay(entry); + new ImportWAL.WALReplayer( + mock(Session.class), + null, + null, + (ignored, treeDelete) -> ImportWAL.WALReplayer.ReplayDecision.SKIP) + .replay(entry); assertFalse(replayed); } + @Test + public void testReplayUnsupportedEntriesSkipAllAfterConfirmation() throws Exception { + final AtomicInteger promptCount = new AtomicInteger(); + final ImportWAL.WALReplayer.ReplayDecisionPrompt skipAllPrompt = + (entry, treeDelete) -> { + promptCount.incrementAndGet(); + return ImportWAL.WALReplayer.ReplayDecision.SKIP_ALL; + }; + final ImportWAL.WALReplayer relationalDeleteReplayer = + new ImportWAL.WALReplayer(mock(Session.class), null, null, skipAllPrompt); + + assertFalse(relationalDeleteReplayer.replay(mockUnsupportedEntry())); + assertFalse(relationalDeleteReplayer.replay(mockUnsupportedEntry())); + + final WALEntry objectEntry = mock(WALEntry.class); + when(objectEntry.getType()).thenReturn(WALEntryType.OBJECT_FILE_NODE); + when(objectEntry.getValue()).thenReturn(mock(ObjectNode.class)); + final ImportWAL.WALReplayer objectNodeReplayer = + new ImportWAL.WALReplayer(mock(Session.class), null, null, skipAllPrompt); + assertFalse(objectNodeReplayer.replay(objectEntry)); + assertEquals(3, promptCount.get()); + } + /** Covers an unsupported entry when the interactive user declines the skip prompt. */ @Test public void testReplayUnsupportedEntryFailsAfterDecliningSkip() { @@ -437,7 +548,11 @@ public void testReplayUnsupportedEntryFailsAfterDecliningSkip() { assertThrows( StatementExecutionException.class, () -> - new ImportWAL.WALReplayer(mock(Session.class), null, null, ignored -> false) + new ImportWAL.WALReplayer( + mock(Session.class), + null, + null, + (ignored, treeDelete) -> ImportWAL.WALReplayer.ReplayDecision.TERMINATE) .replay(entry)); } @@ -448,17 +563,38 @@ public void testReplayUnsupportedEntryFailsWithoutInteractiveInput() { assertThrows( StatementExecutionException.class, - () -> new ImportWAL.WALReplayer(mock(Session.class), null, null, null).replay(entry)); + () -> + new ImportWAL.WALReplayer( + mock(Session.class), + null, + null, + new ImportWAL.WALReplayer.ReplayDecisionController((java.io.Console) null)) + .replay(entry)); } - /** Covers accepted confirmations and the safe default for all other prompt answers. */ @Test - public void testUnsupportedEntrySkipConfirmationParsing() { - assertTrue(ImportWAL.WALReplayer.isSkipConfirmation("y")); - assertTrue(ImportWAL.WALReplayer.isSkipConfirmation(" YES ")); - assertFalse(ImportWAL.WALReplayer.isSkipConfirmation("n")); - assertFalse(ImportWAL.WALReplayer.isSkipConfirmation("")); - assertFalse(ImportWAL.WALReplayer.isSkipConfirmation(null)); + public void testReplayDecisionParsing() { + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.EXECUTE, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("e", true)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.SKIP, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("s", true)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.EXECUTE_ALL, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("a", true)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.SKIP_ALL, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("l", true)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.TERMINATE, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("a", false)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.SKIP_ALL, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("l", false)); + assertEquals( + ImportWAL.WALReplayer.ReplayDecision.TERMINATE, + ImportWAL.WALReplayer.ReplayDecisionController.parseDecision("q", true)); } /** Covers a non-aligned snapshot whose measurements have independent time axes. */