Cosmos: resolve-by-RID, addressing validation, and RID cache (2/3)

sdk/cosmos/.cspell.json

@@ -219,6 +219,7 @@
     "undrained",
     "unfaulted",
     "unparseable",
+    "unrepresentable",
     "unroutable",
     "unsharded",
     "upsert",

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added driver-internal resolution of containers by resource id (RID). `CosmosDriver::resolve_container_by_rid` reads a container's metadata addressing it purely by RID (deriving the parent database RID from the container RID, so no `read_database` round-trip is needed) and caches the result in a by-RID index. References are validated for consistent name/RID addressing in `plan_operation` — the single choke point every executable operation passes through, including multi-page queries — returning a deterministic `CLIENT_MIXED_NAME_RID_ADDRESSING` error before signing instead of letting the gateway reject a mixed name/RID request with an opaque `401`. The new `CLIENT_INVALID_RESOURCE_ID` and `CLIENT_MIXED_NAME_RID_ADDRESSING` client statuses carry searchable names for diagnostics. ([#4663](https://github.com/Azure/azure-sdk-for-rust/pull/4663))
+
 ### Breaking Changes
 
 - Resource-reference accessors now return `Option` to account for RID-addressed references that have no name. `DatabaseReference::name_based_path`, `ContainerReference::database_name`, and `ContainerReference::name_based_path` return `None` when the reference is addressed by RID (previously they returned `&str`/`String` and assumed a name was always present). Use the new `ContainerReference::base_path` to obtain the addressing-appropriate path (RID-based or name-based) when building request URLs. ([#4640](https://github.com/Azure/azure-sdk-for-rust/pull/4640))

sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/container_cache.rs

@@ -5,8 +5,10 @@
 //!
 //! Maintains two lookup indices — by name and by RID — so that a resolved
 //! [`ContainerReference`] can be retrieved efficiently regardless of which
-//! identifier the caller has. When a reference is fetched or inserted,
-//! both caches are cross-populated to keep them in sync.
+//! identifier the caller has. Name-addressed references are cross-populated
+//! into both indices to keep them in sync; RID-addressed references are stored
+//! only in the by-RID index, since they carry no database name to form a safe
+//! by-name key.
 
 use super::AsyncCache;
 use crate::models::ContainerReference;
@@ -25,21 +27,14 @@ struct ContainerNameKey {
 }
 
 impl ContainerNameKey {
-    fn from_container(c: &ContainerReference) -> Self {
-        Self {
+    /// Builds a name key from a container reference, or `None` if the container
+    /// was addressed by RID (no database name is available to key on).
+    fn from_container(c: &ContainerReference) -> Option<Self> {
+        Some(Self {
             account_endpoint: c.account().endpoint().as_str().to_owned(),
-            // TODO(Slice 2): `database_name()` is `None` for a RID-addressed
-            // container, so `unwrap_or_default()` collapses the key to
-            // `{endpoint, "", container_name}`. That is harmless today because
-            // nothing produces RID-addressed references yet, but once Slice 2's
-            // resolution layer starts caching them, two RID-resolved containers
-            // that share a name across different databases would alias to the
-            // same by-name key. Slice 2 must skip the by-name index for
-            // RID-addressed containers (or key it differently) before relying on
-            // this path.
-            db_name: c.database_name().unwrap_or_default().to_owned(),
+            db_name: c.database_name()?.to_owned(),
             container_name: c.name().to_owned(),
-        }
+        })
     }
 }
 
@@ -70,8 +65,10 @@ impl ContainerRidKey {
 ///
 /// Stores fully-resolved [`ContainerReference`] values and indexes them by
 /// both name (`account + db_name + container_name`) and RID
-/// (`account + container_rid`). When a reference is fetched or inserted via
-/// either index, the other index is cross-populated automatically.
+/// (`account + container_rid`). A name-addressed reference fetched or inserted
+/// via either index is cross-populated into the other automatically; a
+/// RID-addressed reference is stored only in the by-RID index, because it has
+/// no database name to form a safe by-name key.
 ///
 /// Uses single-pending-I/O semantics per key — concurrent requests for the
 /// same container share one fetch operation.
@@ -114,6 +111,29 @@ impl ContainerCache {
         self.get_or_fetch_impl(&self.by_name, key, fetch_fn).await
     }
 
+    /// Looks up a container by RID, fetching if not cached.
+    ///
+    /// On a cache miss, calls `fetch_fn` to resolve the container from the
+    /// service. The resolved (RID-addressed) reference is populated into the
+    /// by-RID cache; the by-name cache is left untouched because a RID-addressed
+    /// reference has no database name to key on.
+    pub(crate) async fn get_or_fetch_by_rid<F, Fut>(
+        &self,
+        account_endpoint: &str,
+        container_rid: &str,
+        fetch_fn: F,
+    ) -> crate::error::Result<Arc<ContainerReference>>
+    where
+        F: FnOnce() -> Fut,
+        Fut: std::future::Future<Output = crate::error::Result<ContainerReference>>,
+    {
+        let key = ContainerRidKey {
+            account_endpoint: account_endpoint.to_owned(),
+            container_rid: container_rid.to_owned(),
+        };
+        self.get_or_fetch_impl(&self.by_rid, key, fetch_fn).await
+    }
+
     /// Returns a cached container looked up by name, or `None` if not cached.
     #[allow(dead_code)] // Used in tests; will be called from production code once lookup-by-name is wired up.
     pub(crate) async fn get_by_name(
@@ -197,17 +217,20 @@ impl ContainerCache {
     /// Inserts a known-resolved container reference into both caches.
     ///
     /// If an entry already exists under either key, the existing entry is
-    /// preserved (first-write-wins).
+    /// preserved (first-write-wins). RID-addressed references are inserted only
+    /// into the by-RID cache, since they carry no database name to key on.
     pub(crate) async fn put(&self, container: ContainerReference) {
         let name_key = ContainerNameKey::from_container(&container);
         let rid_key = ContainerRidKey::from_container(&container);
-        let container_for_rid = container.clone();
 
-        self.by_name
-            .get_or_insert_with(name_key, || async { Ok(container) })
-            .await;
+        if let Some(name_key) = name_key {
+            let container_for_name = container.clone();
+            self.by_name
+                .get_or_insert_with(name_key, || async { Ok(container_for_name) })
+                .await;
+        }
         self.by_rid
-            .get_or_insert_with(rid_key, || async { Ok(container_for_rid) })
+            .get_or_insert_with(rid_key, || async { Ok(container) })
             .await;
     }
 }
@@ -261,6 +284,16 @@ mod tests {
         )
     }
 
+    fn test_container_by_rid(db_rid: &str, container_rid: &str) -> ContainerReference {
+        ContainerReference::new_by_rid(
+            test_account(),
+            db_rid.to_owned(),
+            "testcontainer".to_owned(),
+            container_rid.to_owned(),
+            &test_container_props(),
+        )
+    }
+
     // --- get_or_fetch_by_name ---
 
     #[tokio::test]
@@ -343,6 +376,64 @@ mod tests {
         assert!(cache.get_by_rid(ACCOUNT_ENDPOINT, &rid).await.is_some());
     }
 
+    // --- get_or_fetch_by_rid ---
+
+    #[tokio::test]
+    async fn fetch_by_rid_caches_without_name_index() {
+        let cache = ContainerCache::new();
+        let counter = Arc::new(AtomicUsize::new(0));
+
+        let container = test_container_by_rid("db_rid", "coll_rid");
+        let container_clone = container.clone();
+        let counter_clone = counter.clone();
+
+        let resolved = cache
+            .get_or_fetch_by_rid(ACCOUNT_ENDPOINT, "coll_rid", || async move {
+                counter_clone.fetch_add(1, Ordering::SeqCst);
+                Ok(container_clone)
+            })
+            .await
+            .unwrap();
+
+        assert!(resolved.is_by_rid());
+        assert_eq!(counter.load(Ordering::SeqCst), 1);
+
+        // Retrievable by RID.
+        assert!(cache
+            .get_by_rid(ACCOUNT_ENDPOINT, "coll_rid")
+            .await
+            .is_some());
+
+        // A second fetch is served from cache, not the factory.
+        let counter2 = counter.clone();
+        cache
+            .get_or_fetch_by_rid(ACCOUNT_ENDPOINT, "coll_rid", || async move {
+                counter2.fetch_add(1, Ordering::SeqCst);
+                Ok(test_container_by_rid("db_rid", "coll_rid"))
+            })
+            .await
+            .unwrap();
+        assert_eq!(counter.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn put_rid_only_skips_name_cache() {
+        let cache = ContainerCache::new();
+        let container = test_container_by_rid("db_rid", "coll_rid");
+
+        cache.put(container).await;
+
+        // The by-RID index is populated, the by-name index is not (no db name).
+        assert!(cache
+            .get_by_rid(ACCOUNT_ENDPOINT, "coll_rid")
+            .await
+            .is_some());
+        assert!(cache
+            .get_by_name(ACCOUNT_ENDPOINT, "db_rid", "testcontainer")
+            .await
+            .is_none());
+    }
+
     // --- different containers ---
 
     #[tokio::test]

sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs

@@ -1094,6 +1094,87 @@ impl CosmosDriver {
         ))
     }
 
+    /// Fetches a container's metadata from the service addressing it purely by RID.
+    ///
+    /// The parent database RID is derived from the container RID's encoded byte
+    /// layout (the first 4 decoded bytes), so no `read_database` round-trip is
+    /// needed. The read response supplies the container's name and partition key.
+    /// The resulting [`ContainerReference`] is RID-addressed (it carries no
+    /// database name).
+    async fn fetch_container_by_rid(
+        &self,
+        container_rid: &str,
+    ) -> crate::error::Result<ContainerReference> {
+        // A container RID decodes to at least 8 bytes: the first 4 identify the
+        // parent database, the next 4 the container. Anything shorter (e.g. a
+        // bare database RID) is not a container RID — fail fast rather than
+        // issuing a request that the service would reject.
+        let decoded = crate::models::resource_id::decode_rid(container_rid).map_err(|e| {
+            crate::error::CosmosError::builder()
+                .with_status(crate::error::CosmosStatus::CLIENT_INVALID_RESOURCE_ID)
+                .with_message(format!("invalid container RID '{container_rid}'"))
+                .with_source(e)
+                .build()
+        })?;
+        if decoded.len() < 8 {
+            return Err(crate::error::CosmosError::builder()
+                .with_status(crate::error::CosmosStatus::CLIENT_INVALID_RESOURCE_ID)
+                .with_message(format!(
+                    "'{container_rid}' is not a container RID (decodes to {} bytes; a container RID requires at least 8)",
+                    decoded.len()
+                ))
+                .with_source(std::io::Error::other("container RID too short"))
+                .build());
+        }
+        let db_rid = crate::models::resource_id::ResourceId::new(
+            crate::models::resource_id::encode_rid(&decoded[0..4]),
+        );
+
+        let options = OperationOptions::default();
+
+        let container_result = self
+            .execute_singleton_operation(
+                CosmosOperation::read_container_by_rid(
+                    self.account().clone(),
+                    db_rid.as_str().to_owned(),
+                    container_rid.to_owned(),
+                ),
+                options,
+            )
+            .await?;
+        let container_headers = container_result.headers().clone();
+        let container_diagnostics = container_result.diagnostics();
+        let container_props: ContainerProperties =
+            container_result.into_body().into_single().map_err(|e| {
+                crate::error::CosmosError::builder()
+                    .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID)
+                    .with_message("failed to deserialize container response")
+                    .with_response_parts(crate::models::CosmosResponsePayload::new(
+                        crate::models::ResponseBody::NoPayload,
+                        container_headers.clone(),
+                    ))
+                    .with_diagnostics(container_diagnostics.clone())
+                    .with_source(e)
+                    .build()
+            })?;
+
+        // Prefer the authoritative RID echoed back by the service; fall back to
+        // the caller-supplied RID if the response omits it.
+        let resolved_rid = container_props
+            .system_properties
+            .rid
+            .clone()
+            .unwrap_or_else(|| container_rid.to_owned());
+
+        Ok(ContainerReference::new_by_rid(
+            self.account().clone(),
+            db_rid.as_str().to_owned(),
+            container_props.id.clone().into_owned(),
+            resolved_rid,
+            &container_props,
+        ))
+    }
+
     /// Creates a new driver instance.
     ///
     /// This is internal - use [`CosmosDriverRuntime::create_driver()`] instead.
@@ -2141,6 +2222,38 @@ impl CosmosDriver {
         Ok(resolved.as_ref().clone())
     }
 
+    /// Resolves a container by its RID.
+    ///
+    /// Attempts to resolve from `ContainerCache` (by-RID index) first. On a cache
+    /// miss, fetches metadata from the service addressing the container by RID and
+    /// populates the cache. The returned [`ContainerReference`] is RID-addressed
+    /// (it carries no database name).
+    pub async fn resolve_container_by_rid(
+        &self,
+        container_rid: &str,
+    ) -> crate::error::Result<ContainerReference> {
+        let endpoint = self.account().endpoint().as_str().to_owned();
+        let container_rid_owned = container_rid.to_owned();
+
+        let resolved = self
+            .runtime
+            .container_cache()
+            .get_or_fetch_by_rid(&endpoint, container_rid, || async move {
+                self.fetch_container_by_rid(&container_rid_owned)
+                    .await
+                    .map_err(|err| {
+                        crate::error::CosmosErrorBuilder::from_error(err)
+                            .with_context(format!(
+                                "resolve container by rid (container_rid='{container_rid_owned}')"
+                            ))
+                            .build()
+                    })
+            })
+            .await?;
+
+        Ok(resolved.as_ref().clone())
+    }
+
     /// Plans the execution of a Cosmos DB operation.
     ///
     /// For trivial operations (non-query or single-partition), returns a
@@ -2161,6 +2274,19 @@ impl CosmosDriver {
         options: &OperationOptions,
         continuation: Option<&ContinuationToken>,
     ) -> crate::error::Result<OperationPlan> {
+        // Reject mixed name/RID addressing before any IO work is done. A
+        // reference that mixes a name-addressed parent with a RID-addressed child
+        // (or vice versa) signs and routes inconsistently and the gateway rejects
+        // it with an opaque 401. Validating here — the single choke point every
+        // externally executable operation passes through on its way to a plan —
+        // turns that into a deterministic client-side error for references built
+        // through any path (including direct `plan_operation` + `execute_plan`
+        // callers and multi-page queries), not just `execute_operation`. The
+        // check is a cheap in-memory field comparison and a no-op for every
+        // consistently-addressed reference, so it issues no additional network
+        // calls and does not change the request flow.
+        operation.resource_reference().validate_addressing()?;
+
         if !self.initialized.load(Ordering::Acquire) {
             let endpoint = AccountEndpoint::from(self.options.account());
             return Err(crate::error::CosmosError::builder()
@@ -2636,6 +2762,39 @@ mod tests {
         )
     }
 
+    #[tokio::test]
+    async fn plan_operation_rejects_mixed_addressing() {
+        // The mixed name/RID guard lives in plan_operation, the single choke
+        // point every externally executable operation passes through. Calling
+        // plan_operation directly (the path queries and direct callers use,
+        // bypassing execute_operation) must still reject a mixed reference with
+        // the deterministic CLIENT_MIXED_NAME_RID_ADDRESSING error before any
+        // plan is built or signed.
+        let cosmos_runtime = CosmosDriverRuntimeBuilder::new().build().await.unwrap();
+        let driver_options = DriverOptions::builder(test_account()).build();
+        let driver = CosmosDriver::new(cosmos_runtime, driver_options)
+            .expect("CosmosDriver::new should succeed in tests");
+
+        // Name-addressed database parent with a RID-addressed collection leaf.
+        let db = DatabaseReference::from_name(test_account(), "testdb");
+        let mixed = crate::models::CosmosResourceReference::from(db)
+            .with_resource_type(ResourceType::DocumentCollection)
+            .with_rid("Lx1BALxJyZ8=".into());
+        let operation = CosmosOperation::new(crate::models::OperationType::Read, mixed, None);
+
+        let err = match driver
+            .plan_operation(operation, &OperationOptions::default(), None)
+            .await
+        {
+            Ok(_) => panic!("mixed addressing must be rejected before planning"),
+            Err(e) => e,
+        };
+        assert_eq!(
+            err.status(),
+            crate::error::CosmosStatus::CLIENT_MIXED_NAME_RID_ADDRESSING
+        );
+    }
+
     #[tokio::test]
     async fn default_operation_options() {
         let runtime = CosmosDriverRuntimeBuilder::new().build().await.unwrap();