From 004320f4432ca86fc701c7cc0ad8c8793980f618 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Tue, 13 Oct 2026 11:40:00 +0800 Subject: [PATCH] Python: fix swapped session and statement IDs in query result sets --- iotdb-client/client-py/iotdb/Session.py | 128 ++++++++--------- .../tests/unit/test_session_query_ids.py | 132 ++++++++++++++++++ 2 files changed, 196 insertions(+), 64 deletions(-) create mode 100644 iotdb-client/client-py/tests/unit/test_session_query_ids.py diff --git a/iotdb-client/client-py/iotdb/Session.py b/iotdb-client/client-py/iotdb/Session.py index 5a4f2ff23e707..b9bccd0509786 100644 --- a/iotdb-client/client-py/iotdb/Session.py +++ b/iotdb-client/client-py/iotdb/Session.py @@ -1627,22 +1627,22 @@ def execute_query_statement(self, sql, timeout=0): rpc_utils.verify_success(resp.status) return SessionDataSet( - sql, - resp.columns, - resp.dataTypeList, - resp.columnNameIndexMap, - resp.queryId, - self.__session_id, - self.__client, - self.__statement_id, - resp.queryResult, - resp.ignoreTimeStamp, - timeout, - resp.moreData, - self.__fetch_size, - self.__zone_id, - self.__time_precision, - resp.columnIndex2TsBlockColumnIndexList, + sql=sql, + column_name_list=resp.columns, + column_type_list=resp.dataTypeList, + column_name_index=resp.columnNameIndexMap, + query_id=resp.queryId, + statement_id=self.__statement_id, + client=self.__client, + session_id=self.__session_id, + query_result=resp.queryResult, + ignore_timestamp=resp.ignoreTimeStamp, + time_out=timeout, + more_data=resp.moreData, + fetch_size=self.__fetch_size, + zone_id=self.__zone_id, + time_precision=self.__time_precision, + column_index_2_tsblock_column_index_list=resp.columnIndex2TsBlockColumnIndexList, ) def execute_non_query_statement(self, sql): @@ -1698,22 +1698,22 @@ def execute_statement(self, sql: str, timeout=0): rpc_utils.verify_success(resp.status) if resp.columns: return SessionDataSet( - sql, - resp.columns, - resp.dataTypeList, - resp.columnNameIndexMap, - resp.queryId, - self.__session_id, - self.__client, - self.__statement_id, - resp.queryResult, - resp.ignoreTimeStamp, - timeout, - resp.moreData, - self.__fetch_size, - self.__zone_id, - self.__time_precision, - resp.columnIndex2TsBlockColumnIndexList, + sql=sql, + column_name_list=resp.columns, + column_type_list=resp.dataTypeList, + column_name_index=resp.columnNameIndexMap, + query_id=resp.queryId, + statement_id=self.__statement_id, + client=self.__client, + session_id=self.__session_id, + query_result=resp.queryResult, + ignore_timestamp=resp.ignoreTimeStamp, + time_out=timeout, + more_data=resp.moreData, + fetch_size=self.__fetch_size, + zone_id=self.__zone_id, + time_precision=self.__time_precision, + column_index_2_tsblock_column_index_list=resp.columnIndex2TsBlockColumnIndexList, ) else: return None @@ -1861,22 +1861,22 @@ def execute_raw_data_query( raise IoTDBConnectionException(self.connection_error_msg()) from None rpc_utils.verify_success(resp.status) return SessionDataSet( - "", - resp.columns, - resp.dataTypeList, - resp.columnNameIndexMap, - resp.queryId, - self.__session_id, - self.__client, - self.__statement_id, - resp.queryResult, - resp.ignoreTimeStamp, - 0, - resp.moreData, - self.__fetch_size, - self.__zone_id, - self.__time_precision, - resp.columnIndex2TsBlockColumnIndexList, + sql="", + column_name_list=resp.columns, + column_type_list=resp.dataTypeList, + column_name_index=resp.columnNameIndexMap, + query_id=resp.queryId, + statement_id=self.__statement_id, + client=self.__client, + session_id=self.__session_id, + query_result=resp.queryResult, + ignore_timestamp=resp.ignoreTimeStamp, + time_out=0, + more_data=resp.moreData, + fetch_size=self.__fetch_size, + zone_id=self.__zone_id, + time_precision=self.__time_precision, + column_index_2_tsblock_column_index_list=resp.columnIndex2TsBlockColumnIndexList, ) def execute_last_data_query(self, paths: list, last_time: int) -> SessionDataSet: @@ -1908,22 +1908,22 @@ def execute_last_data_query(self, paths: list, last_time: int) -> SessionDataSet raise IoTDBConnectionException(self.connection_error_msg()) from None rpc_utils.verify_success(resp.status) return SessionDataSet( - "", - resp.columns, - resp.dataTypeList, - resp.columnNameIndexMap, - resp.queryId, - self.__statement_id, - self.__client, - self.__session_id, - resp.queryResult, - resp.ignoreTimeStamp, - 0, - resp.moreData, - self.__fetch_size, - self.__zone_id, - self.__time_precision, - resp.columnIndex2TsBlockColumnIndexList, + sql="", + column_name_list=resp.columns, + column_type_list=resp.dataTypeList, + column_name_index=resp.columnNameIndexMap, + query_id=resp.queryId, + statement_id=self.__statement_id, + client=self.__client, + session_id=self.__session_id, + query_result=resp.queryResult, + ignore_timestamp=resp.ignoreTimeStamp, + time_out=0, + more_data=resp.moreData, + fetch_size=self.__fetch_size, + zone_id=self.__zone_id, + time_precision=self.__time_precision, + column_index_2_tsblock_column_index_list=resp.columnIndex2TsBlockColumnIndexList, ) def insert_string_records_of_one_device( diff --git a/iotdb-client/client-py/tests/unit/test_session_query_ids.py b/iotdb-client/client-py/tests/unit/test_session_query_ids.py new file mode 100644 index 0000000000000..17ca6508ed91e --- /dev/null +++ b/iotdb-client/client-py/tests/unit/test_session_query_ids.py @@ -0,0 +1,132 @@ +# 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 struct +from unittest.mock import Mock + +import pytest +from thrift.transport import TTransport + +from iotdb.Session import Session, SessionConnection +from iotdb.thrift.common.ttypes import TSStatus +from iotdb.thrift.rpc.IClientRPCService import Client +from iotdb.thrift.rpc.ttypes import TSExecuteStatementResp, TSFetchResultsResp + + +def int32_block(timestamp, value): + # One INT32 value column, one row, plain INT64/INT32 encodings, no nulls. + return struct.pack(">iBiBB?q?i", 1, 1, 1, 2, 1, False, timestamp, False, value) + + +@pytest.mark.parametrize( + "method,args,rpc_method", + [ + ( + "execute_query_statement", + ("select s from root.test.d",), + "executeQueryStatementV2", + ), + ( + "execute_statement", + ("select s from root.test.d",), + "executeStatementV2", + ), + ( + "execute_raw_data_query", + (["root.test.d.s"], 0, 10), + "executeRawDataQueryV2", + ), + ( + "execute_last_data_query", + (["root.test.d.s"], 0), + "executeLastDataQueryV2", + ), + ], +) +@pytest.mark.parametrize("reconnect", [False, True]) +@pytest.mark.parametrize("early_close", [False, True]) +def test_query_fetch_and_close_keep_connection_ids( + monkeypatch, method, args, rpc_method, reconnect, early_close +): + # The two generators can advance independently; equal IDs hide a swap. + initial_ids = (101, 202) + active_ids = (404, 505) if reconnect else initial_ids + query_id = 303 + client = Mock(spec=Client) + + def execute(request): + assert (request.sessionId, request.statementId) == active_ids + return TSExecuteStatementResp( + status=TSStatus(code=200), + queryId=query_id, + columns=["root.test.d.s"], + dataTypeList=["INT32"], + columnNameIndexMap={"root.test.d.s": 0}, + queryResult=[int32_block(1, 42)], + ignoreTimeStamp=False, + moreData=True, + ) + + def fetch(request): + assert (request.sessionId, request.statementId) == active_ids + assert request.queryId == query_id + return TSFetchResultsResp( + status=TSStatus(code=200), + hasResultSet=True, + queryResult=[int32_block(2, 43)], + moreData=False, + ) + + def close(request): + assert (request.sessionId, request.statementId) == active_ids + assert request.queryId == query_id + return TSStatus(code=200) + + getattr(client, rpc_method).side_effect = execute + client.fetchResultsV2.side_effect = fetch + client.closeOperation.side_effect = close + connections = [] + if reconnect: + failed_client = Mock(spec=Client) + + def fail(request): + assert (request.sessionId, request.statementId) == initial_ids + raise TTransport.TTransportException(message="Connection lost") + + getattr(failed_client, rpc_method).side_effect = fail + connections.append(SessionConnection(failed_client, Mock(), *initial_ids)) + connections.append(SessionConnection(client, Mock(), *active_ids)) + session = Session.init_from_node_urls(["127.0.0.1:6667"]) + init_connection = Mock(side_effect=connections) + monkeypatch.setattr(session, "init_connection", init_connection) + session.open() + + with getattr(session, method)(*args) as data_set: + row = data_set.next() + assert row.get_timestamp() == 1 + assert row.get_fields()[0].get_int_value() == 42 + if not early_close: + row = data_set.next() + assert row.get_timestamp() == 2 + assert row.get_fields()[0].get_int_value() == 43 + assert data_set.next() is None + + assert init_connection.call_count == (2 if reconnect else 1) + getattr(client, rpc_method).assert_called_once() + assert client.fetchResultsV2.call_count == (0 if early_close else 1) + client.closeOperation.assert_called_once()