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()