This is an automated email from the ASF dual-hosted git repository. JackieTien97 pushed a commit to branch rc/2.0.11 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 219f28104cad6b49db570d4039f0ccdf96c3b2d0 Author: Zeyu Zhang <[email protected]> AuthorDate: Thu Jul 16 13:59:12 2026 +0800 fix: optimize reconnect exception as `TTransport.TTransportException` (#18200) --- iotdb-client/client-py/iotdb/Session.py | 7 +- .../tests/unit/test_session_connection.py | 116 +++++++++++++++++++++ 2 files changed, 121 insertions(+), 2 deletions(-) diff --git a/iotdb-client/client-py/iotdb/Session.py b/iotdb-client/client-py/iotdb/Session.py index 0a5cb43fd54..32cdccbe5b9 100644 --- a/iotdb-client/client-py/iotdb/Session.py +++ b/iotdb-client/client-py/iotdb/Session.py @@ -179,7 +179,7 @@ class Session(object): self.__default_connection = self.init_connection( self.__default_endpoint ) - except Exception as e: + except IoTDBConnectionException as e: if not self.reconnect(): if str(e).startswith("Could not connect to any of"): error_msg = ( @@ -240,9 +240,12 @@ class Session(object): session_id = open_resp.sessionId statement_id = client.requestStatementId(session_id) - except Exception as e: + except TTransport.TTransportException as e: transport.close() raise IoTDBConnectionException(e) from None + except Exception: + transport.close() + raise if self.__zone_id is not None: request = TSSetTimeZoneReq(session_id, self.__zone_id) diff --git a/iotdb-client/client-py/tests/unit/test_session_connection.py b/iotdb-client/client-py/tests/unit/test_session_connection.py new file mode 100644 index 00000000000..7a78c697fea --- /dev/null +++ b/iotdb-client/client-py/tests/unit/test_session_connection.py @@ -0,0 +1,116 @@ +# 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. +# + +import importlib +from types import SimpleNamespace + +import pytest +from thrift.transport import TTransport + +from iotdb.thrift.common.ttypes import TSStatus +from iotdb.utils.exception import StatementExecutionException + + +class FakeSocket: + def setTimeout(self, timeout): + self.timeout = timeout + + +class FakeTransport: + def __init__(self, socket): + self.socket = socket + self.opened = False + + def isOpen(self): + return self.opened + + def open(self): + self.opened = True + + def close(self): + self.opened = False + + [email protected]( + "status_code,status_message", + [ + (801, "Authentication failed."), + (822, "Account is blocked due to consecutive failed logins."), + ], +) +def test_session_does_not_retry_iotdb_status_error( + monkeypatch, status_code, status_message +): + session_module = importlib.import_module("iotdb.Session") + + class RejectingClient: + open_session_calls = 0 + + def __init__(self, protocol): + pass + + def openSession(self, request): + RejectingClient.open_session_calls += 1 + return SimpleNamespace( + status=TSStatus(code=status_code, message=status_message) + ) + + monkeypatch.setattr( + session_module.TSocket, "TSocket", lambda host, port: FakeSocket() + ) + monkeypatch.setattr(session_module.TTransport, "TFramedTransport", FakeTransport) + monkeypatch.setattr(session_module, "Client", RejectingClient) + + session = session_module.Session.init_from_node_urls( + ["127.0.0.1:6667"], user="test", password="wrong" + ) + + with pytest.raises(StatementExecutionException) as exc_info: + session.open() + + assert str(exc_info.value) == f"{status_code}: {status_message}" + assert RejectingClient.open_session_calls == 1 + + +def test_session_retries_transport_error(monkeypatch): + session_module = importlib.import_module("iotdb.Session") + + class FailingClient: + open_session_calls = 0 + + def __init__(self, protocol): + pass + + def openSession(self, request): + FailingClient.open_session_calls += 1 + raise TTransport.TTransportException(message="Network is unavailable.") + + monkeypatch.setattr( + session_module.TSocket, "TSocket", lambda host, port: FakeSocket() + ) + monkeypatch.setattr(session_module.TTransport, "TFramedTransport", FakeTransport) + monkeypatch.setattr(session_module, "Client", FailingClient) + + session = session_module.Session.init_from_node_urls( + ["127.0.0.1:6667"], user="test", password="wrong" + ) + + with pytest.raises(session_module.IoTDBConnectionException): + session.open() + + assert FailingClient.open_session_calls == 1 + session.RETRY_NUM
