diff --git a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs index 45dd733824..9c2924e227 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs @@ -270,50 +270,57 @@ public async Task> GetFeedRangesAsync( trace, cancellationToken); - IReadOnlyList partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync( - containerRId, - ContainerCore.allRanges, - trace, - forceRefresh: true); - - if (partitionKeyRanges == null) - { - string refreshedContainerRId; - refreshedContainerRId = await this.GetCachedRIDAsync( - forceRefresh: true, - trace, - cancellationToken); - - if (string.Equals(containerRId, refreshedContainerRId)) - { - throw CosmosExceptionFactory.CreateInternalServerErrorException( - $"Container rid {containerRId} did not have a partition key range after refresh", - headers: new Headers(), - trace: trace); - } - - partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync( - refreshedContainerRId, - ContainerCore.allRanges, - trace, - forceRefresh: true); - - if (partitionKeyRanges == null) - { - throw CosmosExceptionFactory.CreateInternalServerErrorException( - $"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh", - headers: new Headers(), - trace: trace); - } - } - - List feedTokens = new List(partitionKeyRanges.Count); - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange())); + try + { + IReadOnlyList partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync( + containerRId, + ContainerCore.allRanges, + trace, + forceRefresh: true); + + if (partitionKeyRanges == null) + { + string refreshedContainerRId; + refreshedContainerRId = await this.GetCachedRIDAsync( + forceRefresh: true, + trace, + cancellationToken); + + if (string.Equals(containerRId, refreshedContainerRId)) + { + throw CosmosExceptionFactory.CreateInternalServerErrorException( + $"Container rid {containerRId} did not have a partition key range after refresh", + headers: new Headers(), + trace: trace); + } + + partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync( + refreshedContainerRId, + ContainerCore.allRanges, + trace, + forceRefresh: true); + + if (partitionKeyRanges == null) + { + throw CosmosExceptionFactory.CreateInternalServerErrorException( + $"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh", + headers: new Headers(), + trace: trace); + } + } + + List feedTokens = new List(partitionKeyRanges.Count); + foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) + { + feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange())); + } + + return feedTokens; + } + catch (DocumentClientException dce) + { + throw CosmosExceptionFactory.Create(dce, trace); } - - return feedTokens; } public override FeedIterator GetChangeFeedStreamIterator( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs index 8a6b905ac6..3476f567ba 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs @@ -14,6 +14,8 @@ namespace Microsoft.Azure.Cosmos.Tests.FeedRange using Microsoft.Azure.Cosmos.Routing; using Moq; using Microsoft.Azure.Cosmos.Tracing; + using System.Net.Http; + using System.Text; [TestClass] public class FeedRangeTests @@ -220,5 +222,58 @@ public void FeedRangePKRangeId_ToJsonFromJson() Assert.IsNotNull(feedRangePartitionKeyRangeDeserialized); Assert.AreEqual(feedRangePartitionKeyRange.PartitionKeyRangeId, feedRangePartitionKeyRangeDeserialized.PartitionKeyRangeId); } + + /// + /// Upon failures in PartitionKeyRanges calls, the failure should be a CosmosException + /// + [TestMethod] + public async Task GetFeedRangesThrowsCosmosException() + { + Mock mockHttpHandler = new Mock(); + Uri endpoint = MockSetupsHelper.SetupSingleRegionAccount( + "mockAccountInfo", + consistencyLevel: ConsistencyLevel.Session, + mockHttpHandler, + out string primaryRegionEndpoint); + + string databaseName = "mockDbName"; + string containerName = "mockContainerName"; + string containerRid = "ccZ1ANCszwk="; + Documents.ResourceId cRid = Documents.ResourceId.Parse(containerRid); + MockSetupsHelper.SetupContainerProperties( + mockHttpHandler: mockHttpHandler, + regionEndpoint: primaryRegionEndpoint, + databaseName: databaseName, + containerName: containerName, + containerRid: containerRid); + + // Return a 503 on PKRange call + bool invokedPkRanges = false; + Uri partitionKeyUri = new Uri($"{primaryRegionEndpoint}/dbs/{cRid.DatabaseId}/colls/{cRid.DocumentCollectionId}/pkranges"); + mockHttpHandler.Setup(x => x.SendAsync(It.Is(x => x.RequestUri == partitionKeyUri), It.IsAny())) + .Returns(() => Task.FromResult(new HttpResponseMessage() + { + StatusCode = HttpStatusCode.ServiceUnavailable, + Content = new StringContent("ServiceUnavailable") + })) + .Callback(() => invokedPkRanges = true); + + CosmosClientOptions cosmosClientOptions = new CosmosClientOptions() + { + ConsistencyLevel = Cosmos.ConsistencyLevel.Session, + HttpClientFactory = () => new HttpClient(new HttpHandlerHelper(mockHttpHandler.Object)), + }; + + using (CosmosClient customClient = new CosmosClient( + endpoint.ToString(), + Convert.ToBase64String(Encoding.UTF8.GetBytes(Guid.NewGuid().ToString())), + cosmosClientOptions)) + { + Container container = customClient.GetContainer(databaseName, containerName); + CosmosException ex = await Assert.ThrowsExceptionAsync(() => container.GetFeedRangesAsync(CancellationToken.None)); + Assert.AreEqual(HttpStatusCode.ServiceUnavailable, ex.StatusCode); + Assert.IsTrue(invokedPkRanges); + } + } } }