From 3481055e75ce21fc99d864ab66b5c12c769b5225 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 27 Dec 2025 10:44:02 +0000 Subject: [PATCH 1/2] Initial plan From 507b2527263d23c8ffdca9d223027a3b794b2f2c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 29 Dec 2025 17:22:10 +0000 Subject: [PATCH 2/2] Optimize TokenAwarePolicy to call child.make_query_plan only once Co-authored-by: mykaul <4655593+mykaul@users.noreply.github.com> --- cassandra/policies.py | 10 +++--- tests/unit/test_policies.py | 67 ++++++++++++++++++++++++++++++++++--- 2 files changed, 68 insertions(+), 9 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index bcfd797706..72a9ba12d6 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -505,8 +505,12 @@ def make_query_plan(self, working_keyspace=None, query=None): keyspace = query.keyspace if query and query.keyspace else working_keyspace child = self._child_policy + + # Call child.make_query_plan only once and convert to list for reuse + child_plan = list(child.make_query_plan(keyspace, query)) + if query is None or query.routing_key is None or keyspace is None: - for host in child.make_query_plan(keyspace, query): + for host in child_plan: yield host return @@ -517,8 +521,6 @@ def make_query_plan(self, working_keyspace=None, query=None): if tablet is not None: replicas_mapped = set(map(lambda r: r[0], tablet.replicas)) - child_plan = child.make_query_plan(keyspace, query) - replicas = [host for host in child_plan if host.host_id in replicas_mapped] else: replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) @@ -535,7 +537,7 @@ def yield_in_order(hosts): # yield replicas: local_rack, local, remote yield from yield_in_order(replicas) # yield rest of the cluster: local_rack, local, remote - yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]) + yield from yield_in_order([host for host in child_plan if host not in replicas]) def on_up(self, *args, **kwargs): return self._child_policy.on_up(*args, **kwargs) diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index e15705c8f7..25d3689b53 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -945,13 +945,70 @@ def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key): else: assert set(replicas) == set(qplan[:2]) assert hosts[:2] == qplan[2:] - if is_tablets: - child_policy.make_query_plan.assert_called_with(keyspace, query) - assert child_policy.make_query_plan.call_count == 2 - else: - child_policy.make_query_plan.assert_called_once_with(keyspace, query) + # After optimization, child.make_query_plan should be called once for both tablets and vnodes + child_policy.make_query_plan.assert_called_once_with(keyspace, query) assert patched_shuffle.call_count == 1 + def test_child_make_query_plan_called_once(self): + """ + Test to validate that child.make_query_plan is called only once + in all scenarios (with/without tablets, with/without routing key) + + @test_category policy + """ + # Test with vnodes (no tablets) + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)] + for host in hosts: + host.set_up() + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.table_has_tablets.return_value = False + replicas = hosts[2:] + cluster.metadata.get_replicas.return_value = replicas + + child_policy = Mock() + child_policy.make_query_plan.return_value = hosts + child_policy.distance.return_value = HostDistance.LOCAL + + policy = TokenAwarePolicy(child_policy) + policy.populate(cluster, hosts) + + # Test case 1: With routing key and keyspace (should call once) + child_policy.reset_mock() + keyspace = 'keyspace' + routing_key = 'routing_key' + query = Statement(routing_key=routing_key, keyspace=keyspace) + qplan = list(policy.make_query_plan(keyspace, query)) + child_policy.make_query_plan.assert_called_once_with(keyspace, query) + + # Test case 2: Without routing key (should call once) + child_policy.reset_mock() + query = Statement(routing_key=None, keyspace=keyspace) + qplan = list(policy.make_query_plan(keyspace, query)) + child_policy.make_query_plan.assert_called_once_with(keyspace, query) + + # Test case 3: Without keyspace (should call once) + child_policy.reset_mock() + query = Statement(routing_key=routing_key, keyspace=None) + qplan = list(policy.make_query_plan(None, query)) + child_policy.make_query_plan.assert_called_once_with(None, query) + + # Test case 4: With tablets (should call once) + cluster.metadata._tablets.table_has_tablets.return_value = True + tablet = Mock(spec=Tablet) + tablet.replicas = [(hosts[0].host_id, None), (hosts[1].host_id, None)] + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.token_map = Mock() + cluster.metadata.token_map.token_class = Mock() + cluster.metadata.token_map.token_class.from_key.return_value = 'token' + + child_policy.reset_mock() + query = Statement(routing_key=routing_key, keyspace=keyspace, table='test_table') + qplan = list(policy.make_query_plan(keyspace, query)) + child_policy.make_query_plan.assert_called_once_with(keyspace, query) + class ConvictionPolicyTest(unittest.TestCase): def test_not_implemented(self):