Skip to content

This PR fixes UNPREPARED handling in the Python driver execution path… #725

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
calebxyz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

This PR fixes UNPREPARED handling in the Python driver execution path… #725

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 41 additions & 26 deletions cassandra/cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand All @@ -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'):
Expand Down Expand Up @@ -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)

Expand All @@ -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
Expand All @@ -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:
Expand Down
159 changes: 157 additions & 2 deletions tests/unit/test_response_future.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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.
Comment thread
calebxyz marked this conversation as resolved.
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()
Expand Down
Loading