CloudWise-Lukemiao commented on code in PR #14677: URL: https://github.com/apache/iotdb/pull/14677#discussion_r1918238965
########## iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java: ########## @@ -0,0 +1,412 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.tool.data; + +import org.apache.iotdb.cli.utils.IoTPrinter; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; +import org.apache.iotdb.tool.common.Constants; +import org.apache.iotdb.tool.tsfile.ImportTsFileScanTool; + +import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.collections4.MapUtils; +import org.apache.commons.csv.CSVParser; +import org.apache.commons.csv.CSVRecord; +import org.apache.commons.lang3.ObjectUtils; +import org.apache.commons.lang3.StringUtils; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; + +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; +import java.io.FileWriter; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +public class ImportDataTable extends AbstractImportData { + + private static final IoTPrinter ioTPrinter = new IoTPrinter(System.out); + private static ITableSessionPool sessionPool; + private static Map<String, TSDataType> dataTypes = new HashMap<>(); + private static Map<String, Tablet.ColumnCategory> columnCategory = new HashMap<>(); + + public void init() throws InterruptedException { + sessionPool = + new TableSessionPoolBuilder() + .nodeUrls(Collections.singletonList(host + ":" + port)) + .user(username) + .password(password) + .maxSize(threadNum + 1) + .enableCompression(false) + .enableRedirection(false) + .enableAutoFetch(false) + .database(database) + .build(); + final File file = new File(targetPath); + if (!file.isFile() && !file.isDirectory()) { + ioTPrinter.println(String.format("Source file or directory %s does not exist", targetPath)); + System.exit(Constants.CODE_ERROR); + } + // checkDataBase + SessionDataSet sessionDataSet = null; + ITableSession session = null; + try { + List<String> databases = new ArrayList<>(); + session = sessionPool.getSession(); + sessionDataSet = session.executeQueryStatement("show databases"); + while (sessionDataSet.hasNext()) { + RowRecord rowRecord = sessionDataSet.next(); + databases.add(rowRecord.getField(0).getStringValue()); + } + if (!databases.contains(database)) { + ioTPrinter.println(String.format(Constants.TARGET_DATABASE_NOT_EXIST_MSG, database)); + System.exit(1); + } + if (Constants.CSV_SUFFIXS.equals(fileType)) { + if (StringUtils.isNotBlank(table)) { + sessionDataSet = session.executeQueryStatement("show tables"); + List<String> tables = new ArrayList<>(); + while (sessionDataSet.hasNext()) { + RowRecord rowRecord = sessionDataSet.next(); + tables.add(rowRecord.getField(0).getStringValue()); + } + if (!tables.contains(table)) { + ioTPrinter.println(String.format(Constants.TARGET_TABLE_NOT_EXIST_MSG, table)); + System.exit(1); + } + sessionDataSet = session.executeQueryStatement("describe " + table); + while (sessionDataSet.hasNext()) { + RowRecord rowRecord = sessionDataSet.next(); + final String columnName = rowRecord.getField(0).getStringValue(); + final String category = rowRecord.getField(2).getStringValue(); + if (!timeColumn.equalsIgnoreCase(category)) { + dataTypes.put(columnName, getType(rowRecord.getField(1).getStringValue())); + columnCategory.put(columnName, getColumnCategory(category)); + } + } + } else { + ioTPrinter.println(String.format(Constants.TARGET_TABLE_NOT_EXIST_MSG, null)); + System.exit(1); + } + } + } catch (StatementExecutionException e) { + ioTPrinter.println(Constants.INSERT_CSV_MEET_ERROR_MSG + e.getMessage()); + System.exit(1); + } catch (IoTDBConnectionException e) { + throw new RuntimeException(e); + } finally { + if (ObjectUtils.isNotEmpty(sessionDataSet)) { + try { + sessionDataSet.close(); + } catch (Exception e) { + ; + } + } + if (ObjectUtils.isNotEmpty(session)) { + try { + session.close(); + } catch (IoTDBConnectionException e) { + ; + } + } + } + if (Constants.TSFILE_SUFFIXS.equalsIgnoreCase(fileType)) { + ImportTsFileScanTool.setSourceFullPath(targetPath); + ImportTsFileScanTool.traverseAndCollectFiles(); + ImportTsFileScanTool.addNoResourceOrModsToQueue(); + } else { + ImportDataScanTool.setSourceFullPath(targetPath); + ImportDataScanTool.traverseAndCollectFiles(); + } + } + + @Override + protected Runnable getAsyncImportRunnable() { + return new ImportDataTable(); // 返回子类1的Runnable对象 + } + + protected static void processSuccessFile() { + loadFileSuccessfulNum.increment(); + } + + @SuppressWarnings("java:S2259") + protected void importFromSqlFile(File file) { + ArrayList<List<Object>> failedRecords = new ArrayList<>(); + String failedFilePath; + if (failedFileDirectory == null) { + failedFilePath = file.getAbsolutePath() + ".failed"; + } else { + failedFilePath = failedFileDirectory + file.getName() + ".failed"; + } + ITableSession session = null; + try (BufferedReader br = new BufferedReader(new FileReader(file.getAbsolutePath()))) { + String sql; + while ((sql = br.readLine()) != null) { + try { + sql = sql.replace(";", ""); + session = sessionPool.getSession(); + session.executeNonQueryStatement(sql); + } catch (IoTDBConnectionException | StatementExecutionException e) { + failedRecords.add(Arrays.asList(sql)); + } finally { + if (ObjectUtils.isNotEmpty(session)) { + try { + session.close(); + } catch (IoTDBConnectionException e) { + ; + } + } + } + } + processSuccessFile(); + } catch (IOException e) { + ioTPrinter.println("SQL file read exception because: " + e.getMessage()); + } + if (!failedRecords.isEmpty()) { + FileWriter writer = null; + try { + writer = new FileWriter(failedFilePath); + for (List<Object> failedRecord : failedRecords) { + writer.write(failedRecord.get(0).toString() + "\n"); + } + } catch (IOException e) { + ioTPrinter.println("Cannot dump fail result because: " + e.getMessage()); + } finally { + if (ObjectUtils.isNotEmpty(writer)) { + try { + writer.flush(); + writer.close(); + } catch (IOException e) { + ; + } + } + } + } + } + + protected void importFromTsFile(File file) { + final String sql = "load '" + file + "'"; + try { + sessionPool.getSession().executeNonQueryStatement(sql); + processSuccessFile(file.getPath()); + } catch (final Exception e) { + processFailFile(file.getPath(), e); + } + } + + protected void importFromCsvFile(File file) { + if (file.getName().endsWith(Constants.CSV_SUFFIXS) + || file.getName().endsWith(Constants.TXT_SUFFIXS)) { + try { + CSVParser csvRecords = readCsvFile(file.getAbsolutePath()); + List<String> headerNames = csvRecords.getHeaderNames(); + Stream<CSVRecord> records = csvRecords.stream(); + if (headerNames.isEmpty()) { + ioTPrinter.println("Empty file!"); + return; + } + if (!timeColumn.toLowerCase().equalsIgnoreCase(filterBomHeader(headerNames.get(0)))) { + ioTPrinter.println("The first field of header must be `time`!"); + return; + } + String failedFilePath; + if (failedFileDirectory == null) { + failedFilePath = file.getAbsolutePath() + ".failed"; + } else { + failedFilePath = failedFileDirectory + file.getName() + ".failed"; + } + writeData(headerNames, records.collect(Collectors.toList()), failedFilePath); + processSuccessFile(); + } catch (IOException e) { + ioTPrinter.println("CSV file read exception because: " + e.getMessage()); + } + } else { + ioTPrinter.println("The file name must end with \"csv\" or \"txt\"!"); + } + } + + protected void writeData( + List<String> headerNames, List<CSVRecord> records, String failedFilePath) { + Map<String, TSDataType> headerTypeMap = new HashMap<>(); + Map<String, String> headerNameMap = new HashMap<>(); + parseHeaders(headerNames, headerTypeMap, headerNameMap); + queryType(headerTypeMap); + List<List<Object>> failedRecords = new ArrayList<>(); + List<IMeasurementSchema> schemas = new ArrayList<>(); + headerNames.stream() + .forEach( + t -> { + if (dataTypes.keySet().contains(t)) { + schemas.add(new MeasurementSchema(t, dataTypes.get(t))); + } + }); + List<String> headNames = new LinkedList<>(dataTypes.keySet()); + List<TSDataType> columnTypes = new LinkedList<>(dataTypes.values()); + List<Tablet.ColumnCategory> columnCategorys = new LinkedList<>(columnCategory.values()); + Tablet tablet = new Tablet(table, headNames, columnTypes, columnCategorys); + for (CSVRecord recordObj : records) { + boolean isFail = false; + final int rowSize = tablet.getRowSize(); + long rowTimeStamp = parseTimestamp(recordObj.get(0)); + for (String headerName : headerNameMap.keySet()) { + String value = recordObj.get(headerNameMap.get(headerName)); + if (!"".equals(value)) { + TSDataType type; + if (timeColumn.equalsIgnoreCase(headerName)) { + // time列 + tablet.addTimestamp(rowSize, rowTimeStamp); + continue; + } else if (!headerTypeMap.containsKey(headerName)) { + // 未定义列 Review Comment: english -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
