Fix TokenAwarePolicy to call child.make_query_plan once

[Copilot]

Pre-review checklist

• [x] I have split my patch into logically separate commits.
• [x] All commit messages clearly explain what they change and why.
• [x] I added relevant tests for new features and bug fixes.
• [x] All commits compile, pass static checks and pass test.
• [x] 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.

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.

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)
• Added comprehensive comments explaining the optimization rationale and when list materialization is necessary

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

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)

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.

https://github.com/scylladb/python-driver/blob/2739b29052c5a1745b7cb93ed31ff7af8063bd6e/cassandra/policies.py#L506-L539 </issue_description>

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

• Fixes scylladb/python-driver#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.