Skip to content

Commit ababbf3

Browse files
authored
Support SSL-related parameters and interactive password input in export/import tools (apache#16698)
* Add ssl param + Support interactive pw input * fix spotless * Changed params processing order: ts > tpw > pw
1 parent f096b82 commit ababbf3

File tree

12 files changed

+190
-37
lines changed

12 files changed

+190
-37
lines changed

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/Constants.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,18 @@ public class Constants {
5656
public static final String USERNAME_DESC = "Username (optional)";
5757
public static final String USERNAME_DEFAULT_VALUE = "root";
5858

59+
public static final String USE_SSL_ARGS = "usessl";
60+
public static final String USE_SSL_NAME = "use_ssl";
61+
public static final String USE_SSL_DESC = "Use SSL statement. (optional)";
62+
63+
public static final String TRUST_STORE_ARGS = "ts";
64+
public static final String TRUST_STORE_NAME = "trust_store";
65+
public static final String TRUST_STORE_DESC = "Trust store. (optional)";
66+
67+
public static final String TRUST_STORE_PWD_ARGS = "tpw";
68+
public static final String TRUST_STORE_PWD_NAME = "trust_store_password";
69+
public static final String TRUST_STORE_PWD_DESC = "Trust store password. (optional)";
70+
5971
public static final String FILE_TYPE_ARGS = "ft";
6072
public static final String FILE_TYPE_NAME = "file_type";
6173
public static final String FILE_TYPE_ARGS_NAME = "format";

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/common/OptionsUtil.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,36 @@ public static Options createCommonOptions(Options options) {
103103
.build();
104104
options.addOption(opPassword);
105105

106+
Option opUseSSL =
107+
Option.builder(USE_SSL_ARGS)
108+
.longOpt(USE_SSL_NAME)
109+
.optionalArg(true)
110+
.argName(USE_SSL_NAME)
111+
.hasArg()
112+
.desc(USE_SSL_DESC)
113+
.build();
114+
options.addOption(opUseSSL);
115+
116+
Option opTrustStore =
117+
Option.builder(TRUST_STORE_ARGS)
118+
.longOpt(TRUST_STORE_NAME)
119+
.optionalArg(true)
120+
.argName(TRUST_STORE_NAME)
121+
.hasArg()
122+
.desc(TRUST_STORE_DESC)
123+
.build();
124+
options.addOption(opTrustStore);
125+
126+
Option opTrustStorePwd =
127+
Option.builder(TRUST_STORE_PWD_ARGS)
128+
.longOpt(TRUST_STORE_PWD_NAME)
129+
.optionalArg(true)
130+
.argName(TRUST_STORE_PWD_NAME)
131+
.hasArg()
132+
.desc(TRUST_STORE_PWD_DESC)
133+
.build();
134+
options.addOption(opTrustStorePwd);
135+
106136
return options;
107137
}
108138

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/AbstractDataTool.java

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@
1919

2020
package org.apache.iotdb.tool.data;
2121

22+
import org.apache.iotdb.cli.type.ExitType;
23+
import org.apache.iotdb.cli.utils.CliContext;
2224
import org.apache.iotdb.cli.utils.IoTPrinter;
25+
import org.apache.iotdb.cli.utils.JlineUtils;
2326
import org.apache.iotdb.commons.exception.IllegalPathException;
2427
import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant;
2528
import org.apache.iotdb.commons.utils.PathUtils;
@@ -50,6 +53,7 @@
5053
import org.apache.tsfile.read.common.Field;
5154
import org.apache.tsfile.read.common.RowRecord;
5255
import org.apache.tsfile.utils.Binary;
56+
import org.jline.reader.LineReader;
5357
import org.slf4j.Logger;
5458
import org.slf4j.LoggerFactory;
5559

@@ -86,6 +90,9 @@ public abstract class AbstractDataTool {
8690
protected static String endTime;
8791
protected static String username;
8892
protected static String password;
93+
protected static Boolean useSsl;
94+
protected static String trustStore;
95+
protected static String trustStorePwd;
8996
protected static Boolean aligned;
9097
protected static String database;
9198
protected static String startTime;
@@ -142,7 +149,8 @@ protected static String checkRequiredArg(
142149
return str;
143150
}
144151

145-
protected static void parseBasicParams(CommandLine commandLine) throws ArgsErrorException {
152+
protected static void parseBasicParams(CommandLine commandLine)
153+
throws ArgsErrorException, IOException {
146154
host =
147155
checkRequiredArg(
148156
Constants.HOST_ARGS, Constants.HOST_NAME, commandLine, Constants.HOST_DEFAULT_VALUE);
@@ -155,7 +163,36 @@ protected static void parseBasicParams(CommandLine commandLine) throws ArgsError
155163
Constants.USERNAME_NAME,
156164
commandLine,
157165
Constants.USERNAME_DEFAULT_VALUE);
158-
password = commandLine.getOptionValue(Constants.PW_ARGS, Constants.PW_DEFAULT_VALUE);
166+
CliContext cliCtx = new CliContext(System.in, System.out, System.err, ExitType.SYSTEM_EXIT);
167+
LineReader lineReader = JlineUtils.getLineReader(cliCtx, username, host, port);
168+
cliCtx.setLineReader(lineReader);
169+
String useSslStr = commandLine.getOptionValue(Constants.USE_SSL_ARGS);
170+
useSsl = Boolean.parseBoolean(useSslStr);
171+
if (useSsl) {
172+
String givenTS = commandLine.getOptionValue(Constants.TRUST_STORE_ARGS);
173+
if (givenTS != null) {
174+
trustStore = givenTS;
175+
} else {
176+
trustStore = cliCtx.getLineReader().readLine("please input your trust_store:", '\0');
177+
}
178+
String givenTPW = commandLine.getOptionValue(Constants.TRUST_STORE_PWD_ARGS);
179+
if (givenTPW != null) {
180+
trustStorePwd = givenTPW;
181+
} else {
182+
trustStorePwd = cliCtx.getLineReader().readLine("please input your trust_store_pwd:", '\0');
183+
}
184+
}
185+
boolean hasPw = commandLine.hasOption(Constants.PW_ARGS);
186+
if (hasPw) {
187+
String inputPassword = commandLine.getOptionValue(Constants.PW_ARGS);
188+
if (inputPassword != null) {
189+
password = inputPassword;
190+
} else {
191+
password = cliCtx.getLineReader().readLine("please input your password:", '\0');
192+
}
193+
} else {
194+
password = Constants.PW_DEFAULT_VALUE;
195+
}
159196
}
160197

161198
protected static void printHelpOptions(

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTable.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,18 @@ public class ExportDataTable extends AbstractExportData {
6363

6464
@Override
6565
public void init() throws IoTDBConnectionException, StatementExecutionException {
66-
tableSession =
66+
TableSessionBuilder tableSessionBuilder =
6767
new TableSessionBuilder()
6868
.nodeUrls(Collections.singletonList(host + ":" + port))
6969
.username(username)
7070
.password(password)
7171
.database(database)
72-
.thriftMaxFrameSize(rpcMaxFrameSize)
73-
.build();
72+
.thriftMaxFrameSize(rpcMaxFrameSize);
73+
if (useSsl) {
74+
tableSessionBuilder =
75+
tableSessionBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
76+
}
77+
tableSession = tableSessionBuilder.build();
7478
SessionDataSet sessionDataSet = tableSession.executeQueryStatement("show databases", timeout);
7579
List<String> databases = new ArrayList<>();
7680
while (sessionDataSet.hasNext()) {

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ExportDataTree.java

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -67,18 +67,23 @@ public class ExportDataTree extends AbstractExportData {
6767

6868
@Override
6969
public void init() throws IoTDBConnectionException, StatementExecutionException, TException {
70-
session =
71-
new Session(
72-
host,
73-
Integer.parseInt(port),
74-
username,
75-
password,
76-
SessionConfig.DEFAULT_FETCH_SIZE,
77-
null,
78-
SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY,
79-
rpcMaxFrameSize,
80-
SessionConfig.DEFAULT_REDIRECTION_MODE,
81-
SessionConfig.DEFAULT_VERSION);
70+
Session.Builder sessionBuilder =
71+
new Session.Builder()
72+
.host(host)
73+
.port(Integer.parseInt(port))
74+
.username(username)
75+
.password(password)
76+
.fetchSize(SessionConfig.DEFAULT_FETCH_SIZE)
77+
.zoneId(null)
78+
.thriftDefaultBufferSize(SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY)
79+
.thriftMaxFrameSize(rpcMaxFrameSize)
80+
.enableRedirection(SessionConfig.DEFAULT_REDIRECTION_MODE)
81+
.version(SessionConfig.DEFAULT_VERSION);
82+
if (useSsl) {
83+
sessionBuilder =
84+
sessionBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
85+
}
86+
session = sessionBuilder.build();
8287
session.open(false);
8388
timestampPrecision = session.getTimestampPrecision();
8489
if (timeZoneID != null) {

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTable.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ public class ImportDataTable extends AbstractImportData {
6666
private static Map<String, ColumnCategory> columnCategory = new HashMap<>();
6767

6868
public void init() throws InterruptedException {
69-
sessionPool =
69+
TableSessionPoolBuilder tableSessionPoolBuilder =
7070
new TableSessionPoolBuilder()
7171
.nodeUrls(Collections.singletonList(host + ":" + port))
7272
.user(username)
@@ -75,8 +75,12 @@ public void init() throws InterruptedException {
7575
.enableThriftCompression(false)
7676
.enableRedirection(false)
7777
.enableAutoFetch(false)
78-
.database(database)
79-
.build();
78+
.database(database);
79+
if (useSsl) {
80+
tableSessionPoolBuilder =
81+
tableSessionPoolBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
82+
}
83+
sessionPool = tableSessionPoolBuilder.build();
8084
final File file = new File(targetPath);
8185
if (!file.isFile() && !file.isDirectory()) {
8286
ioTPrinter.println(String.format("Source file or directory %s does not exist", targetPath));

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/data/ImportDataTree.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ public class ImportDataTree extends AbstractImportData {
6161

6262
public void init()
6363
throws InterruptedException, IoTDBConnectionException, StatementExecutionException {
64-
sessionPool =
64+
SessionPool.Builder sessionPoolBuilder =
6565
new SessionPool.Builder()
6666
.host(host)
6767
.port(Integer.parseInt(port))
@@ -70,8 +70,12 @@ public void init()
7070
.maxSize(threadNum + 1)
7171
.enableIoTDBRpcCompression(false)
7272
.enableRedirection(false)
73-
.enableAutoFetch(false)
74-
.build();
73+
.enableAutoFetch(false);
74+
if (useSsl) {
75+
sessionPoolBuilder =
76+
sessionPoolBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
77+
}
78+
sessionPool = sessionPoolBuilder.build();
7579
sessionPool.setEnableQueryRedirection(false);
7680
if (timeZoneID != null) {
7781
sessionPool.setTimeZone(timeZoneID);

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/schema/AbstractSchemaTool.java

Lines changed: 39 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@
1919

2020
package org.apache.iotdb.tool.schema;
2121

22+
import org.apache.iotdb.cli.type.ExitType;
23+
import org.apache.iotdb.cli.utils.CliContext;
2224
import org.apache.iotdb.cli.utils.IoTPrinter;
25+
import org.apache.iotdb.cli.utils.JlineUtils;
2326
import org.apache.iotdb.exception.ArgsErrorException;
2427
import org.apache.iotdb.session.Session;
2528
import org.apache.iotdb.tool.common.Constants;
@@ -29,6 +32,7 @@
2932
import org.apache.commons.csv.CSVPrinter;
3033
import org.apache.commons.csv.QuoteMode;
3134
import org.apache.tsfile.external.commons.lang3.StringUtils;
35+
import org.jline.reader.LineReader;
3236
import org.slf4j.Logger;
3337
import org.slf4j.LoggerFactory;
3438

@@ -46,6 +50,9 @@ public abstract class AbstractSchemaTool {
4650
protected static String database;
4751
protected static String username;
4852
protected static String password;
53+
protected static Boolean useSsl;
54+
protected static String trustStore;
55+
protected static String trustStorePwd;
4956
protected static Session session;
5057
protected static String queryPath;
5158
protected static int threadNum = 8;
@@ -80,7 +87,8 @@ protected static String checkRequiredArg(
8087
return str;
8188
}
8289

83-
protected static void parseBasicParams(CommandLine commandLine) throws ArgsErrorException {
90+
protected static void parseBasicParams(CommandLine commandLine)
91+
throws ArgsErrorException, IOException {
8492
host =
8593
checkRequiredArg(
8694
Constants.HOST_ARGS, Constants.HOST_NAME, commandLine, Constants.HOST_DEFAULT_VALUE);
@@ -93,9 +101,36 @@ protected static void parseBasicParams(CommandLine commandLine) throws ArgsError
93101
Constants.USERNAME_NAME,
94102
commandLine,
95103
Constants.USERNAME_DEFAULT_VALUE);
96-
password =
97-
checkRequiredArg(
98-
Constants.PW_ARGS, Constants.PW_NAME, commandLine, Constants.PW_DEFAULT_VALUE);
104+
CliContext cliCtx = new CliContext(System.in, System.out, System.err, ExitType.SYSTEM_EXIT);
105+
LineReader lineReader = JlineUtils.getLineReader(cliCtx, username, host, port);
106+
cliCtx.setLineReader(lineReader);
107+
String useSslStr = commandLine.getOptionValue(Constants.USE_SSL_ARGS);
108+
useSsl = Boolean.parseBoolean(useSslStr);
109+
if (useSsl) {
110+
String givenTS = commandLine.getOptionValue(Constants.TRUST_STORE_ARGS);
111+
if (givenTS != null) {
112+
trustStore = givenTS;
113+
} else {
114+
trustStore = cliCtx.getLineReader().readLine("please input your trust_store:", '\0');
115+
}
116+
String givenTPW = commandLine.getOptionValue(Constants.TRUST_STORE_PWD_ARGS);
117+
if (givenTPW != null) {
118+
trustStorePwd = givenTPW;
119+
} else {
120+
trustStorePwd = cliCtx.getLineReader().readLine("please input your trust_store_pwd:", '\0');
121+
}
122+
}
123+
boolean hasPw = commandLine.hasOption(Constants.PW_ARGS);
124+
if (hasPw) {
125+
String inputPassword = commandLine.getOptionValue(Constants.PW_ARGS);
126+
if (inputPassword != null) {
127+
password = inputPassword;
128+
} else {
129+
password = cliCtx.getLineReader().readLine("please input your password:", '\0');
130+
}
131+
} else {
132+
password = Constants.PW_DEFAULT_VALUE;
133+
}
99134
}
100135

101136
/**

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/schema/ExportSchemaTable.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class ExportSchemaTable extends AbstractExportSchema {
5252
private static Map<String, String> tableCommentList = new HashMap<>();
5353

5454
public void init() throws InterruptedException {
55-
sessionPool =
55+
TableSessionPoolBuilder tableSessionPoolBuilder =
5656
new TableSessionPoolBuilder()
5757
.nodeUrls(Collections.singletonList(host + ":" + port))
5858
.user(username)
@@ -61,8 +61,12 @@ public void init() throws InterruptedException {
6161
.enableThriftCompression(false)
6262
.enableRedirection(false)
6363
.enableAutoFetch(false)
64-
.database(database)
65-
.build();
64+
.database(database);
65+
if (useSsl) {
66+
tableSessionPoolBuilder =
67+
tableSessionPoolBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
68+
}
69+
sessionPool = tableSessionPoolBuilder.build();
6670
checkDatabase();
6771
try {
6872
parseTablesBySelectSchema(String.format(Constants.EXPORT_SCHEMA_TABLES_SELECT, database));

iotdb-client/cli/src/main/java/org/apache/iotdb/tool/schema/ExportSchemaTree.java

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,17 @@ public class ExportSchemaTree extends AbstractExportSchema {
4343

4444
public void init()
4545
throws InterruptedException, IoTDBConnectionException, StatementExecutionException {
46-
session = new Session(host, Integer.parseInt(port), username, password);
46+
Session.Builder sessionBuilder =
47+
new Session.Builder()
48+
.host(host)
49+
.port(Integer.parseInt(port))
50+
.username(username)
51+
.password(password);
52+
if (useSsl) {
53+
sessionBuilder =
54+
sessionBuilder.useSSL(true).trustStore(trustStore).trustStorePwd(trustStorePwd);
55+
}
56+
session = sessionBuilder.build();
4757
session.open(false);
4858
}
4959

0 commit comments

Comments
 (0)