From d9db800e6c5b38b2e2db1dbee5a9f19cc94b6119 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 4 Mar 2026 11:08:03 +0200 Subject: [PATCH] This PR fixes UNPREPARED handling in the Python driver execution path for prepared statements. ### Problem When the server returned UNPREPARED (PreparedQueryNotFound) during an EXECUTE, the driver could fail to recover in some valid cases (cache miss, race, coordinator restart, id mismatch edge cases), even when the in-flight request still had enough PreparedStatement context to recover. ### Expected behavior If the driver still has the original PreparedStatement context, it should: 1. issue PREPARE, 2. then retry the original EXECUTE, regardless of why the coordinator no longer recognizes the prepared id. ### What this PR changes In ResponseFuture PreparedQueryNotFound handling: 1. Cache the in-flight self.prepared_statement first (when present). 2. Attempt lookup by the returned UNPREPARED id from cluster._prepared_statements. 3. Reprepare using the resolved statement. 4. If returned-id lookup fails: - fallback to in-flight context if available, - otherwise fail as unknown prepared statement. This preserves correctness while improving recovery in real-world coordinator invalidation/race/restart scenarios. ### Tests Added/updated unit coverage in tests/unit/test_response_future.py: - validates reprepare path for PreparedQueryNotFound, - validates fallback to in-flight context, - validates preference for statement resolved by returned UNPREPARED id when available. Fixes: https://github.com/scylladb/scylladb/issues/27657 --- cassandra/cluster.py | 67 +++++++----- tests/unit/test_response_future.py | 159 ++++++++++++++++++++++++++++- 2 files changed, 198 insertions(+), 28 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 1962f95a65..53bc7ce5b9 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -8153,8 +8153,9 @@ def start_fetching_next_page(self): self._start_timer() self.send_request() - def _reprepare(self, prepare_message, host, connection, pool): - cb = partial(self.session.submit, self._execute_after_prepare, host, connection, pool) + def _reprepare(self, prepare_message, host, connection, pool, prepared_statement=_NOT_SET): + cb = partial(self.session.submit, self._execute_after_prepare, host, connection, pool, + prepared_statement=prepared_statement) if pool is None and connection is not None and connection.is_control_connection: request_id = self._query_control_connection(prepare_message, cb=cb, connection=connection, host=host) @@ -8320,28 +8321,31 @@ def _set_result(self, host, connection, pool, response, page_generation=None): self.query, cl, error=response, retry_num=self._query_retries) elif isinstance(response, PreparedQueryNotFound): + query_id = response.info + if self.prepared_statement: - query_id = self.prepared_statement.query_id - assert query_id == response.info, \ - "Got different query ID in server response (%s) than we " \ - "had before (%s)" % (response.info, query_id) - else: - query_id = response.info + # Cache local in-flight context first so lookup by either id can succeed. + self.session.cluster.add_prepared(self.prepared_statement.query_id, self.prepared_statement) + # Resolve the statement to reprepare into a local: self.prepared_statement + # is the execution context shared by every attempt of this future + # (retries, speculative executions), so it must not be replaced by + # the statement one attempt found missing. try: - prepared_statement = self.session.cluster._prepared_statements[query_id] + statement_to_prepare = self.session.cluster._prepared_statements[query_id] except KeyError: if not self.prepared_statement: - log.error("Tried to execute unknown prepared statement: id=%s", - query_id.encode('hex')) + log.error("Tried to execute unknown prepared statement: id=%s", hexlify(query_id)) self._set_final_exception(response) return - else: - prepared_statement = self.prepared_statement - self.session.cluster._prepared_statements[query_id] = prepared_statement + log.warning( + "UNPREPARED for query id %s while executing statement id %s. " + "Could not resolve returned id in cache, proceeding with in-flight context.", + hexlify(query_id), hexlify(self.prepared_statement.query_id)) + statement_to_prepare = self.prepared_statement current_keyspace = self._connection.keyspace - prepared_keyspace = prepared_statement.keyspace + prepared_keyspace = statement_to_prepare.keyspace if not ProtocolVersion.uses_keyspace_flag(self.session.cluster.protocol_version) \ and prepared_keyspace and current_keyspace != prepared_keyspace: self._set_final_exception( @@ -8351,15 +8355,18 @@ def _set_result(self, host, connection, pool, response, page_generation=None): (current_keyspace, prepared_keyspace))) return - log.debug("Re-preparing unrecognized prepared statement against host %s: %s", - host, prepared_statement.query_string) - prepared_keyspace = prepared_statement.keyspace \ + log.debug( + "Re-preparing unrecognized prepared statement against host %s: %s", + host, statement_to_prepare.query_string + ) + prepared_keyspace = statement_to_prepare.keyspace \ if ProtocolVersion.uses_keyspace_flag(self.session.cluster.protocol_version) else None - prepare_message = PrepareMessage(query=prepared_statement.query_string, + prepare_message = PrepareMessage(query=statement_to_prepare.query_string, keyspace=prepared_keyspace) # since this might block, run on the executor to avoid hanging # the event loop thread - self.session.submit(self._reprepare, prepare_message, host, connection, pool) + self.session.submit(self._reprepare, prepare_message, host, connection, pool, + prepared_statement=statement_to_prepare) return else: if hasattr(response, 'to_exception'): @@ -8410,11 +8417,18 @@ def _set_keyspace_completed(self, errors): self._set_final_exception(ConnectionException( "Failed to set keyspace on all hosts: %s" % (errors,))) - def _execute_after_prepare(self, host, connection, pool, response): + def _execute_after_prepare(self, host, connection, pool, response, prepared_statement=_NOT_SET): """ Handle the response to our attempt to prepare a statement. If it succeeded, run the original query again against the same host. + + ``prepared_statement`` is the statement this PREPARE was sent for; the + response is validated against it and its result metadata is updated. + It defaults to the future's own prepared statement. """ + if prepared_statement is _NOT_SET: + prepared_statement = self.prepared_statement + if pool: pool.return_connection(connection) @@ -8423,16 +8437,17 @@ def _execute_after_prepare(self, host, connection, pool, response): if isinstance(response, ResultMessage): if response.kind == RESULT_KIND_PREPARED: - if self.prepared_statement: - if self.prepared_statement.query_id != response.query_id: + if prepared_statement: + if prepared_statement.query_id != response.query_id: self._set_final_exception(DriverException( "ID mismatch while trying to reprepare (expected {expected}, got {got}). " "This prepared statement won't work anymore. " "This usually happens when you run a 'USE...' " "query after the statement was prepared.".format( - expected=hexlify(self.prepared_statement.query_id), got=hexlify(response.query_id) + expected=hexlify(prepared_statement.query_id), got=hexlify(response.query_id) ) )) + return # Update the metadata/id pair atomically from exactly what this # reprepare response carries. Falling back to the previously # cached id when this response has none would risk pairing it @@ -8442,9 +8457,9 @@ def _execute_after_prepare(self, host, connection, pool, response): # id-aware execute could send without the server detecting the # mismatch. Dropping it instead triggers the same self-healing # b'' sentinel path a never-prepared id would. - self.prepared_statement.update_result_metadata( + prepared_statement.update_result_metadata( response.column_metadata, response.result_metadata_id) - + # use self._query to re-use the same host and # at the same time properly borrow the connection if pool is None and connection is not None and connection.is_control_connection: diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index 1e7a568168..279d9a9583 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -20,7 +20,7 @@ from threading import Barrier, Event, RLock, Thread from unittest.mock import Mock, MagicMock, ANY, patch -from cassandra import ConsistencyLevel, InvalidRequest, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut +from cassandra import ConsistencyLevel, DriverException, InvalidRequest, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut from cassandra.cluster import (Session, ResponseFuture, NoHostAvailable, ProtocolVersion, ControlConnection, ControlConnectionQueryFallback, _NOT_SET) from cassandra.connection import (Connection, ConnectionBusy, ConnectionException, @@ -2307,7 +2307,8 @@ def test_control_connection_fallback_reprepare_sets_keyspace_first(self): prepared_response = Mock(spec=ResultMessage, kind=RESULT_KIND_PREPARED) connection.send_msg.call_args_list[1][1]['cb'](prepared_response) session.submit.assert_called_once_with( - rf._execute_after_prepare, host, connection, None, prepared_response) + rf._execute_after_prepare, host, connection, None, prepared_response, + prepared_statement=_NOT_SET) def test_control_connection_fallback_not_used_when_pool_can_serve(self): session = self.make_basic_session() @@ -2715,6 +2716,160 @@ def test_prepared_query_not_found_bad_keyspace(self): with pytest.raises(ValueError): rf.result() + def test_prepared_query_not_found_uses_local_prepared_context(self): + session = self.make_session() + pool = session._pools.get.return_value + connection = Mock(spec=Connection) + pool.borrow_connection.return_value = (connection, 1) + + rf = self.make_response_future(session) + rf.send_request() + + session.cluster.protocol_version = ProtocolVersion.V4 + session.cluster._prepared_statements = {} + rf._connection.keyspace = "FooKeyspace" + + rf.prepared_statement = Mock() + rf.prepared_statement.query_id = b"known-query-id" + rf.prepared_statement.query_string = "SELECT * FROM foobar" + rf.prepared_statement.keyspace = "FooKeyspace" + + # Different query id in UNPREPARED response should not prevent reprepare when local context exists. + result = Mock(spec=PreparedQueryNotFound, info=b"other-query-id") + rf._set_result(None, None, None, result) + + assert session.submit.call_args + args, _ = session.submit.call_args + assert rf._reprepare == args[-5] + assert isinstance(args[-4], PrepareMessage) + assert args[-4].query == "SELECT * FROM foobar" + + def test_prepared_query_not_found_prefers_returned_id_from_cache(self): + session = self.make_session() + pool = session._pools.get.return_value + connection = Mock(spec=Connection) + pool.borrow_connection.return_value = (connection, 1) + + rf = self.make_response_future(session) + rf.send_request() + + session.cluster.protocol_version = ProtocolVersion.V4 + rf._connection.keyspace = "FooKeyspace" + + rf.prepared_statement = Mock() + rf.prepared_statement.query_id = b"local-id" + rf.prepared_statement.query_string = "SELECT * FROM local_ctx" + rf.prepared_statement.keyspace = "FooKeyspace" + + cached_stmt = Mock() + cached_stmt.query_id = b"returned-id" + cached_stmt.query_string = "SELECT * FROM returned_ctx" + cached_stmt.keyspace = "FooKeyspace" + session.cluster._prepared_statements = {cached_stmt.query_id: cached_stmt} + + result = Mock(spec=PreparedQueryNotFound, info=cached_stmt.query_id) + rf._set_result(None, None, None, result) + + assert session.submit.call_args + args, _ = session.submit.call_args + assert rf._reprepare == args[-5] + assert isinstance(args[-4], PrepareMessage) + assert args[-4].query == "SELECT * FROM returned_ctx" + + def _reprepare_and_capture_callbacks(self, rf, session, missing_ids): + """ + Feed one UNPREPARED response per id into ``rf`` (as separate attempts of + the same future would) and return the PREPARE callbacks, in order. + """ + session.submit.side_effect = lambda fn, *args, **kwargs: fn(*args, **kwargs) + rf._query = Mock(return_value=1) + for query_id in missing_ids: + rf._set_result('host', None, None, Mock(spec=PreparedQueryNotFound, info=query_id)) + prepare_calls = [c for c in rf._query.call_args_list if 'cb' in c[1]] + assert len(prepare_calls) == len(missing_ids) + rf._query.reset_mock() + return [c[1]['cb'] for c in prepare_calls] + + def _prepared_response(self, query_id, column_metadata): + return Mock(spec=ResultMessage, kind=RESULT_KIND_PREPARED, query_id=query_id, + column_metadata=column_metadata, result_metadata_id=query_id + b'-meta') + + def test_reprepare_of_two_statements_on_one_future_keeps_completions_apart(self): + """ + Two attempts of one future (e.g. speculative executions of a batch) can each + get UNPREPARED for a different statement before either PREPARE completes. + Each PREPARED response must be matched against the statement it was sent + for, not against whichever statement was reported missing last. + """ + session = self.make_session() + session.cluster.protocol_version = ProtocolVersion.V4 + stmt_a = self._make_prepared_statement([], None, query_id=b'id-a') + stmt_b = self._make_prepared_statement([], None, query_id=b'id-b') + session.cluster._prepared_statements = {b'id-a': stmt_a, b'id-b': stmt_b} + + rf = self.make_response_future(session) + rf.send_request() + rf._connection.keyspace = 'ks' + assert rf.prepared_statement is None + + cb_a, cb_b = self._reprepare_and_capture_callbacks(rf, session, [b'id-a', b'id-b']) + assert rf.prepared_statement is None + + meta_a = [('ks', 'tb', 'a', Mock())] + meta_b = [('ks', 'tb', 'b', Mock())] + cb_a(self._prepared_response(b'id-a', meta_a)) + cb_b(self._prepared_response(b'id-b', meta_b)) + + assert rf._final_exception is None + assert stmt_a.result_metadata is meta_a + assert stmt_b.result_metadata is meta_b + assert rf._query.call_count == 2 + + def test_reprepare_of_returned_id_keeps_execution_context(self): + """ + EXECUTE(A) answered with UNPREPARED(B) for a cached B reprepares B, but the + future keeps executing A: A's statement and metadata must stay untouched. + """ + session = self.make_session() + session.cluster.protocol_version = ProtocolVersion.V4 + meta_a = [('ks', 'tb', 'a', Mock())] + stmt_a = self._make_prepared_statement(meta_a, b'meta-a', query_id=b'id-a') + stmt_b = self._make_prepared_statement([], None, query_id=b'id-b') + session.cluster._prepared_statements = {b'id-b': stmt_b} + + rf = self.make_response_future(session) + rf.prepared_statement = stmt_a + rf.send_request() + rf._connection.keyspace = 'ks' + + cb_b, = self._reprepare_and_capture_callbacks(rf, session, [b'id-b']) + assert rf.prepared_statement is stmt_a + + meta_b = [('ks', 'tb', 'b', Mock())] + cb_b(self._prepared_response(b'id-b', meta_b)) + + assert rf._final_exception is None + assert rf.prepared_statement is stmt_a + assert stmt_a.result_metadata_and_id == (meta_a, b'meta-a') + assert stmt_b.result_metadata is meta_b + rf._query.assert_called_once_with('host') + + def test_reprepare_id_mismatch_does_not_retry_or_update_metadata(self): + session = self.make_session() + meta = [('ks', 'tb', 'old', Mock())] + stmt = self._make_prepared_statement(meta, b'old-meta', query_id=b'id-a') + rf = self.make_response_future(session) + rf._query = Mock(return_value=1) + + rf._execute_after_prepare('host', None, None, + self._prepared_response(b'id-other', []), + prepared_statement=stmt) + + with pytest.raises(DriverException, match="ID mismatch"): + rf.result() + assert stmt.result_metadata_and_id == (meta, b'old-meta') + rf._query.assert_not_called() + def test_repeat_orig_query_after_succesful_reprepare(self): query_id = b'abc123' # Just a random binary string so we don't hit id mismatch exception session = self.make_session()