Fix TokenAwarePolicy to call child.make_query_plan once#634
Closed
Fix TokenAwarePolicy to call child.make_query_plan once#634
Conversation
|
@copilot please complete the work on this PR. |
Co-authored-by: mykaul <4655593+mykaul@users.noreply.github.com>
Author
Completed! The optimization ensures Summary of changes:
|
Copilot
AI
changed the title
[WIP] Fix TokenAwarePolicy to call child.make_query_plan once
Fix TokenAwarePolicy to call child.make_query_plan once
Dec 29, 2025
|
@copilot address review comments from your friendly AI codeQL above. |
Co-authored-by: mykaul <4655593+mykaul@users.noreply.github.com>
Co-authored-by: mykaul <4655593+mykaul@users.noreply.github.com>
Author
Addressed the memory usage concern from the code review. See commits ec8203e and 85c6ee8. Changes made:
All 83 tests pass and no security issues detected. |
mykaul
reviewed
Jan 20, 2026
…ts call Co-authored-by: mykaul <4655593+mykaul@users.noreply.github.com>
| keyspace = 'keyspace' | ||
| routing_key = 'routing_key' | ||
| query = Statement(routing_key=routing_key, keyspace=keyspace) | ||
| qplan = list(policy.make_query_plan(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)) |
| # 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.reset_mock() | ||
| query = Statement(routing_key=routing_key, keyspace=keyspace, table='test_table') | ||
| qplan = list(policy.make_query_plan(keyspace, query)) |
mykaul
reviewed
Jan 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pre-review checklist
./docs/source/.Description
Optimizes
TokenAwarePolicy.make_query_planto callchild.make_query_planonly 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_plancould callchild.make_query_planup to three times:This redundant computation was unnecessary and could impact performance, especially with complex child policies. Additionally, the tablet routing code had a redundant
table_has_tabletscheck before callingget_tablet_for_key.Changes Made
cassandra/policies.py
TokenAwarePolicy.make_query_planto callchild.make_query_planonly oncechild.make_query_planonce and convert to list for reuse (needed for replica identification and yielding remaining hosts)table_has_tabletscheck - now directly callsget_tablet_for_keyand checks if the result is not None (sinceget_tablet_for_keyalready returns None when there are no tablets)tests/unit/test_policies.py
test_child_make_query_plan_called_oncethat validates the optimization works in all scenarios:_assert_shuffleto expect one call for both tablet and vnode scenariosget_tablet_for_key.return_valueinstead of the removedtable_has_tablets.return_valuePerformance & Memory Impact
child.make_query_planin worst-case scenarioschild.make_query_planTesting
tests/unit/test_policies.pypassOriginal prompt
💡 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.