Skip to content

[INTERNAL] AsyncCache: Adds Parameter to Enable Exception Handling Using CosmosClientOptions #5073

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kirankumarkolli merged 12 commits into from
Mar 20, 2025
Merged

[INTERNAL] AsyncCache: Adds Parameter to Enable Exception Handling Using CosmosClientOptions #5073

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
02e2b42
Code changes to hook in cosmos client options to enable exception han…
kundadebdatta Mar 19, 2025
57f8e87
Few misc changes
kirankumarkolli Mar 20, 2025
7905e13
Some more changes
kirankumarkolli Mar 20, 2025
1da325b
One more missing place
kirankumarkolli Mar 20, 2025
92239a7
Brinign it back default argument (tests failing)
kirankumarkolli Mar 20, 2025
240534f
One more defaults back
kirankumarkolli Mar 20, 2025
a9f800c
Brining more defaults
kirankumarkolli Mar 20, 2025
ebab9c6
Fixing the mocks
kirankumarkolli Mar 20, 2025
0808545
Code changes to fix mocked document client.
kundadebdatta Mar 20, 2025
3e9b97b
Code changes to fix test failures in build pipeline.
kundadebdatta Mar 20, 2025
6e744cc
Merge branch 'master' into users/kundadebdatta/openai_asynccache_exce…
kundadebdatta Mar 20, 2025
38f7146
Fixing change feed tests.
kundadebdatta Mar 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ...t.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Mocks/MockDocumentClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ ValueTask<string> ICosmosAuthorizationTokenProvider.GetUserAuthorizationTokenAsy

private void Init()
{
this.collectionCache = new Mock<ClientCollectionCache>(null, new ServerStoreModel(null), null, null, null);
this.collectionCache = new Mock<ClientCollectionCache>(null, new ServerStoreModel(null), null, null, null, false);

ContainerProperties containerProperties = ContainerProperties.CreateWithResourceId("test");
containerProperties.PartitionKey = partitionKeyDefinition;
Expand Down Expand Up @@ -181,7 +181,7 @@ private void Init()
},
string.Empty);

this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null, null, false);
this.partitionKeyRangeCache.Setup(
m => m.TryLookupAsync(
It.IsAny<string>(),
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,7 +799,7 @@ private class ClientWithSplitDetection : MockDocumentClient

public ClientWithSplitDetection()
{
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null, false);
this.partitionKeyRangeCache.Setup(
m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down
2 changes: 1 addition & 1 deletion ...Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,7 @@ private class ClientWithSplitDetection : MockDocumentClient

public ClientWithSplitDetection()
{
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null, false);
this.partitionKeyRangeCache.Setup(
m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down
2 changes: 1 addition & 1 deletion ....Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ private class ClientWithSplitDetection : MockDocumentClient

public ClientWithSplitDetection()
{
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null, false);
this.partitionKeyRangeCache.Setup(
m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down
2 changes: 1 addition & 1 deletion ...re.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BulkPartitionKeyRangeGoneRetryPolicyTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ private class ClientWithSplitDetection : MockDocumentClient

public ClientWithSplitDetection()
{
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null, false);
this.partitionKeyRangeCache.Setup(
m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CancellationTokenTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ private static GatewayStoreModel MockGatewayStoreModel(Func<HttpRequestMessage,
Mock<IDocumentClientInternal> mockDocumentClient = new Mock<IDocumentClientInternal>();
mockDocumentClient.Setup(client => client.ServiceEndpoint).Returns(new Uri("https://foo"));

Mock<GlobalEndpointManager> endpointManager = new Mock<GlobalEndpointManager>(mockDocumentClient.Object, new ConnectionPolicy());
Mock<GlobalEndpointManager> endpointManager = new Mock<GlobalEndpointManager>(mockDocumentClient.Object, new ConnectionPolicy(), false);
endpointManager.Setup(gep => gep.ResolveServiceEndpoint(It.IsAny<DocumentServiceRequest>())).Returns(new Uri("http://localhost"));
ISessionContainer sessionContainer = new SessionContainer(string.Empty);
HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
Expand Down
14 changes: 7 additions & 7 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayAddressCacheTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ public GatewayAddressCacheTests()
}
};

this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null, null);
this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null, null, false);
this.partitionKeyRangeCache
.Setup(m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down Expand Up @@ -636,7 +636,7 @@ public async Task GlobalAddressResolver_OpenConnectionsToAllReplicasAsync_WithVa
containerProperties.Id = "TestId";
containerProperties.PartitionKeyPath = "/pk";

Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict, false);
mockCollectionCahce
.Setup(x => x.ResolveByNameAsync(
It.IsAny<string>(),
Expand Down Expand Up @@ -715,7 +715,7 @@ public async Task GlobalAddressResolver_OpenConnectionsToAllReplicasAsync_WhenHa
containerProperties.Id = "TestId";
containerProperties.PartitionKeyPath = "/pk";

Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict, false);
mockCollectionCahce
.Setup(x => x.ResolveByNameAsync(
It.IsAny<string>(),
Expand Down Expand Up @@ -794,7 +794,7 @@ public async Task GlobalAddressResolver_OpenConnectionsToAllReplicasAsync_WhenIn
containerProperties.Id = "TestId";
containerProperties.PartitionKeyPath = "/pk";

Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict, false);
mockCollectionCahce
.Setup(x => x.ResolveByNameAsync(
It.IsAny<string>(),
Expand All @@ -806,7 +806,7 @@ public async Task GlobalAddressResolver_OpenConnectionsToAllReplicasAsync_WhenIn
.Returns(Task.FromResult(containerProperties));

string exceptionMessage = "Failed to lookup partition key ranges.";
Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new(null, null, null, null);
Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new(null, null, null, null, false);
partitionKeyRangeCache
.Setup(m => m.TryGetOverlappingRangesAsync(
It.IsAny<string>(),
Expand Down Expand Up @@ -882,7 +882,7 @@ public async Task GlobalAddressResolver_OpenConnectionsToAllReplicasAsync_WhenNu
RequestTimeout = TimeSpan.FromSeconds(120)
};

Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict, false);
Comment thread
kirankumarkolli marked this conversation as resolved.
mockCollectionCahce
.Setup(x => x.ResolveByNameAsync(
It.IsAny<string>(),
Expand Down Expand Up @@ -1644,7 +1644,7 @@ public async Task GlobalAddressResolver_TryOpenConnectionToUnhealthyEndpointsAsy
containerProperties.Id = "TestId";
containerProperties.PartitionKeyPath = "/pk";

Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new(MockBehavior.Strict, false);
mockCollectionCahce
.Setup(x => x.ResolveByNameAsync(
It.IsAny<string>(),
Expand Down
48 changes: 24 additions & 24 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,8 +280,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsr,
ConsistencyLevel.Session,
new Mock<ISessionContainer>().Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: Mock.Of<IGlobalEndpointManager>());

Assert.IsNull(dsr.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand All @@ -307,8 +307,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsrQueryPlan,
ConsistencyLevel.Session,
new Mock<ISessionContainer>().Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: Mock.Of<IGlobalEndpointManager>());

Assert.IsNull(dsrQueryPlan.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand Down Expand Up @@ -361,8 +361,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsr,
ConsistencyLevel.Session,
new Mock<ISessionContainer>().Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: Mock.Of<IGlobalEndpointManager>());

Assert.AreEqual(dsrSessionToken, dsr.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand Down Expand Up @@ -391,8 +391,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsrNoSessionToken,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: globalEndpointManager.Object);

if (dsrNoSessionToken.IsReadOnlyRequest || dsrNoSessionToken.OperationType == OperationType.Batch || multiMaster)
Expand Down Expand Up @@ -426,13 +426,13 @@ await GatewayStoreModel.ApplySessionTokenAsync(
containerProperties.Id = "TestId";
containerProperties.PartitionKeyPath = "/pk";

Mock<CollectionCache> mockCollectionCahce = new Mock<CollectionCache>(MockBehavior.Strict);
Mock<CollectionCache> mockCollectionCahce = new Mock<CollectionCache>(MockBehavior.Strict, false);
mockCollectionCahce.Setup(x => x.ResolveCollectionAsync(
dsr,
It.IsAny<CancellationToken>(),
NoOpTrace.Singleton)).Returns(Task.FromResult(containerProperties));

Mock<PartitionKeyRangeCache> mockPartitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
Mock<PartitionKeyRangeCache> mockPartitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null, false);
mockPartitionKeyRangeCache.Setup(x => x.TryGetPartitionKeyRangeByIdAsync(
containerProperties.ResourceId,
partitionKeyRangeId,
Expand Down Expand Up @@ -479,8 +479,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsrSprocExecute,
ConsistencyLevel.Session,
new Mock<ISessionContainer>().Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: Mock.Of<IGlobalEndpointManager>());

Assert.AreEqual(sessionToken, dsrSprocExecute.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand Down Expand Up @@ -518,8 +518,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
dsrNoSessionToken,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: globalEndpointManager.Object);

if (isWriteRequest && multiMaster)
Expand Down Expand Up @@ -968,8 +968,8 @@ public async Task GatewayStoreModel_AvoidGlobalSessionToken()
eventSource,
null,
MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient()));
Mock<ClientCollectionCache> clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null);
Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache.Object, endpointManager);
Mock<ClientCollectionCache> clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null, false);
Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache.Object, endpointManager, false);

sessionContainer.SetSessionToken(
ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString(),
Expand Down Expand Up @@ -1063,9 +1063,9 @@ Task<HttpResponseMessage> sendFunc(HttpRequestMessage request)
null,
MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler)));

Mock<ClientCollectionCache> clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null);
Mock<ClientCollectionCache> clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null, false);

Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache.Object, endpointManager);
Mock<PartitionKeyRangeCache> partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache.Object, endpointManager, false);
storeModel.SetCaches(partitionKeyRangeCache.Object, clientCollectionCache.Object);

INameValueCollection headers = new RequestNameValueCollection();
Expand Down Expand Up @@ -1133,8 +1133,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
documentServiceRequestToChild,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: globalEndpointManager.Object);

Assert.AreEqual($"{childPKRangeId}:{parentSession}", documentServiceRequestToChild.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand Down Expand Up @@ -1199,8 +1199,8 @@ await GatewayStoreModel.ApplySessionTokenAsync(
documentServiceRequestToChild,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null, null).Object,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null, null, false).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null, null, false).Object,
globalEndpointManager: globalEndpointManager.Object);

Assert.AreEqual($"{childPKRangeId}:{tokenWithAllMax}", documentServiceRequestToChild.Headers[HttpConstants.HttpHeaders.SessionToken]);
Expand Down Expand Up @@ -1269,8 +1269,8 @@ static async Task<HttpResponseMessage> messageHandler(HttpRequestMessage request
null,
MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(httpMessageHandler)));

ClientCollectionCache clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null).Object;
PartitionKeyRangeCache partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache, endpointManager).Object;
ClientCollectionCache clientCollectionCache = new Mock<ClientCollectionCache>(new SessionContainer("testhost"), storeModel, null, null, null, false).Object;
PartitionKeyRangeCache partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, storeModel, clientCollectionCache, endpointManager, false).Object;
storeModel.SetCaches(partitionKeyRangeCache, clientCollectionCache);

await executeWithGatewayStoreModel(storeModel);
Expand Down
Loading