Skip to content

[IOTDB-17179] CLI: support automatic reconnection when connection is lost #17181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
miantalha45 wants to merge 2 commits into
base: master
Choose a base branch
from
+206 −23
Open

[IOTDB-17179] CLI: support automatic reconnection when connection is lost #17181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9877a70
[IOTDB-XXXX] CLI: support automatic reconnection when connection is lost
miantalha45 Feb 7, 2026
ba46ce3
CLI auto-reconnect: handle session/statement errors after reconnect
miantalha45 Feb 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 83 additions & 5 deletions iotdb-client/cli/src/main/java/org/apache/iotdb/cli/AbstractCli.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,59 @@

public abstract class AbstractCli {

/**
* Returns true if the SQLException is likely due to connection loss. Used so that CLI can rethrow
* and trigger reconnection.
*/
static boolean isConnectionRelated(SQLException e) {
if (e == null) {
return false;
}
if (matchesConnectionFailure(e.getMessage())) {
return true;
}
Throwable cause = e.getCause();
return cause != null && matchesConnectionFailure(cause.getMessage());
}

private static boolean matchesConnectionFailure(String msg) {
if (msg == null) {
return false;
}
String lower = msg.toLowerCase();
return lower.contains("connection")
|| lower.contains("refused")
|| lower.contains("timeout")
|| lower.contains("closed")
|| lower.contains("reset")
|| lower.contains("network")
|| lower.contains("broken pipe");
}

/**
* Returns true if the SQLException indicates a session/statement state error (e.g. statement ID
* no longer valid after reconnect). Used to show a friendly message instead of the raw exception.
*/
static boolean isSessionOrStatementError(SQLException e) {
if (e == null) {
return false;
}
if (e.getMessage() != null && matchesSessionOrStatementFailure(e.getMessage())) {
return true;
}
Throwable cause = e.getCause();
return cause != null
&& cause.getMessage() != null
&& matchesSessionOrStatementFailure(cause.getMessage());
}

private static boolean matchesSessionOrStatementFailure(String msg) {
String lower = msg.toLowerCase();
return lower.contains("doesn't exist in this session")
|| lower.contains("statementid")
|| lower.contains("statement id");
}

static final String HOST_ARGS = "h";
static final String HOST_NAME = "host";

Expand Down Expand Up @@ -395,7 +448,8 @@ static void echoStarting(CliContext ctx) {
ctx.getPrinter().println("---------------------");
}

static OperationResult handleInputCmd(CliContext ctx, String cmd, IoTDBConnection connection) {
static OperationResult handleInputCmd(CliContext ctx, String cmd, IoTDBConnection connection)
throws SQLException {
lastProcessStatus = CODE_OK;
String specialCmd = cmd.toLowerCase().trim();

Expand Down Expand Up @@ -496,7 +550,8 @@ private static int setTimestampDisplay(CliContext ctx, String specialCmd, String
* @return execute result code
*/
private static int setTimeZone(
CliContext ctx, String specialCmd, String cmd, IoTDBConnection connection) {
CliContext ctx, String specialCmd, String cmd, IoTDBConnection connection)
throws SQLException {
String[] values = specialCmd.split("=");
if (values.length != 2) {
ctx.getPrinter()
Expand All @@ -506,6 +561,12 @@ private static int setTimeZone(
}
try {
connection.setTimeZone(cmd.split("=")[1].trim());
} catch (SQLException e) {
if (isConnectionRelated(e)) {
throw e;
}
ctx.getPrinter().println(String.format("Time zone format error: %s", e.getMessage()));
return CODE_ERROR;
} catch (Exception e) {
ctx.getPrinter().println(String.format("Time zone format error: %s", e.getMessage()));
return CODE_ERROR;
Expand All @@ -514,10 +575,13 @@ private static int setTimeZone(
return CODE_OK;
}

private static int showTimeZone(CliContext ctx, IoTDBConnection connection) {
private static int showTimeZone(CliContext ctx, IoTDBConnection connection) throws SQLException {
try {
ctx.getPrinter().println("Current time zone: " + connection.getTimeZone());
} catch (Exception e) {
if (e instanceof SQLException && isConnectionRelated((SQLException) e)) {
throw (SQLException) e;
}
ctx.getPrinter().println("Cannot get time zone from server side because: " + e.getMessage());
return CODE_ERROR;
}
Expand Down Expand Up @@ -550,7 +614,8 @@ private static int importCmd(
}

@SuppressWarnings({"squid:S3776"}) // Suppress high Cognitive Complexity warning
private static int executeQuery(CliContext ctx, IoTDBConnection connection, String cmd) {
private static int executeQuery(CliContext ctx, IoTDBConnection connection, String cmd)
throws SQLException {
int executeStatus = CODE_OK;
long startTime = System.currentTimeMillis();
try (Statement statement = connection.createStatement()) {
Expand Down Expand Up @@ -611,6 +676,18 @@ private static int executeQuery(CliContext ctx, IoTDBConnection connection, Stri
} else {
ctx.getPrinter().println("Msg: " + SUCCESS_MESSAGE);
}
} catch (SQLException e) {
if (isConnectionRelated(e)) {
throw e;
}
if (isSessionOrStatementError(e)) {
ctx.getPrinter()
.println(
"Reconnected, but the previous command could not be completed. Please run your command again.");
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

executeQuery prints "Reconnected, but the previous command could not be completed..." for session/statement-state errors, but this method cannot know whether a reconnect actually happened. If the same server-side error occurs without any reconnect, this message is misleading. Suggestion: make the message not assume reconnection (e.g., "Session state was reset..."), or only print this in the reconnect-and-retry path where reconnection is known to have occurred.

Suggested change
"Reconnected, but the previous command could not be completed. Please run your command again.");
"Session state was reset and the previous command could not be completed. Please run your command again.");

Copilot uses AI. Check for mistakes.
} else {
ctx.getPrinter().println("Msg: " + e);
}
executeStatus = CODE_ERROR;
} catch (Exception e) {
ctx.getPrinter().println("Msg: " + e);
executeStatus = CODE_ERROR;
Expand Down Expand Up @@ -859,7 +936,8 @@ enum OperationResult {
NO_OPER
}

static boolean processCommand(CliContext ctx, String s, IoTDBConnection connection) {
static boolean processCommand(CliContext ctx, String s, IoTDBConnection connection)
throws SQLException {
if (s == null) {
return true;
}
Expand Down
137 changes: 120 additions & 17 deletions iotdb-client/cli/src/main/java/org/apache/iotdb/cli/Cli.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,35 @@ public class Cli extends AbstractCli {
// TODO: Make non-static
private static final Properties info = new Properties();

/** Number of reconnection attempts when connection is lost during interactive session. */
private static final int RECONNECT_RETRY_NUM = 3;

/** Delay in ms between reconnection attempts. */
private static final long RECONNECT_RETRY_INTERVAL_MS = 1000;

/** Result of reading and processing one line; used to support reconnection. */
private static class ReadLineResult {
final boolean stop;
final String failedCommand;

ReadLineResult(boolean stop, String failedCommand) {
this.stop = stop;
this.failedCommand = failedCommand;
}

static ReadLineResult continueLoop() {
return new ReadLineResult(false, null);
}

static ReadLineResult stopLoop() {
return new ReadLineResult(true, null);
}

static ReadLineResult reconnectAndRetry(String command) {
return new ReadLineResult(false, command);
}
}

/**
* IoTDB Client main function.
*
Expand Down Expand Up @@ -155,6 +184,29 @@ private static boolean parseCommandLine(
return true;
}

private static IoTDBConnection openConnection() throws SQLException {
return (IoTDBConnection)
DriverManager.getConnection(Config.IOTDB_URL_PREFIX + host + ":" + port + "/", info);
}

private static void setupConnection(IoTDBConnection connection)
throws java.sql.SQLException, org.apache.thrift.TException {
connection.setQueryTimeout(queryTimeout);
properties = connection.getServerProperties();
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

setupConnection() appends to AGGREGRATE_TIME_LIST every time a reconnect happens, but the list is static and never cleared/deduped. During repeated reconnects this will grow unbounded and may introduce duplicate entries. Suggestion: clear the list before adding, or change it to a Set (if ordering is not required).

Suggested change
properties = connection.getServerProperties();
properties = connection.getServerProperties();
AGGREGRATE_TIME_LIST.clear();

Copilot uses AI. Check for mistakes.
AGGREGRATE_TIME_LIST.addAll(properties.getSupportedTimeAggregationOperations());
timestampPrecision = properties.getTimestampPrecision();
}

private static void closeConnectionQuietly(IoTDBConnection connection) {
if (connection != null) {
try {
connection.close();
} catch (SQLException ignored) {
// ignore
}
}
}

private static void serve(CliContext ctx) {
try {
useSsl = commandLine.getOptionValue(USE_SSL_ARGS);
Expand Down Expand Up @@ -195,36 +247,87 @@ private static void executeSql(CliContext ctx) throws TException {
}

private static void receiveCommands(CliContext ctx) throws TException {
try (IoTDBConnection connection =
(IoTDBConnection)
DriverManager.getConnection(Config.IOTDB_URL_PREFIX + host + ":" + port + "/", info)) {
connection.setQueryTimeout(queryTimeout);
properties = connection.getServerProperties();
AGGREGRATE_TIME_LIST.addAll(properties.getSupportedTimeAggregationOperations());
timestampPrecision = properties.getTimestampPrecision();

IoTDBConnection connection = null;
try {
connection = openConnection();
setupConnection(connection);
echoStarting(ctx);
displayLogo(ctx, properties.getLogo(), properties.getVersion(), properties.getBuildInfo());
ctx.getPrinter().println(String.format("Successfully login at %s:%s", host, port));
while (true) {
boolean readLine = readerReadLine(ctx, connection);
if (readLine) {
ReadLineResult result = readerReadLine(ctx, connection);
if (result.stop) {
break;
}
if (result.failedCommand != null) {
// Connection failed during processCommand; try to reconnect and retry the command.
closeConnectionQuietly(connection);
connection = null;
boolean reconnected = false;
for (int attempt = 1; attempt <= RECONNECT_RETRY_NUM; attempt++) {
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Test is always true, because of this condition.
Test is always true, because of this condition.

Copilot uses AI. Check for mistakes.
if (attempt > 1) {
try {
Thread.sleep(RECONNECT_RETRY_INTERVAL_MS);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
ctx.getErr().printf("%s: Reconnection interrupted.%n", IOTDB_ERROR_PREFIX);
ctx.exit(CODE_ERROR);
}
}
try {
connection = openConnection();
setupConnection(connection);
ctx.getPrinter().println("Connection lost. Reconnected. Retrying command.");
processCommand(ctx, result.failedCommand, connection);
reconnected = true;
break;
Comment on lines +277 to +283
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reconnect retry try { connection = openConnection(); setupConnection(connection); ... } can throw TException from setupConnection(), but the retry loop only catches SQLException. This means some connection-loss cases can bypass the reconnect logic and bubble out immediately. Suggestion: include TException (or a common supertype) in the reconnect-attempt catch and handle it as a reconnect failure (with the same retry/backoff), while still distinguishing non-connection SQL errors from command execution.

Copilot uses AI. Check for mistakes.
Comment on lines +280 to +283
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the reconnect-and-retry path, the return value of processCommand(...) is ignored. If the retried input contains an exit/quit statement (or otherwise returns false), the CLI will continue the main loop instead of stopping as it normally would. Suggestion: capture the boolean return from processCommand and convert it into a result.stop / break out of the outer loop when it indicates the user requested exit.

Copilot uses AI. Check for mistakes.
} catch (SQLException e) {
if (isSessionOrStatementError(e)) {
// Reconnect succeeded but retry failed due to session/statement state; ask user to
// run the command again.
ctx.getPrinter()
.println(
"Reconnected, but the previous command could not be completed. Please run your command again.");
reconnected = true;
break;
}
if (attempt == RECONNECT_RETRY_NUM) {
ctx.getErr()
.printf(
"%s: Could not reconnect after %d attempts. Please check that the server is running and try again.%n",
IOTDB_ERROR_PREFIX, RECONNECT_RETRY_NUM);
ctx.exit(CODE_ERROR);
}
Comment on lines +284 to +300
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the reconnect loop, any SQLException thrown while retrying the failed command is treated as a reconnect failure. This can incorrectly trigger further reconnect attempts (and eventually exit with "Could not reconnect") for non-connection SQL errors (e.g., syntax/permission errors) that happen after the reconnect succeeds. Suggestion: only retry when isConnectionRelated(e) is true; otherwise report the SQL error from the retried command and continue the main loop (or stop) without attempting further reconnects.

Copilot uses AI. Check for mistakes.
}
Comment on lines +277 to +301
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On reconnect attempts, if openConnection() succeeds but setupConnection() or processCommand() throws, the newly created connection is not closed before the next attempt (the variable is overwritten on the next iteration). This can leak sockets/sessions across retries. Suggestion: in the catch path, call closeConnectionQuietly(connection) and set connection = null before continuing to the next attempt.

Copilot uses AI. Check for mistakes.
}
if (!reconnected) {
break;
}
}
}
} catch (SQLException e) {
ctx.getErr().printf("%s: %s%n", IOTDB_ERROR_PREFIX, e.getMessage());
ctx.exit(CODE_ERROR);
} finally {
closeConnectionQuietly(connection);
}
}

private static boolean readerReadLine(CliContext ctx, IoTDBConnection connection) {
private static ReadLineResult readerReadLine(CliContext ctx, IoTDBConnection connection) {
String s;
try {
s = ctx.getLineReader().readLine(cliPrefix + "> ", null);
boolean continues = processCommand(ctx, s, connection);
if (!continues) {
return true;
try {
boolean continues = processCommand(ctx, s, connection);
if (!continues) {
return ReadLineResult.stopLoop();
}
} catch (SQLException e) {
if (isConnectionRelated(e)) {
return ReadLineResult.reconnectAndRetry(s);
}
Comment on lines +325 to +328
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reconnection trigger/loop is a significant behavior change but there are no unit tests exercising it (e.g., simulating a connection-related SQLException from processCommand, verifying retry count/backoff, and ensuring non-connection SQLExceptions do not trigger reconnect or an incorrect "Could not reconnect" exit). Consider adding focused tests around this new flow to prevent regressions.

Copilot uses AI. Check for mistakes.
Comment on lines 319 to +328
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When a connection-related SQLException occurs, the CLI retries the entire raw input line s. If the line contains multiple statements separated by ;, some statements may already have executed successfully before the failure, and retrying the whole line can re-run those statements (duplicate writes / side effects). Suggestion: either only retry the specific statement that failed (track progress inside processCommand), or disable auto-retry for multi-statement input and ask the user to rerun manually.

Copilot uses AI. Check for mistakes.
ctx.getErr().printf("%s: %s%n", IOTDB_ERROR_PREFIX, e.getMessage());
return ReadLineResult.stopLoop();
}
} catch (UserInterruptException e) {
// Exit on signal INT requires confirmation.
Expand All @@ -233,12 +336,12 @@ private static boolean readerReadLine(CliContext ctx, IoTDBConnection connection
// Exit on EOF (usually by pressing CTRL+D).
ctx.exit(CODE_OK);
} catch (IllegalArgumentException e) {
if (e.getMessage().contains("history")) {
return false;
if (e.getMessage() != null && e.getMessage().contains("history")) {
return ReadLineResult.continueLoop();
}
throw e;
}
return false;
return ReadLineResult.continueLoop();
}

private static void readLine(CliContext ctx) {
Expand Down
4 changes: 3 additions & 1 deletion iotdb-client/cli/src/test/java/org/apache/iotdb/cli/AbstractCliTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.when;
Expand Down Expand Up @@ -150,7 +152,7 @@ private void isTwoStringArrayEqual(String[] expected, String[] actual) {
}

@Test
public void testHandleInputInputCmd() {
public void testHandleInputInputCmd() throws SQLException {
CliContext ctx = new CliContext(System.in, System.out, System.err, ExitType.EXCEPTION);
assertEquals(
OperationResult.STOP_OPER,
Expand Down
Loading