|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | +# |
| 18 | + |
| 19 | +import struct |
| 20 | +from unittest.mock import Mock |
| 21 | + |
| 22 | +import pytest |
| 23 | +from thrift.transport import TTransport |
| 24 | + |
| 25 | +from iotdb.Session import Session, SessionConnection |
| 26 | +from iotdb.thrift.common.ttypes import TSStatus |
| 27 | +from iotdb.thrift.rpc.IClientRPCService import Client |
| 28 | +from iotdb.thrift.rpc.ttypes import TSExecuteStatementResp, TSFetchResultsResp |
| 29 | + |
| 30 | + |
| 31 | +def int32_block(timestamp, value): |
| 32 | + # One INT32 value column, one row, plain INT64/INT32 encodings, no nulls. |
| 33 | + return struct.pack(">iBiBB?q?i", 1, 1, 1, 2, 1, False, timestamp, False, value) |
| 34 | + |
| 35 | + |
| 36 | +@pytest.mark.parametrize( |
| 37 | + "method,args,rpc_method", |
| 38 | + [ |
| 39 | + ( |
| 40 | + "execute_query_statement", |
| 41 | + ("select s from root.test.d",), |
| 42 | + "executeQueryStatementV2", |
| 43 | + ), |
| 44 | + ( |
| 45 | + "execute_statement", |
| 46 | + ("select s from root.test.d",), |
| 47 | + "executeStatementV2", |
| 48 | + ), |
| 49 | + ( |
| 50 | + "execute_raw_data_query", |
| 51 | + (["root.test.d.s"], 0, 10), |
| 52 | + "executeRawDataQueryV2", |
| 53 | + ), |
| 54 | + ( |
| 55 | + "execute_last_data_query", |
| 56 | + (["root.test.d.s"], 0), |
| 57 | + "executeLastDataQueryV2", |
| 58 | + ), |
| 59 | + ], |
| 60 | +) |
| 61 | +@pytest.mark.parametrize("reconnect", [False, True]) |
| 62 | +@pytest.mark.parametrize("early_close", [False, True]) |
| 63 | +def test_query_fetch_and_close_keep_connection_ids( |
| 64 | + monkeypatch, method, args, rpc_method, reconnect, early_close |
| 65 | +): |
| 66 | + # The two generators can advance independently; equal IDs hide a swap. |
| 67 | + initial_ids = (101, 202) |
| 68 | + active_ids = (404, 505) if reconnect else initial_ids |
| 69 | + query_id = 303 |
| 70 | + client = Mock(spec=Client) |
| 71 | + |
| 72 | + def execute(request): |
| 73 | + assert (request.sessionId, request.statementId) == active_ids |
| 74 | + return TSExecuteStatementResp( |
| 75 | + status=TSStatus(code=200), |
| 76 | + queryId=query_id, |
| 77 | + columns=["root.test.d.s"], |
| 78 | + dataTypeList=["INT32"], |
| 79 | + columnNameIndexMap={"root.test.d.s": 0}, |
| 80 | + queryResult=[int32_block(1, 42)], |
| 81 | + ignoreTimeStamp=False, |
| 82 | + moreData=True, |
| 83 | + ) |
| 84 | + |
| 85 | + def fetch(request): |
| 86 | + assert (request.sessionId, request.statementId) == active_ids |
| 87 | + assert request.queryId == query_id |
| 88 | + return TSFetchResultsResp( |
| 89 | + status=TSStatus(code=200), |
| 90 | + hasResultSet=True, |
| 91 | + queryResult=[int32_block(2, 43)], |
| 92 | + moreData=False, |
| 93 | + ) |
| 94 | + |
| 95 | + def close(request): |
| 96 | + assert (request.sessionId, request.statementId) == active_ids |
| 97 | + assert request.queryId == query_id |
| 98 | + return TSStatus(code=200) |
| 99 | + |
| 100 | + getattr(client, rpc_method).side_effect = execute |
| 101 | + client.fetchResultsV2.side_effect = fetch |
| 102 | + client.closeOperation.side_effect = close |
| 103 | + connections = [] |
| 104 | + if reconnect: |
| 105 | + failed_client = Mock(spec=Client) |
| 106 | + |
| 107 | + def fail(request): |
| 108 | + assert (request.sessionId, request.statementId) == initial_ids |
| 109 | + raise TTransport.TTransportException(message="Connection lost") |
| 110 | + |
| 111 | + getattr(failed_client, rpc_method).side_effect = fail |
| 112 | + connections.append(SessionConnection(failed_client, Mock(), *initial_ids)) |
| 113 | + connections.append(SessionConnection(client, Mock(), *active_ids)) |
| 114 | + session = Session.init_from_node_urls(["127.0.0.1:6667"]) |
| 115 | + init_connection = Mock(side_effect=connections) |
| 116 | + monkeypatch.setattr(session, "init_connection", init_connection) |
| 117 | + session.open() |
| 118 | + |
| 119 | + with getattr(session, method)(*args) as data_set: |
| 120 | + row = data_set.next() |
| 121 | + assert row.get_timestamp() == 1 |
| 122 | + assert row.get_fields()[0].get_int_value() == 42 |
| 123 | + if not early_close: |
| 124 | + row = data_set.next() |
| 125 | + assert row.get_timestamp() == 2 |
| 126 | + assert row.get_fields()[0].get_int_value() == 43 |
| 127 | + assert data_set.next() is None |
| 128 | + |
| 129 | + assert init_connection.call_count == (2 if reconnect else 1) |
| 130 | + getattr(client, rpc_method).assert_called_once() |
| 131 | + assert client.fetchResultsV2.call_count == (0 if early_close else 1) |
| 132 | + client.closeOperation.assert_called_once() |
0 commit comments