Fix TokenAwarePolicy to call child.make_query_plan once

[Copilot]

Pre-review checklist

• I have split my patch into logically separate commits.
• All commit messages clearly explain what they change and why.
• I added relevant tests for new features and bug fixes.
• All commits compile, pass static checks and pass test.
• PR description sums up the changes and reasons why they should be introduced.
• I have provided docstrings for the public items that I want to introduce.
• I have adjusted the documentation in ./docs/source/.

Description

Optimizes TokenAwarePolicy.make_query_plan to call child.make_query_plan only once instead of up to 3 times in some scenarios, improving performance by eliminating redundant query plan generation while preserving memory efficiency through lazy evaluation. Also simplifies the tablet routing logic by removing redundant checks.

Problem

In the original implementation, TokenAwarePolicy.make_query_plan could call child.make_query_plan up to three times:

1. Line 509: When query/routing_key/keyspace is None (early return)
2. Line 520: When using tablet-based routing
3. Line 538: To get the rest of the cluster hosts

This redundant computation was unnecessary and could impact performance, especially with complex child policies. Additionally, the tablet routing code had a redundant table_has_tablets check before calling get_tablet_for_key.

Changes Made

cassandra/policies.py

• Modified TokenAwarePolicy.make_query_plan to call child.make_query_plan only once
• For early return cases (no routing key/keyspace), pass through the generator directly to preserve lazy evaluation and avoid unnecessary memory usage
• For token-aware routing cases, call child.make_query_plan once and convert to list for reuse (needed for replica identification and yielding remaining hosts)
• Simplified tablet routing by removing redundant table_has_tablets check - now directly calls get_tablet_for_key and checks if the result is not None (since get_tablet_for_key already returns None when there are no tablets)
• Added comprehensive comments explaining the optimization rationale and when list materialization is necessary, including:
  • Why lazy evaluation is preserved for early return cases
  • When and why list conversion occurs (for double iteration during replica identification and final host yielding)

tests/unit/test_policies.py

• Added a new comprehensive test test_child_make_query_plan_called_once that validates the optimization works in all scenarios:
  • With routing key and keyspace
  • Without routing key
  • Without keyspace
  • With tablets enabled
• Updated existing test assertion in _assert_shuffle to expect one call for both tablet and vnode scenarios
• Updated all test mocks to use get_tablet_for_key.return_value instead of the removed table_has_tablets.return_value

Performance & Memory Impact

• Before: Up to 3 calls to child.make_query_plan in worst-case scenarios
• After: Always exactly 1 call to child.make_query_plan
• Memory: Lazy evaluation preserved for early return cases; list materialization only occurs when necessary for token-aware routing (when the query plan must be iterated twice: once for replica identification, once for yielding remaining hosts)
• Code Simplification: Reduced nesting and eliminated redundant tablet check, making the code more maintainable

Testing

• ✅ All 83 tests in tests/unit/test_policies.py pass
• ✅ New test specifically validates the optimization
• ✅ No security vulnerabilities detected by CodeQL scanner
• ✅ Code review verified no memory concerns with the implementation

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Call child.make_query_plan in TokenAwarePolicy.make_query_plan only once</issue_title>
<issue_description>In some scenarios TokenAwarePolicy.make_query_plan can call child.make_query_plan three times.
Let's make sure it is called only once.

python-driver/cassandra/policies.py

Lines 506 to 539 in 2739b29

	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
	if query is None or query.routing_key is None or keyspace is None:
	for host in child.make_query_plan(keyspace, query):
	yield host
	return
	replicas = []
	if self._tablets_routing_v1:
	tablet = self._cluster_metadata._tablets.get_tablet_for_key(
	keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key))
	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]
	if not replicas:
	replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
	if self.shuffle_replicas:
	shuffle(replicas)
	for replica in replicas:
	if replica.is_up and child.distance(replica) in [HostDistance.LOCAL, HostDistance.LOCAL_RACK]:
	yield replica
	for host in child.make_query_plan(keyspace, query):
	# skip if we've already listed this host
	if host not in replicas or child.distance(host) == HostDistance.REMOTE:
	yield host

</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes Call child.make_query_plan in TokenAwarePolicy.make_query_plan only once #358

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[mykaul]

@copilot please complete the work on this PR.

[Copilot]

@copilot please complete the work on this PR.

Completed! The optimization ensures child.make_query_plan is called only once in all scenarios. See commit 507b252.

Summary of changes:

• Modified TokenAwarePolicy.make_query_plan to call child.make_query_plan once and reuse the result
• Added comprehensive test coverage to verify the optimization
• All 83 existing tests pass
• No security issues detected

[mykaul]

@copilot address review comments from your friendly AI codeQL above.

[Copilot]

@copilot address review comments from your friendly AI codeQL above.

Addressed the memory usage concern from the code review. See commits ec8203e and 85c6ee8.

Changes made:

• Early return case now preserves lazy evaluation by passing through the generator directly, avoiding unnecessary memory usage
• List materialization only occurs when token-aware routing is needed (requires two iterations over the plan)
• Added clear comments explaining the optimization rationale

All 83 tests pass and no security issues detected.

[mykaul]

@copilot - this line (if self._cluster_metadata._tablets.table_has_tablets(keyspace, query.table):) could be replaced with the next line (tablet = ...) and then if tablet is not None: that we have already. Can you create a follow-up commit with this change?

[Copilot]

Done. The table_has_tablets check has been replaced with a direct call to get_tablet_for_key, followed by the existing if tablet is not None check. This simplifies the code since get_tablet_for_key already returns None when there are no tablets for the table. See commit d49f4dd.