Revert "PartitionKey: Fixes EnsureIdGetAppendedToPartitionKeyIfNeededAsync to use in-memory cache only"

This reverts commit 2f8b90b.

Author: NaluTripician

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -616,12 +616,6 @@ internal virtual async Task<ClientCollectionCache> GetCollectionCacheAsync(ITrac
             }
         }
 
-        /// <summary>
-        /// Returns the collection cache if the client is already initialized.
-        /// Does not trigger initialization. Returns null if the client has not been initialized yet.
-        /// </summary>
-        internal ClientCollectionCache CollectionCacheIfInitialized => this.collectionCache;
-
         internal virtual async Task<PartitionKeyRangeCache> GetPartitionKeyRangeCacheAsync(ITrace trace)
         {
             using (ITrace childTrace = trace.StartChild("Get Partition Key Range Cache", TraceComponent.Routing, Tracing.TraceLevel.Info))

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs

@@ -441,21 +441,6 @@ public override async Task<ContainerProperties> GetCachedContainerPropertiesAsyn
             }
         }
 
-        internal override bool TryGetCachedContainerPropertiesFromMemory(out ContainerProperties containerProperties)
-        {
-            containerProperties = null;
-            ClientCollectionCache cache = this.ClientContext.DocumentClient.CollectionCacheIfInitialized;
-            if (cache == null)
-            {
-                return false;
-            }
-
-            return cache.TryGetCachedContainerProperties(
-                HttpConstants.Versions.CurrentVersion,
-                this.LinkUri,
-                out containerProperties);
-        }
-
         // Name based look-up, needs re-computation and can't be cached
         public override async Task<string> GetCachedRIDAsync(
             bool forceRefresh,

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerInternal.cs

@@ -55,17 +55,6 @@ public abstract Task<ContainerProperties> GetCachedContainerPropertiesAsync(
             ITrace trace,
             CancellationToken cancellationToken);
 
-        /// <summary>
-        /// Attempts to retrieve container properties from in-memory cache only.
-        /// Does not trigger client initialization or make HTTP calls.
-        /// Returns false if the cache is not populated.
-        /// </summary>
-        internal virtual bool TryGetCachedContainerPropertiesFromMemory(out ContainerProperties containerProperties)
-        {
-            containerProperties = null;
-            return false;
-        }
-
         public abstract Task<IReadOnlyList<IReadOnlyList<string>>> GetPartitionKeyPathTokensAsync(
             ITrace trace,
             CancellationToken cancellationToken = default);

Microsoft.Azure.Cosmos/src/Resource/Settings/ContainerPropertiesExtensions.cs

@@ -19,28 +19,35 @@ internal static class ContainerPropertiesExtensions
     {
         private const int DefaultStreamBufferSize = 4096;
 
-        internal static Task<PartitionKey?> EnsureIdGetAppendedToPartitionKeyIfNeededAsync(
+        internal static async Task<PartitionKey?> EnsureIdGetAppendedToPartitionKeyIfNeededAsync(
             this ContainerInternal container,
             PartitionKey? partitionKey,
             string itemId,
             CancellationToken cancellationToken)
         {
             if (container == null)
             {
-                return Task.FromResult(partitionKey);
+                return partitionKey;
             }
 
-            // Only use container properties that are already cached in memory.
-            // This avoids triggering client initialization (EnsureValidClientAsync)
-            // or making HTTP calls (ResolveByNameAsync) before the handler pipeline
-            // has had a chance to initialize the client.
-            // On subsequent requests, the cache will be populated by the normal
-            // request pipeline and this optimization will take effect.
-            if (!container.TryGetCachedContainerPropertiesFromMemory(out ContainerProperties containerProperties)
-                || containerProperties == null
-                || !containerProperties.IsLastPartitionKeyPathId)
+            ContainerProperties containerProperties;
+            try
+            {
+                containerProperties = await container.GetCachedContainerPropertiesAsync(
+                    forceRefresh: false,
+                    trace: NoOpTrace.Singleton,
+                    cancellationToken: cancellationToken);
+            }
+            catch (Exception)
+            {
+                // Container may not exist yet (e.g. operations on non-existent containers).
+                // Swallow the exception and let the actual operation surface the proper error.
+                return partitionKey;
+            }
+
+            if (containerProperties == null || !containerProperties.IsLastPartitionKeyPathId)
             {
-                return Task.FromResult(partitionKey);
+                return partitionKey;
             }
 
             if (string.IsNullOrEmpty(itemId))
@@ -56,7 +63,7 @@ internal static class ContainerPropertiesExtensions
 
                 if (existingComponents.Count != partitionKeyPaths.Count - 1)
                 {
-                    return Task.FromResult(partitionKey);
+                    return partitionKey;
                 }
             }
 
@@ -86,7 +93,7 @@ internal static class ContainerPropertiesExtensions
 
             Documents.Routing.PartitionKeyInternal partitionKeyInternal = new Documents.Routing.PartitionKeyInternal(allComponentsList);
 
-            return Task.FromResult<PartitionKey?>(new PartitionKey(partitionKeyInternal));
+            return new PartitionKey(partitionKeyInternal);
         }
 
         internal static async Task<(string, Stream)> GetItemIdFromStreamIfRequiredAsync(

Microsoft.Azure.Cosmos/src/Routing/AsyncCache.cs

@@ -198,27 +198,6 @@ public void Remove(TKey key)
             }
         }
 
-        /// <summary>
-        /// Attempts to get a value from the cache without triggering any async operations.
-        /// Returns true only if the value is already cached and completed successfully.
-        /// </summary>
-        public bool TryGetCachedValue(TKey key, out TValue value)
-        {
-            value = default;
-            if (this.values.TryGetValue(key, out AsyncLazy<TValue> lazyValue)
-                && lazyValue.IsValueCreated
-                && lazyValue.Value.IsCompleted
-                && lazyValue.Value.Status == TaskStatus.RanToCompletion)
-            {
-#pragma warning disable VSTHRD002 // Safe: task is already completed successfully
-                value = lazyValue.Value.Result;
-#pragma warning restore VSTHRD002
-                return true;
-            }
-
-            return false;
-        }
-
         public bool TryRemoveIfCompleted(TKey key)
         {
             AsyncLazy<TValue> initialLazyValue;

Microsoft.Azure.Cosmos/src/Routing/CollectionCache.cs

@@ -271,21 +271,6 @@ private Task<ContainerProperties> ResolveByRidAsync(
                 cancellationToken);
         }
 
-        /// <summary>
-        /// Attempts to retrieve cached container properties from in-memory cache only.
-        /// Does not trigger any HTTP calls or async initialization.
-        /// Returns true if the properties are already cached and available.
-        /// </summary>
-        internal bool TryGetCachedContainerProperties(
-            string apiVersion,
-            string resourceAddress,
-            out ContainerProperties containerProperties)
-        {
-            string resourceFullName = PathsHelper.GetCollectionPath(resourceAddress);
-            InternalCache cache = this.GetCache(apiVersion);
-            return cache.collectionInfoByName.TryGetCachedValue(resourceFullName, out containerProperties);
-        }
-
         internal virtual async Task<ContainerProperties> ResolveByNameAsync(
             string apiVersion,
             string resourceAddress,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs

@@ -1208,8 +1208,6 @@ private static Mock<ContainerInternal> CreateMockContainerWithProperties(Contain
                     It.IsAny<ITrace>(),
                     It.IsAny<CancellationToken>()))
                 .ReturnsAsync(containerProperties);
-            mock.Setup(x => x.TryGetCachedContainerPropertiesFromMemory(out containerProperties))
-                .Returns(true);
             return mock;
         }