diff --git a/benchmarks/bench_was_applied.py b/benchmarks/bench_was_applied.py new file mode 100644 index 0000000000..467dd5d296 --- /dev/null +++ b/benchmarks/bench_was_applied.py @@ -0,0 +1,86 @@ +# Copyright DataStax, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Micro-benchmark: was_applied fast path for known LWT statements. + +Measures the speedup from skipping regex batch detection when the +query already knows it's an LWT statement (is_lwt() returns True). + +This benchmarks the real `cassandra.cluster.ResultSet.was_applied` property +(not a simplified stand-in for it): a minimal `ResultSet` is constructed with +a mocked `response_future` -- was_applied only reads +`response_future.row_factory`/`response_future.query` -- and real +`cassandra.query` statement objects, so both the fast-path `is_lwt()` check +and the slow-path `ResultSet.batch_regex` match run exactly as they do in +production. + +Run: + python benchmarks/bench_was_applied.py +""" +import timeit +from unittest.mock import Mock + +from cassandra.cluster import ResultSet +from cassandra.query import named_tuple_factory, SimpleStatement, PreparedStatement, BoundStatement + + +def _make_result_set(query, row): + """Build a minimal ResultSet with a mocked response_future, mirroring what + Session.execute()/ResponseFuture.result() construct in production.""" + response_future = Mock(row_factory=named_tuple_factory, query=query, + _col_names=None, _col_types=None) + return ResultSet(response_future, [row]) + + +def bench_was_applied(): + """Benchmark ResultSet.was_applied: fast path vs slow path.""" + # Fast path: a BoundStatement bound from a PreparedStatement whose LWT + # status was already resolved from the server's PREPARE response, so + # was_applied can skip batch/regex detection entirely. + prepared = PreparedStatement( + column_metadata=None, query_id=b'\x00', routing_key_indexes=None, + query="UPDATE t SET v=1 WHERE k=1 IF v=0", keyspace=None, + protocol_version=4, result_metadata=None, result_metadata_id=None, + is_lwt=True) + lwt_query = BoundStatement(prepared) + fast_rs = _make_result_set(lwt_query, (True,)) + + # Slow path: a plain SimpleStatement with unknown LWT status, so + # was_applied must match the query string against the real + # ResultSet.batch_regex to rule out a BEGIN BATCH. + non_lwt_query = SimpleStatement("INSERT INTO t (k, v) VALUES (1, 2) IF NOT EXISTS") + slow_rs = _make_result_set(non_lwt_query, (True,)) + + def fast_path(): + _ = fast_rs.was_applied + + def slow_path(): + _ = slow_rs.was_applied + + n = 500_000 + t_fast = timeit.timeit(fast_path, number=n) + t_slow = timeit.timeit(slow_path, number=n) + + print(f"Fast path (known LWT, {n} iters): {t_fast:.3f}s ({t_fast / n * 1e6:.2f} us/call)") + print(f"Slow path (regex check, {n} iters): {t_slow:.3f}s ({t_slow / n * 1e6:.2f} us/call)") + print(f"Speedup: {t_slow / t_fast:.1f}x") + + +def main(): + bench_was_applied() + + +if __name__ == '__main__': + main() diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 88c8d2707a..e736a32927 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -5909,13 +5909,29 @@ def was_applied(self): if self.response_future.row_factory not in (named_tuple_factory, dict_factory, tuple_factory): raise RuntimeError("Cannot determine LWT result with row factory %s" % (self.response_future.row_factory,)) - is_batch_statement = isinstance(self.response_future.query, BatchStatement) \ - or (isinstance(self.response_future.query, SimpleStatement) and self.batch_regex.match(self.response_future.query.query_string)) - if is_batch_statement and (not self.column_names or self.column_names[0] != "[applied]"): - raise RuntimeError("No LWT were present in the BatchStatement") + query = self.response_future.query + + # Fast path: BoundStatement/PreparedStatement with known LWT status + # from the server PREPARE response avoids batch detection entirely. + # `isinstance` is checked first so it always short-circuits before + # is_lwt() is looked up: BatchStatement is never eligible for the fast + # path anyway, and the getattr/callable guard means a query object + # that doesn't implement is_lwt() (e.g. a None query, or some other + # non-Statement/duck-typed object) falls through to the slow path + # below instead of raising AttributeError. + is_lwt = getattr(query, 'is_lwt', None) + if not isinstance(query, BatchStatement) and callable(is_lwt) and is_lwt(): + # Known single LWT statement - skip batch detection + if len(self.current_rows) != 1: + raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows))) + else: + is_batch_statement = isinstance(query, BatchStatement) \ + or (isinstance(query, SimpleStatement) and self.batch_regex.match(query.query_string)) + if is_batch_statement and (not self.column_names or self.column_names[0] != "[applied]"): + raise RuntimeError("No LWT were present in the BatchStatement") - if not is_batch_statement and len(self.current_rows) != 1: - raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows))) + if not is_batch_statement and len(self.current_rows) != 1: + raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows))) row = self.current_rows[0] if isinstance(row, tuple): diff --git a/tests/unit/test_resultset.py b/tests/unit/test_resultset.py index 80e9c21ff9..5017fb998e 100644 --- a/tests/unit/test_resultset.py +++ b/tests/unit/test_resultset.py @@ -16,7 +16,7 @@ from unittest.mock import Mock, PropertyMock, patch from cassandra.cluster import ResultSet -from cassandra.query import named_tuple_factory, dict_factory, tuple_factory +from cassandra.query import named_tuple_factory, dict_factory, tuple_factory, SimpleStatement, BatchStatement from tests.util import assertListEqual import pytest @@ -175,11 +175,18 @@ def test_bool(self): assert ResultSet(Mock(has_more_pages=False), [1]) def test_was_applied(self): + # Create a non-LWT query so these assertions exercise the slow (regex) path. + # Without this, Mock().query.is_lwt() returns a truthy Mock, accidentally + # routing all checks through the fast path. + non_lwt_query = Mock(spec=SimpleStatement) + non_lwt_query.is_lwt.return_value = False + non_lwt_query.query_string = "INSERT INTO t (k) VALUES (1)" + # unknown row factory raises with pytest.raises(RuntimeError): - ResultSet(Mock(), []).was_applied + ResultSet(Mock(query=non_lwt_query), []).was_applied - response_future = Mock(row_factory=named_tuple_factory) + response_future = Mock(row_factory=named_tuple_factory, query=non_lwt_query) # no row with pytest.raises(RuntimeError): @@ -192,14 +199,94 @@ def test_was_applied(self): # various internal row factories for row_factory in (named_tuple_factory, tuple_factory): for applied in (True, False): - rs = ResultSet(Mock(row_factory=row_factory), [(applied,)]) + rs = ResultSet(Mock(row_factory=row_factory, query=non_lwt_query), [(applied,)]) assert rs.was_applied == applied row_factory = dict_factory for applied in (True, False): - rs = ResultSet(Mock(row_factory=row_factory), [{'[applied]': applied}]) + rs = ResultSet(Mock(row_factory=row_factory, query=non_lwt_query), [{'[applied]': applied}]) assert rs.was_applied == applied + + def test_was_applied_lwt_fast_path(self): + """Test that was_applied uses fast path for known LWT statements.""" + # BoundStatement-like query with is_lwt() = True (fast path) + lwt_query = Mock() + lwt_query.is_lwt.return_value = True + for row_factory in (named_tuple_factory, tuple_factory): + for applied in (True, False): + rf = Mock(row_factory=row_factory, query=lwt_query) + rs = ResultSet(rf, [(applied,)]) + assert rs.was_applied == applied + + for applied in (True, False): + rf = Mock(row_factory=dict_factory, query=lwt_query) + rs = ResultSet(rf, [{'[applied]': applied}]) + assert rs.was_applied == applied + + # Fast path with too many rows should raise + rf = Mock(row_factory=named_tuple_factory, query=lwt_query) + with pytest.raises(RuntimeError, match="exactly one row"): + ResultSet(rf, [tuple(), tuple()]).was_applied + + def test_was_applied_non_lwt_fallback(self): + """Test that was_applied falls back to slow path for non-LWT statements.""" + # SimpleStatement-like query with is_lwt() = False (slow path, non-batch) + non_lwt_query = Mock(spec=SimpleStatement) + non_lwt_query.is_lwt.return_value = False + non_lwt_query.query_string = "INSERT INTO t (k) VALUES (1)" + + for applied in (True, False): + rf = Mock(row_factory=tuple_factory, query=non_lwt_query) + rs = ResultSet(rf, [(applied,)]) + assert rs.was_applied == applied + + def test_was_applied_batch_statement(self): + """Test that was_applied handles BatchStatement correctly (slow path).""" + # BatchStatement with LWT should check column_names + batch_query = Mock(spec=BatchStatement) + batch_query.is_lwt.return_value = True + + # Batch with [applied] column -- pass _col_names so ResultSet.__init__ + # sets column_names correctly (instead of post-construction override). + rf = Mock(row_factory=tuple_factory, query=batch_query, + _col_names=['[applied]'], _col_types=None) + rs = ResultSet(rf, [(True,)]) + assert rs.was_applied + + # Batch without [applied] column raises + rf = Mock(row_factory=tuple_factory, query=batch_query, + _col_names=['other'], _col_types=None) + rs = ResultSet(rf, [(True,)]) + with pytest.raises(RuntimeError, match="No LWT were present"): + rs.was_applied + + def test_was_applied_query_without_is_lwt(self): + """Regression test: was_applied must not raise AttributeError when + response_future.query doesn't implement is_lwt(). + + The fast path used to call query.is_lwt() before confirming the query + even has that method (only guarding against BatchStatement). Any other + object without is_lwt() -- including a bare `None`, which is a real, + reachable value (ResponseFuture.query defaults to None, and + Session.prepare()/prepare_on_all_hosts construct ResponseFuture(..., + query=None, ...) explicitly) -- would blow up instead of falling + through to the same slow-path handling as any other non-LWT query. + """ + rf = Mock(row_factory=named_tuple_factory, query=None, _col_names=None, _col_types=None) + rs = ResultSet(rf, [(True,)]) + assert rs.was_applied is True + + # A generic duck-typed statement predating the is_lwt() fast path, + # with no is_lwt() method at all. + class _StatementWithoutIsLwt: + query_string = "SELECT 1" + + rf2 = Mock(row_factory=named_tuple_factory, query=_StatementWithoutIsLwt(), + _col_names=None, _col_types=None) + rs2 = ResultSet(rf2, [(False,)]) + assert rs2.was_applied is False + def test_one(self): # no pages first, second = Mock(), Mock()