Skip to content

Commit 65e2db1

Browse files
committed
perf: optimize was_applied fast path for known LWT statements
Add a fast path in ResultSet.was_applied that skips batch detection (isinstance checks + regex match) when the query has a known LWT status from the server PREPARE response. For BoundStatement queries where is_lwt() returns True, the batch_regex match on the query string is entirely avoided. This benefits the most common LWT use case: prepared INSERT/UPDATE IF statements executed via BoundStatement, where the driver already knows from the PREPARE response whether the statement is an LWT. The slow path (isinstance + regex) is preserved for: - BatchStatement queries (detected via isinstance) - SimpleStatement batch queries (detected via regex) - Any query where is_lwt() returns False Also adds explicit tests for the fast path, non-LWT fallback, and BatchStatement handling in was_applied. Part of: scylladb#751
1 parent 9c53d78 commit 65e2db1

2 files changed

Lines changed: 69 additions & 7 deletions

File tree

cassandra/cluster.py

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5336,13 +5336,22 @@ def was_applied(self):
53365336
if self.response_future.row_factory not in (named_tuple_factory, dict_factory, tuple_factory):
53375337
raise RuntimeError("Cannot determine LWT result with row factory %s" % (self.response_future.row_factory,))
53385338

5339-
is_batch_statement = isinstance(self.response_future.query, BatchStatement) \
5340-
or (isinstance(self.response_future.query, SimpleStatement) and self.batch_regex.match(self.response_future.query.query_string))
5341-
if is_batch_statement and (not self.column_names or self.column_names[0] != "[applied]"):
5342-
raise RuntimeError("No LWT were present in the BatchStatement")
5339+
query = self.response_future.query
5340+
5341+
# Fast path: BoundStatement/PreparedStatement with known LWT status
5342+
# from the server PREPARE response avoids batch detection entirely.
5343+
if query.is_lwt() and not isinstance(query, BatchStatement):
5344+
# Known single LWT statement - skip batch detection
5345+
if len(self.current_rows) != 1:
5346+
raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows)))
5347+
else:
5348+
is_batch_statement = isinstance(query, BatchStatement) \
5349+
or (isinstance(query, SimpleStatement) and self.batch_regex.match(query.query_string))
5350+
if is_batch_statement and (not self.column_names or self.column_names[0] != "[applied]"):
5351+
raise RuntimeError("No LWT were present in the BatchStatement")
53435352

5344-
if not is_batch_statement and len(self.current_rows) != 1:
5345-
raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows)))
5353+
if not is_batch_statement and len(self.current_rows) != 1:
5354+
raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows)))
53465355

53475356
row = self.current_rows[0]
53485357
if isinstance(row, tuple):

tests/unit/test_resultset.py

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from unittest.mock import Mock, PropertyMock, patch
1717

1818
from cassandra.cluster import ResultSet
19-
from cassandra.query import named_tuple_factory, dict_factory, tuple_factory
19+
from cassandra.query import named_tuple_factory, dict_factory, tuple_factory, SimpleStatement, BatchStatement
2020

2121
from tests.util import assertListEqual
2222
import pytest
@@ -200,6 +200,59 @@ def test_was_applied(self):
200200
rs = ResultSet(Mock(row_factory=row_factory), [{'[applied]': applied}])
201201
assert rs.was_applied == applied
202202

203+
204+
def test_was_applied_lwt_fast_path(self):
205+
"""Test that was_applied uses fast path for known LWT statements."""
206+
# BoundStatement-like query with is_lwt() = True (fast path)
207+
lwt_query = Mock()
208+
lwt_query.is_lwt.return_value = True
209+
for row_factory in (named_tuple_factory, tuple_factory):
210+
for applied in (True, False):
211+
rf = Mock(row_factory=row_factory, query=lwt_query)
212+
rs = ResultSet(rf, [(applied,)])
213+
assert rs.was_applied == applied
214+
215+
for applied in (True, False):
216+
rf = Mock(row_factory=dict_factory, query=lwt_query)
217+
rs = ResultSet(rf, [{'[applied]': applied}])
218+
assert rs.was_applied == applied
219+
220+
# Fast path with too many rows should raise
221+
rf = Mock(row_factory=named_tuple_factory, query=lwt_query)
222+
with pytest.raises(RuntimeError, match="exactly one row"):
223+
ResultSet(rf, [tuple(), tuple()]).was_applied
224+
225+
def test_was_applied_non_lwt_fallback(self):
226+
"""Test that was_applied falls back to slow path for non-LWT statements."""
227+
# SimpleStatement-like query with is_lwt() = False (slow path, non-batch)
228+
non_lwt_query = Mock(spec=SimpleStatement)
229+
non_lwt_query.is_lwt.return_value = False
230+
non_lwt_query.query_string = "INSERT INTO t (k) VALUES (1)"
231+
232+
for applied in (True, False):
233+
rf = Mock(row_factory=tuple_factory, query=non_lwt_query)
234+
rs = ResultSet(rf, [(applied,)])
235+
assert rs.was_applied == applied
236+
237+
def test_was_applied_batch_statement(self):
238+
"""Test that was_applied handles BatchStatement correctly (slow path)."""
239+
# BatchStatement with LWT should check column_names
240+
batch_query = Mock(spec=BatchStatement)
241+
batch_query.is_lwt.return_value = True
242+
243+
# Batch with [applied] column
244+
rf = Mock(row_factory=tuple_factory, query=batch_query)
245+
rs = ResultSet(rf, [(True,)])
246+
rs.column_names = ['[applied]']
247+
assert rs.was_applied == True
248+
249+
# Batch without [applied] column raises
250+
rf = Mock(row_factory=tuple_factory, query=batch_query)
251+
rs = ResultSet(rf, [(True,)])
252+
rs.column_names = ['other']
253+
with pytest.raises(RuntimeError, match="No LWT were present"):
254+
rs.was_applied
255+
203256
def test_one(self):
204257
# no pages
205258
first, second = Mock(), Mock()

0 commit comments

Comments
 (0)