diff --git a/tests/table/table_test.py b/tests/table/table_test.py index be527ee1..e0dc03fc 100644 --- a/tests/table/table_test.py +++ b/tests/table/table_test.py @@ -2,6 +2,20 @@ import ydb +class TestSession: + def test_keep_in_cache_dataquery(self, driver_sync): + # Keep prepared query for keep_in_cache only, not store query_id if client query cache disabled + tc_settings = ydb.TableClientSettings().with_client_query_cache(enabled=False) + table_client = ydb.TableClient(driver_sync, tc_settings) + data_query = ydb.DataQuery("select 1", {}) + session = ydb.retry_operation_sync(lambda: table_client.session().create()) + session.transaction().execute(data_query, commit_tx=True) + assert session.has_prepared("select 1") + assert session._state.lookup("select 1")[1] is None + assert session.has_prepared(data_query) + assert session._state.lookup(data_query)[1] is None + + class TestSessionPool: def test_checkout_from_stopped_pool(self, driver_sync): pool = ydb.SessionPool(driver_sync, 1) diff --git a/ydb/_tx_ctx_impl.py b/ydb/_tx_ctx_impl.py index 969a4dc4..3c7eec39 100644 --- a/ydb/_tx_ctx_impl.py +++ b/ydb/_tx_ctx_impl.py @@ -110,6 +110,8 @@ def execute_request_factory(session_state, tx_state, query, parameters, commit_t data_query, query_id = session_state.lookup(query) parameters_types = {} + is_data_query = False + if query_id is not None: query_pb = _apis.ydb_table.Query(id=query_id) parameters_types = data_query.parameters_types @@ -118,9 +120,11 @@ def execute_request_factory(session_state, tx_state, query, parameters, commit_t # client cache disabled for send query text every time yql_text = data_query.yql_text parameters_types = data_query.parameters_types + is_data_query = True elif isinstance(query, types.DataQuery): yql_text = query.yql_text parameters_types = query.parameters_types + is_data_query = True else: yql_text = query query_pb = _apis.ydb_table.Query(yql_text=yql_text) @@ -133,6 +137,8 @@ def execute_request_factory(session_state, tx_state, query, parameters, commit_t keep_in_cache = settings.keep_in_cache elif parameters: keep_in_cache = True + elif is_data_query: + keep_in_cache = True else: keep_in_cache = False @@ -159,5 +165,7 @@ def wrap_result_and_tx_id(rpc_state, response_pb, session_state, tx_state, query issues._process_response(response_pb.operation) message = _apis.ydb_table.ExecuteQueryResult() response_pb.operation.result.Unpack(message) + if message.query_meta.id and isinstance(query, types.DataQuery): + session_state.keep(query, message.query_meta.id) tx_state.tx_id = None if not message.tx_meta.id else message.tx_meta.id return convert.ResultSets(message.result_sets, session_state.table_client_settings)