Cosmos: Route PKRange resolution through driver by simorenoh · Pull Request #4342 · Azure/azure-sdk-for-rust · GitHub

simorenoh · 2026-05-04T22:06:29Z

Migrate read_feed_ranges() and feed_range_from_partition_key() in ContainerClient to use the driver's partition key range cache instead of the SDK-layer cache. This is the first step toward removing the SDK's redundant cache infrastructure.

Changes

cosmos_driver.rs: Add resolve_all_partition_key_ranges() and resolve_partition_key_ranges_for_key() public methods to CosmosDriver, exposing the driver's internal PKRange cache for SDK consumption.
container_client.rs: Rewrite read_feed_ranges() and feed_range_from_partition_key() to delegate to the driver instead of maintaining a separate SDK-layer cache with its own HTTP calls.
feed_range.rs: Remove SDK-specific adapter methods (from_partition_key_range, from_epk_range) that are no longer needed now that the driver returns ranges directly.

Motivation

The SDK previously maintained its own PartitionKeyRangeCache that duplicated the driver's cache. After the PPAF/PPCB cutover (#4156), both caches existed in parallel. This PR routes all PKRange resolution through the driver, making the SDK cache unused and ready for removal in a follow-up PR.

Follow-up

A subsequent PR will remove the now-dead SDK-layer cache infrastructure (partition_key_range_cache.rs, collection_routing_map.rs, old pipeline, retry policies, etc.).

Migrate read_feed_ranges() and feed_range_from_partition_key() in ContainerClient to use the driver's partition key range cache instead of the SDK-layer cache. This is the first step toward removing the SDK's redundant cache infrastructure. Changes: - Add resolve_all_partition_key_ranges() and resolve_partition_key_ranges_for_key() public methods to CosmosDriver - Rewrite ContainerClient::read_feed_ranges() to call driver - Rewrite ContainerClient::feed_range_from_partition_key() to call driver - Remove SDK-specific FeedRange adapter code (from_partition_key_range, from_epk_range) that are no longer needed Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Documents the PR split strategy, remaining deletions, and test coverage gaps for the SDK-to-driver migration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Routes ContainerClient feed-range APIs through the azure_data_cosmos_driver partition key range cache, removing the SDK-layer routing-map adapters as a first step toward deleting redundant SDK caching infrastructure.

Changes:

Add CosmosDriver::{resolve_all_partition_key_ranges, resolve_partition_key_ranges_for_key} as public APIs to expose PK range resolution via the driver cache.
Update ContainerClient::{read_feed_ranges, feed_range_from_partition_key} to delegate PK range resolution to the driver.
Remove SDK-only FeedRange adapter helpers that were tied to the SDK routing map types.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
sdk/cosmos/azure_data_cosmos/src/feed_range.rs	Removes SDK-routing-map-specific conversion helpers/tests; keeps driver `PartitionKeyRange` conversion.
sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs	Switches feed-range APIs to call into the driver for PK range resolution.
sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs	Adds public driver methods to resolve all PK ranges or PK ranges for a given partition key.

Return an error instead of an empty Vec when the driver's PKRange cache resolves to zero ranges after a forced refresh. A valid container always has at least one partition key range, so an empty result indicates a service/auth failure that should not be silently swallowed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Address review findings: 1. Full-key point-lookup miss now returns Some(vec![]) instead of None, allowing the caller's empty-check retry to fire after partition splits. Previously, None short-circuited to an error before the retry logic. 2. Skip redundant retry when force_refresh was already true — avoids hitting the service twice with identical parameters. 3. Log a warning when EPK computation fails instead of silently discarding the error via .ok()?. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

A valid container always has at least one partition key range. When the driver's cache falls back to an empty routing map (e.g., service unreachable or parse failure), return None instead of Some(vec![]) so callers get a clear failure signal rather than a silently empty result. This applies to both resolve_all_partition_key_ranges (empty ranges slice) and the full-key branch of resolve_partition_key_ranges_for_key (empty routing map before EPK lookup). The prefix-key branch already returns None via resolve_overlapping_ranges when the underlying map has no entries. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add tests covering the key scenarios for resolve_all/resolve_for_key: - Full key point lookup returns single range - Prefix key (MultiHash) returns overlapping ranges - Empty routing map (service failure) returns empty vec - Fetch failure returns None from try_lookup - Force refresh repopulates after transient empty cache - All ranges returned from try_lookup in correct order These tests validate the cache behavior exercised by CosmosDriver::resolve_all_partition_key_ranges() and resolve_partition_key_ranges_for_key(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

simorenoh · 2026-05-05T16:19:01Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-05-05T16:19:21Z

Azure Pipelines successfully started running 1 pipeline(s).

analogrelay · 2026-05-05T18:37:47Z

+    #[allow(dead_code)]
    pub(crate) global_endpoint_manager: Arc<GlobalEndpointManager>,
+    #[allow(dead_code)]
    pub(crate) global_partition_endpoint_manager: Arc<GlobalPartitionEndpointManager>,


Same here, let's yeet the dead code now, rather than waiting.

analogrelay · 2026-05-05T18:38:59Z

+    /// Uses the driver's internal `PartitionKeyRangeCache`. When `force_refresh`
+    /// is `true`, the cached routing map is refreshed from the service before
+    /// returning results. Returns `None` if the routing map cannot be resolved.
+    pub async fn resolve_all_partition_key_ranges(


We should consider using FeedRange terminology here. I'm going to be moving FeedRange to the driver as part of the Feed Operations, since it needs to be accessible there. I'd argue that at this point, the driver is the only component that needs to see the actual physical partition metadata (PK Range ID) and everything above it can and should work on FeedRanges.

will be tracked in a follow-up

FabianMeiswinkel · 2026-05-05T19:33:38Z

-            .ordered_partition_key_ranges()
+
+        if ranges.is_empty() && !options.force_refresh() {
+            // A valid container always has at least one partition key range.


Would we even allow an empty range in the cache? Why?

This is part of the existing logic on the pk range cache that was already merged to the driver, was just accommodating to it here - but you're right this shouldn't happen. Will be tracked in a follow-up.

FabianMeiswinkel · 2026-05-05T19:38:05Z

    ///
    /// Full keys return a single-element `Vec`. Prefix keys on MultiHash
    /// containers return one or more feed ranges.
    pub async fn feed_range_from_partition_key(


This name is terrible IMO - maybe get_overlapping_feed_range? Why do we even need this API? We will need one API to produce a FeedRane for a PartitionKey - but just the EPK of teh PK - then a helper function that checks whether FeedRanges are overlapping - maybe whether one feedrange is fully conatined in another - but that is all you need - thsi API gives the wrong impression of being able to identify the "physical partition" for a PK - and that is nothing that ever works reliably - so, we should also nozt give that impression? IMO this APi should just be removed - if customers use it replace it by scalable helper functions we have in Java and .Net as well.

yep absolutely - and no, we don't need it like this. already discussed offline but this should be what you're describing, not what it is currently. will be tracked in a follow-up.

Mostly just summarizing offline conversations here for posterity

I think we can land on FeedRange::for_partition_key (this would basically be the only way to create a FeedRange as a user, but it makes sense for it to live on FeedRange since it doesn't involve any I/O, just hashing). That also aligns with the Python behavior (whereas this Rust code is actually incorrect because it does indeed find the physical partition that currently contains that PK; which isn't what we want)

FabianMeiswinkel · 2026-05-05T19:39:31Z

-                    .collect())
+        if ranges.is_empty() && !options.force_refresh() {
+            // Empty result may indicate a stale cache — retry with refresh.
+            let ranges = self


Like above - when would this ever happen? We should prevent it instead of trying to wroakround it.

ditto, will be tracked in a follow up

Remove global_endpoint_manager and global_partition_endpoint_manager fields from ClientContext — they were stored but never read through this struct. The GatewayPipeline holds its own Arc references. Remove _database_link parameter from ContainerClient::new() since the driver resolves container metadata by ID directly. Remove DatabaseClient.link field (ResourceLink) which was only used to pass to ContainerClient::new(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The cutover-remaining-work.md is a planning artifact that should stay local rather than in the repository. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

analogrelay

LGTM given the discussed follow-ups

analogrelay · 2026-05-06T18:13:38Z

    ///
    /// Full keys return a single-element `Vec`. Prefix keys on MultiHash
    /// containers return one or more feed ranges.
    pub async fn feed_range_from_partition_key(


Mostly just summarizing offline conversations here for posterity

I think we can land on FeedRange::for_partition_key (this would basically be the only way to create a FeedRange as a user, but it makes sense for it to live on FeedRange since it doesn't involve any I/O, just hashing). That also aligns with the Python behavior (whereas this Rust code is actually incorrect because it does indeed find the physical partition that currently contains that PK; which isn't what we want)

simorenoh added the Created by copilot label May 4, 2026

simorenoh added this to CosmosDB Go/Rust Crew May 4, 2026

github-project-automation Bot moved this to Todo in CosmosDB Go/Rust Crew May 4, 2026

simorenoh added the Cosmos The azure_cosmos crate label May 4, 2026

simorenoh linked an issue May 4, 2026 that may be closed by this pull request

Cosmos Driver: Switch read_feed_ranges / feed_range_from_partition_key to driver routing map #4138

Open

4 tasks

simorenoh marked this pull request as ready for review May 4, 2026 22:07

simorenoh requested a review from a team as a code owner May 4, 2026 22:07

Copilot AI review requested due to automatic review settings May 4, 2026 22:07

Copilot started reviewing on behalf of simorenoh May 4, 2026 22:08 View session

Add cutover tracking document

dd52ddb

Documents the PR split strategy, remaining deletions, and test coverage gaps for the SDK-to-driver migration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI reviewed May 4, 2026

View reviewed changes

simorenoh and others added 4 commits May 4, 2026 18:24

simorenoh force-pushed the pkrange-cache-cutover branch from c771c30 to 98b9250 Compare May 5, 2026 00:53

analogrelay reviewed May 5, 2026

View reviewed changes

FabianMeiswinkel reviewed May 5, 2026

View reviewed changes

FabianMeiswinkel requested changes May 5, 2026

View reviewed changes

github-project-automation Bot moved this from Todo to Changes Requested in CosmosDB Go/Rust Crew May 5, 2026

simorenoh and others added 2 commits May 6, 2026 10:38

Remove tracking document from repository

392cf2b

The cutover-remaining-work.md is a planning artifact that should stay local rather than in the repository. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

simorenoh requested review from FabianMeiswinkel and analogrelay May 6, 2026 18:07

analogrelay approved these changes May 6, 2026

View reviewed changes

Conversation

simorenoh commented May 4, 2026

Changes

Motivation

Follow-up

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

simorenoh commented May 5, 2026

Uh oh!

azure-pipelines Bot commented May 5, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simorenoh May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FabianMeiswinkel May 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simorenoh May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

analogrelay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

simorenoh May 6, 2026 •

edited

Loading

FabianMeiswinkel May 5, 2026 •

edited

Loading

simorenoh May 6, 2026 •

edited

Loading