From 949dbedb92cf403786da60c6e9df35f10bdfbde8 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Mon, 23 Jun 2025 19:12:31 -0500 Subject: [PATCH 1/6] fix: HPK QueryRequest Option fix --- .../src/Handler/RequestInvokerHandler.cs | 29 ++++++++++++------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs index 7499ac178e..586336cbbb 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs @@ -305,17 +305,24 @@ public virtual async Task SendAsync( // For epk range filtering we can end up in one of 3 cases: if (overlappingRanges.Count > 1) { - // 1) The EpkRange spans more than one physical partition - // In this case it means we have encountered a split and - // we need to bubble that up to the higher layers to update their datastructures - CosmosException goneException = new CosmosException( - message: $"Epk Range: {feedRangeEpk.Range} is gone.", - statusCode: System.Net.HttpStatusCode.Gone, - subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: Guid.NewGuid().ToString(), - requestCharge: default); - - return goneException.ToCosmosResponseMessage(request); + bool isQueryOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan; + if (!isQueryOperation) + { + // 1) The EpkRange spans more than one physical partition + // In this case it means we have encountered a split and + // we need to bubble that up to the higher layers to update their datastructures + CosmosException goneException = new CosmosException( + message: $"Epk Range: {feedRangeEpk.Range} is gone.", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + return goneException.ToCosmosResponseMessage(request); + } + // For query operations spanning multiple partitions, we don't set PartitionKeyRangeId + // because the query engine needs to handle this case differently + // Just continue without setting PartitionKeyRangeId } // overlappingRanges.Count == 1 else From 02258871c0a2271372687e4bf2c6eaa7d1954db4 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Mon, 23 Jun 2025 19:50:31 -0500 Subject: [PATCH 2/6] fix: adding comments --- Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs index 586336cbbb..9929cb309d 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs @@ -305,6 +305,7 @@ public virtual async Task SendAsync( // For epk range filtering we can end up in one of 3 cases: if (overlappingRanges.Count > 1) { + //If we are running a query and our provided partition key results in a hash that resolves to more than one EPKRanges then its a valid use case bool isQueryOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan; if (!isQueryOperation) { From 8952be82b419af65fd8004a3dc7673502219ec05 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Tue, 24 Jun 2025 11:16:29 -0500 Subject: [PATCH 3/6] fix: adding unit test --- .../RetryHandlerTests.cs | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs index dd771fac1a..d4d35ddfc8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs @@ -5,11 +5,17 @@ namespace Microsoft.Azure.Cosmos.Tests { using System; + using System.Collections.Generic; + using System.IO; using System.Net; using System.Net.Http; using System.Threading; using System.Threading.Tasks; + using global::Azure; + using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.Handlers; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; @@ -18,6 +24,74 @@ namespace Microsoft.Azure.Cosmos.Tests public class RetryHandlerTests { private static readonly Uri TestUri = new Uri("https://dummy.documents.azure.com:443/dbs"); + [TestMethod] + public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnException() + { + using CosmosClient client = MockCosmosUtil.CreateMockCosmosClient(); + + // Create mock container + Mock containerMock = MockCosmosUtil.CreateMockContainer("testDb", "testColl"); + + // Setup container properties + ContainerProperties containerProps = new ContainerProperties("testColl", "/pk"); + var resourceIdProperty = typeof(ContainerProperties).GetProperty( + "ResourceId", + System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + resourceIdProperty.SetValue(containerProps, "testCollRid"); + + // Set up additional mocks as needed + containerMock.Setup(c => c.GetCachedContainerPropertiesAsync( + It.IsAny(), It.IsAny(), It.IsAny())) + .ReturnsAsync(containerProps); + + Mock databaseMock = new Mock(); + databaseMock.Setup(d => d.Id).Returns("testDb"); + containerMock.Setup(c => c.Database).Returns(databaseMock.Object); + + // Mock PartitionKeyRangeCache + List overlappingRanges = new List + { + new PartitionKeyRange { Id = "0", MinInclusive = "A", MaxExclusive = "Z" }, + new PartitionKeyRange { Id = "1", MinInclusive = "M", MaxExclusive = "Z" } + }; + PartitionKeyRangeCache pkRangeCache = new TestPartitionKeyRangeCache(overlappingRanges); + + // FeedRangeEpk for the test + FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range("A", "Z", true, false)); + RequestInvokerHandler invoker = new RequestInvokerHandler(client, null, null, null) + { + InnerHandler = new TestHandler((request, token) => TestHandler.ReturnSuccess()) + }; + + // Act + ResponseMessage response = await invoker.SendAsync("dbs/testDb/colls/testColl", ResourceType.Document, OperationType.QueryPlan, null, containerMock.Object, feedRange, + null,null, NoOpTrace.Singleton, CancellationToken.None); + + //Assert + Assert.IsNotNull(response, "Response should not be null."); + Assert.IsTrue(response.IsSuccessStatusCode, $"Expected a successful status code, but got {response.StatusCode}."); + } + + private class TestPartitionKeyRangeCache : PartitionKeyRangeCache + { + private readonly IReadOnlyList overlappingRanges; + + public TestPartitionKeyRangeCache(IReadOnlyList overlappingRanges) + : base(null, null, null, null) // Pass nulls or mocks as needed for base constructor + { + this.overlappingRanges = overlappingRanges; + } + + public override Task> TryGetOverlappingRangesAsync( + string collectionRid, + Documents.Routing.Range range, + ITrace trace, + bool forceRefresh) + { + return Task.FromResult(this.overlappingRanges); + } + } + [TestMethod] public async Task RetryHandlerDoesNotRetryOnSuccess() From ee5330b9feb63400320f336ae4ecfe2f4028d87e Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Tue, 24 Jun 2025 18:10:39 -0500 Subject: [PATCH 4/6] fix: addressing comments --- Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs | 5 +---- .../tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs | 6 +++--- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs index 9929cb309d..8060586955 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs @@ -305,7 +305,7 @@ public virtual async Task SendAsync( // For epk range filtering we can end up in one of 3 cases: if (overlappingRanges.Count > 1) { - //If we are running a query and our provided partition key results in a hash that resolves to more than one EPKRanges then its a valid use case + //If we are running a query plan and our provided partition key results in a hash that resolves to more than one EPKRanges then its a valid use case bool isQueryOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan; if (!isQueryOperation) { @@ -321,9 +321,6 @@ public virtual async Task SendAsync( return goneException.ToCosmosResponseMessage(request); } - // For query operations spanning multiple partitions, we don't set PartitionKeyRangeId - // because the query engine needs to handle this case differently - // Just continue without setting PartitionKeyRangeId } // overlappingRanges.Count == 1 else diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs index d4d35ddfc8..e81a349e0a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs @@ -51,13 +51,13 @@ public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnExcept // Mock PartitionKeyRangeCache List overlappingRanges = new List { - new PartitionKeyRange { Id = "0", MinInclusive = "A", MaxExclusive = "Z" }, - new PartitionKeyRange { Id = "1", MinInclusive = "M", MaxExclusive = "Z" } + new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" }, + new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" } }; PartitionKeyRangeCache pkRangeCache = new TestPartitionKeyRangeCache(overlappingRanges); // FeedRangeEpk for the test - FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range("A", "Z", true, false)); + FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range("0DCEB8CE51C6BFE84F4BD9409F69B9BB", "0DCEB8CE51C6BFE84F4BD9409F69B9BBFF", true, false)); RequestInvokerHandler invoker = new RequestInvokerHandler(client, null, null, null) { InnerHandler = new TestHandler((request, token) => TestHandler.ReturnSuccess()) From 0060ade767496b0b03cc69b4ba9e34af9f156f15 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Wed, 25 Jun 2025 13:12:13 -0500 Subject: [PATCH 5/6] fix: refactor code --- .../src/Handler/RequestInvokerHandler.cs | 4 +- .../RetryHandlerTests.cs | 63 +++++++++++++++---- 2 files changed, 52 insertions(+), 15 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs index 8060586955..48cee8bf34 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs @@ -306,8 +306,8 @@ public virtual async Task SendAsync( if (overlappingRanges.Count > 1) { //If we are running a query plan and our provided partition key results in a hash that resolves to more than one EPKRanges then its a valid use case - bool isQueryOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan; - if (!isQueryOperation) + bool isQueryPlanOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan; + if (!isQueryPlanOperation) { // 1) The EpkRange spans more than one physical partition // In this case it means we have encountered a split and diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs index e81a349e0a..357c026a7c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs @@ -12,7 +12,6 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Threading; using System.Threading.Tasks; using global::Azure; - using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.Handlers; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.Tracing; @@ -27,7 +26,23 @@ public class RetryHandlerTests [TestMethod] public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnException() { - using CosmosClient client = MockCosmosUtil.CreateMockCosmosClient(); + // Create overlapping ranges for the test + List overlappingRanges = new List + { + new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" }, + new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" } + }; + + // Create a custom document client with our TestPartitionKeyRangeCache + var testPartitionKeyRangeCache = new TestPartitionKeyRangeCache(overlappingRanges); + var customDocClient = new CustomMockDocumentClient(testPartitionKeyRangeCache); + + // Create CosmosClient with our custom document client + using CosmosClient client = new CosmosClient( + "https://localhost:8081", + MockCosmosUtil.RandomInvalidCorrectlyFormatedAuthKey, + new CosmosClientOptions(), + customDocClient); // Create mock container Mock containerMock = MockCosmosUtil.CreateMockContainer("testDb", "testColl"); @@ -48,30 +63,52 @@ public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnExcept databaseMock.Setup(d => d.Id).Returns("testDb"); containerMock.Setup(c => c.Database).Returns(databaseMock.Object); - // Mock PartitionKeyRangeCache - List overlappingRanges = new List - { - new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" }, - new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" } - }; - PartitionKeyRangeCache pkRangeCache = new TestPartitionKeyRangeCache(overlappingRanges); + // FeedRangeEpk for the test - use a range that overlaps both partition key ranges + FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range( + "0DCEB8CE51C6BFE84F4BD9409F69B9BB", // Start just before the boundary + "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579EE", // End just after the boundary + true, false)); - // FeedRangeEpk for the test - FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range("0DCEB8CE51C6BFE84F4BD9409F69B9BB", "0DCEB8CE51C6BFE84F4BD9409F69B9BBFF", true, false)); RequestInvokerHandler invoker = new RequestInvokerHandler(client, null, null, null) { InnerHandler = new TestHandler((request, token) => TestHandler.ReturnSuccess()) }; // Act - ResponseMessage response = await invoker.SendAsync("dbs/testDb/colls/testColl", ResourceType.Document, OperationType.QueryPlan, null, containerMock.Object, feedRange, - null,null, NoOpTrace.Singleton, CancellationToken.None); + ResponseMessage response = await invoker.SendAsync( + "dbs/testDb/colls/testColl", + ResourceType.Document, + OperationType.QueryPlan, + null, + containerMock.Object, + feedRange, + null, + null, + NoOpTrace.Singleton, + CancellationToken.None); //Assert Assert.IsNotNull(response, "Response should not be null."); Assert.IsTrue(response.IsSuccessStatusCode, $"Expected a successful status code, but got {response.StatusCode}."); } + // Custom MockDocumentClient that allows injecting our TestPartitionKeyRangeCache + private class CustomMockDocumentClient : MockDocumentClient + { + private readonly TestPartitionKeyRangeCache testPartitionKeyRangeCache; + + public CustomMockDocumentClient(TestPartitionKeyRangeCache testPartitionKeyRangeCache) + : base(new ConnectionPolicy()) + { + this.testPartitionKeyRangeCache = testPartitionKeyRangeCache; + } + + internal override Task GetPartitionKeyRangeCacheAsync(ITrace trace) + { + return Task.FromResult(this.testPartitionKeyRangeCache); + } + } + private class TestPartitionKeyRangeCache : PartitionKeyRangeCache { private readonly IReadOnlyList overlappingRanges; From d46fb4efab010021e28f8e0d8b4852f362046869 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Wed, 25 Jun 2025 13:44:04 -0500 Subject: [PATCH 6/6] fix: adding unit test --- .../RetryHandlerTests.cs | 43 +++++++++++++++---- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs index 357c026a7c..ca52e5d356 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs @@ -24,13 +24,30 @@ public class RetryHandlerTests { private static readonly Uri TestUri = new Uri("https://dummy.documents.azure.com:443/dbs"); [TestMethod] - public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnException() + public async Task ValidateQueryPlanDoesNotThrowExceptionForOverlappingRanges() + { + await this.ValidateOverlappingRangesBehaviorAsync( + operationType: OperationType.QueryPlan, + shouldThrowGoneException: false); + } + + [TestMethod] + public async Task ValidateQueryThrowsGoneExceptionForOverlappingRanges() + { + await this.ValidateOverlappingRangesBehaviorAsync( + operationType: OperationType.Query, + shouldThrowGoneException: true); + } + + private async Task ValidateOverlappingRangesBehaviorAsync( + OperationType operationType, + bool shouldThrowGoneException) { // Create overlapping ranges for the test List overlappingRanges = new List { - new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" }, - new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" } + new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" }, + new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" } }; // Create a custom document client with our TestPartitionKeyRangeCache @@ -65,8 +82,8 @@ public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnExcept // FeedRangeEpk for the test - use a range that overlaps both partition key ranges FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range( - "0DCEB8CE51C6BFE84F4BD9409F69B9BB", // Start just before the boundary - "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579EE", // End just after the boundary + "0DCEB8CE51C6BFE84F4BD9409F69B9BB", + "0DCEB8CE51C6BFE84F4BD9409F69B9BBFF", true, false)); RequestInvokerHandler invoker = new RequestInvokerHandler(client, null, null, null) @@ -78,7 +95,7 @@ public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnExcept ResponseMessage response = await invoker.SendAsync( "dbs/testDb/colls/testColl", ResourceType.Document, - OperationType.QueryPlan, + operationType, null, containerMock.Object, feedRange, @@ -87,9 +104,19 @@ public async Task ValidatePassingOverlappingRangesInQueryPlanDoesntThrowAnExcept NoOpTrace.Singleton, CancellationToken.None); - //Assert + // Assert Assert.IsNotNull(response, "Response should not be null."); - Assert.IsTrue(response.IsSuccessStatusCode, $"Expected a successful status code, but got {response.StatusCode}."); + + if (shouldThrowGoneException) + { + Assert.IsFalse(response.IsSuccessStatusCode, "Expected a failure status code for Query operation."); + Assert.AreEqual(HttpStatusCode.Gone, response.StatusCode, "Expected a 410 Gone status code."); + Assert.AreEqual((int)SubStatusCodes.PartitionKeyRangeGone, (int)response.Headers.SubStatusCode, "Expected PartitionKeyRangeGone sub-status code."); + } + else + { + Assert.IsTrue(response.IsSuccessStatusCode, $"Expected a successful status code, but got {response.StatusCode}."); + } } // Custom MockDocumentClient that allows injecting our TestPartitionKeyRangeCache