diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 3123eed0ba..db8a549387 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -101,10 +101,10 @@ await GatewayStoreModel.ApplySessionTokenAsync( request.RequestContext.RegionName = regionName; } - bool isPPAFEnabled = this.IsPartitionLevelFailoverEnabled(); - // This is applicable for both per partition automatic failover and per partition circuit breaker. - if ((isPPAFEnabled || this.isThinClientEnabled) - && !ReplicatedResourceClient.IsMasterResource(request.ResourceType) + bool isPPAFEnabled = this.IsPartitionLevelFailoverEnabled(); + // This is applicable for both per partition automatic failover and per partition circuit breaker. + if ((isPPAFEnabled || this.isThinClientEnabled) + && !ReplicatedResourceClient.IsMasterResource(request.ResourceType) && (request.ResourceType.IsPartitioned() || request.ResourceType == ResourceType.StoredProcedure)) { (bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync( @@ -563,7 +563,8 @@ internal static bool IsOperationSupportedByThinClient(DocumentServiceRequest req || request.OperationType == OperationType.Upsert || request.OperationType == OperationType.Replace || request.OperationType == OperationType.Delete - || request.OperationType == OperationType.Query)) + || request.OperationType == OperationType.Query + || request.OperationType == OperationType.QueryPlan)) { return true; } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 8f3c7ca366..d210dae18d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -173,8 +173,8 @@ public override async Task> ExecuteItemQueryAsync( resourceType, message, trace); - } - + } + public override async Task ExecuteQueryPlanRequestAsync( string resourceUri, ResourceType resourceType, @@ -209,7 +209,20 @@ public override async Task ExecuteQueryPlanReques { // Syntax exception are argument exceptions and thrown to the user. message.EnsureSuccessStatusCode(); - partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); + + if (this.documentClient.isThinClientEnabled) + { + ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( + resourceUri, trace, cancellationToken); + + partitionedQueryExecutionInfo = ThinClientQueryPlanHelper.DeserializeQueryPlanResponse( + message.Content, + containerProperties.PartitionKey); + } + else + { + partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); + } } return partitionedQueryExecutionInfo; diff --git a/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs new file mode 100644 index 0000000000..6edb6c21a0 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs @@ -0,0 +1,162 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Text.Json; + using Newtonsoft.Json; + using PartitionKeyDefinition = Documents.PartitionKeyDefinition; + using PartitionKeyInternal = Documents.Routing.PartitionKeyInternal; + + /// + /// Handles conversion of thin client query plan responses where query ranges + /// are returned in PartitionKeyInternal format instead of EPK hex strings. + /// Mirrors the conversion logic in . + /// + /// + /// Uses System.Text.Json for primary parsing and structural validation. + /// Newtonsoft.Json is used only for deserializing QueryInfo, HybridSearchQueryInfo, + /// and Range<PartitionKeyInternal> because these types and their deep type hierarchies + /// (including the external Direct package types) use Newtonsoft [JsonProperty] attributes + /// and [JsonObject(MemberSerialization.OptIn)] semantics that have no System.Text.Json equivalent. + /// + internal static class ThinClientQueryPlanHelper + { + private static readonly Newtonsoft.Json.JsonSerializerSettings NewtonsoftSettings = + new Newtonsoft.Json.JsonSerializerSettings + { + DateParseHandling = Newtonsoft.Json.DateParseHandling.None, + MaxDepth = 64, + }; + + /// + /// Deserializes a thin client query plan response stream into a + /// with EPK string ranges. + /// The response contains query ranges in PartitionKeyInternal format + /// which are converted to EPK hex strings and sorted. + /// + /// The response stream containing the raw query plan JSON. + /// The partition key definition for the container. + /// with sorted EPK string ranges. + /// Thrown when or is null. + /// Thrown when the response JSON is malformed or missing required properties. + public static PartitionedQueryExecutionInfo DeserializeQueryPlanResponse( + Stream stream, + PartitionKeyDefinition partitionKeyDefinition) + { + if (stream == null) + { + throw new ArgumentNullException(nameof(stream)); + } + + if (partitionKeyDefinition == null) + { + throw new ArgumentNullException(nameof(partitionKeyDefinition)); + } + + using JsonDocument doc = JsonDocument.Parse(stream); + JsonElement root = doc.RootElement; + + if (root.ValueKind != JsonValueKind.Object) + { + throw new FormatException( + $"Thin client query plan response must be a JSON object, but was {root.ValueKind}."); + } + + // Validate and extract queryRanges (required) + if (!root.TryGetProperty("queryRanges", out JsonElement queryRangesElement)) + { + throw new FormatException( + "Thin client query plan response is missing the required 'queryRanges' property."); + } + + if (queryRangesElement.ValueKind != JsonValueKind.Array) + { + throw new FormatException( + $"Expected 'queryRanges' to be a JSON array, but was {queryRangesElement.ValueKind}."); + } + + if (queryRangesElement.GetArrayLength() == 0) + { + throw new FormatException( + "Thin client query plan response 'queryRanges' array must not be empty."); + } + + // Deserialize QueryInfo using Newtonsoft because QueryInfo uses + // [JsonObject(MemberSerialization.OptIn)] and Newtonsoft-only [JsonProperty] attributes. + QueryInfo queryInfo = null; + if (root.TryGetProperty("queryInfo", out JsonElement queryInfoElement) + && queryInfoElement.ValueKind != JsonValueKind.Null) + { + queryInfo = Newtonsoft.Json.JsonConvert.DeserializeObject( + queryInfoElement.GetRawText(), + ThinClientQueryPlanHelper.NewtonsoftSettings); + } + + // Deserialize HybridSearchQueryInfo using Newtonsoft (same constraint as QueryInfo). + HybridSearchQueryInfo hybridSearchQueryInfo = null; + if (root.TryGetProperty("hybridSearchQueryInfo", out JsonElement hybridElement) + && hybridElement.ValueKind != JsonValueKind.Null) + { + hybridSearchQueryInfo = Newtonsoft.Json.JsonConvert.DeserializeObject( + hybridElement.GetRawText(), + ThinClientQueryPlanHelper.NewtonsoftSettings); + } + + // Parse and convert query ranges to EPK string ranges. + // Range requires Newtonsoft because PartitionKeyInternal + // is from the external Direct package with Newtonsoft-based serialization. + List> effectiveRanges = + new List>(queryRangesElement.GetArrayLength()); + + foreach (JsonElement rangeElement in queryRangesElement.EnumerateArray()) + { + if (rangeElement.ValueKind != JsonValueKind.Object) + { + throw new FormatException( + $"Each query range must be a JSON object, but was {rangeElement.ValueKind}."); + } + + if (!rangeElement.TryGetProperty("min", out _)) + { + throw new FormatException( + "Query range is missing the required 'min' property."); + } + + if (!rangeElement.TryGetProperty("max", out _)) + { + throw new FormatException( + "Query range is missing the required 'max' property."); + } + + Documents.Routing.Range internalRange = + Newtonsoft.Json.JsonConvert.DeserializeObject>( + rangeElement.GetRawText(), + ThinClientQueryPlanHelper.NewtonsoftSettings); + + if (internalRange == null) + { + throw new FormatException( + "Failed to deserialize query range from thin client response."); + } + + effectiveRanges.Add(PartitionKeyInternal.GetEffectivePartitionKeyRange( + partitionKeyDefinition, + internalRange)); + } + + effectiveRanges.Sort(Documents.Routing.Range.MinComparer.Instance); + + return new PartitionedQueryExecutionInfo() + { + QueryInfo = queryInfo, + QueryRanges = effectiveRanges, + HybridSearchQueryInfo = hybridSearchQueryInfo, + }; + } + } +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs index 7cbd053101..d34866a6ca 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs @@ -31,12 +31,12 @@ public class CosmosItemThinClientTests private MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer cosmosSystemTextJsonSerializer; private const int ItemCount = 100; - [TestInitialize] - public async Task TestInitAsync() - { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); + [TestInitialize] + public async Task TestInitAsync() + { + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); this.connectionString = Environment.GetEnvironmentVariable("COSMOSDB_THINCLIENT"); - + if (string.IsNullOrEmpty(this.connectionString)) { Assert.Fail("Set environment variable COSMOSDB_THINCLIENT to run the tests"); @@ -188,183 +188,227 @@ public async Task RegionalDatabaseAccountNameIsEmptyInPayload() [TestCategory("ThinClient")] public async Task TestThinClientWithExecuteStoredProcedureAsync() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); + CosmosClient localClient = null; + Database localDatabase = null; - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + try { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, - }); + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + localClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = localSerializer, + }); + string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - string sprocId = "testSproc_" + Guid.NewGuid().ToString(); - string sprocBody = @"function(itemToCreate) { - var context = getContext(); - var collection = context.getCollection(); - var response = context.getResponse(); + + string sprocId = "testSproc_" + Guid.NewGuid().ToString(); + string sprocBody = @"function(itemToCreate) { + var context = getContext(); + var collection = context.getCollection(); + var response = context.getResponse(); - if (!itemToCreate) throw new Error('Item is undefined or null.'); + if (!itemToCreate) throw new Error('Item is undefined or null.'); - // Create a document - var accepted = collection.createDocument( - collection.getSelfLink(), - itemToCreate, - function(err, newItem) { - if (err) throw err; - - // Query the created document - var query = 'SELECT * FROM c WHERE c.id = ""' + newItem.id + '""'; - var isAccepted = collection.queryDocuments( + // Create a document + var accepted = collection.createDocument( collection.getSelfLink(), - query, - function(queryErr, documents) { - if (queryErr) throw queryErr; - response.setBody({ - created: newItem, - queried: documents[0] - }); - } - ); - if (!isAccepted) throw 'Query not accepted'; - }); + itemToCreate, + function(err, newItem) { + if (err) throw err; + + // Query the created document + var query = 'SELECT * FROM c WHERE c.id = ""' + newItem.id + '""'; + var isAccepted = collection.queryDocuments( + collection.getSelfLink(), + query, + function(queryErr, documents) { + if (queryErr) throw queryErr; + response.setBody({ + created: newItem, + queried: documents[0] + }); + } + ); + if (!isAccepted) throw 'Query not accepted'; + }); - if (!accepted) throw new Error('Create was not accepted.'); - }"; + if (!accepted) throw new Error('Create was not accepted.'); + }"; - // Create stored procedure - Scripts.StoredProcedureResponse createResponse = await this.container.Scripts.CreateStoredProcedureAsync( - new Scripts.StoredProcedureProperties(sprocId, sprocBody)); - Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); + // Create stored procedure + Scripts.StoredProcedureResponse createResponse = await localContainer.Scripts.CreateStoredProcedureAsync( + new Scripts.StoredProcedureProperties(sprocId, sprocBody)); + Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); - // Execute stored procedure - string testPartitionId = Guid.NewGuid().ToString(); - TestObject testItem = new TestObject - { - Id = Guid.NewGuid().ToString(), - Pk = testPartitionId, - Other = "Created by Stored Procedure" - }; + // Execute stored procedure + string testPartitionId = Guid.NewGuid().ToString(); + TestObject testItem = new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = testPartitionId, + Other = "Created by Stored Procedure" + }; - Scripts.StoredProcedureExecuteResponse executeResponse = - await this.container.Scripts.ExecuteStoredProcedureAsync( - sprocId, - new PartitionKey(testPartitionId), - new dynamic[] { testItem }); + Scripts.StoredProcedureExecuteResponse executeResponse = + await localContainer.Scripts.ExecuteStoredProcedureAsync( + sprocId, + new PartitionKey(testPartitionId), + new dynamic[] { testItem }); - Assert.AreEqual(HttpStatusCode.OK, executeResponse.StatusCode); - Assert.IsNotNull(executeResponse.Resource); - string diagnostics = executeResponse.Diagnostics.ToString(); - Assert.IsTrue(diagnostics.Contains("|F4"), "Diagnostics User Agent should contain '|F4' for ThinClient"); + Assert.AreEqual(HttpStatusCode.OK, executeResponse.StatusCode); + Assert.IsNotNull(executeResponse.Resource); + string diagnostics = executeResponse.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), "Diagnostics User Agent should contain '|F4' for ThinClient"); - // Delete stored procedure - await this.container.Scripts.DeleteStoredProcedureAsync(sprocId); + // Delete stored procedure + await localContainer.Scripts.DeleteStoredProcedureAsync(sprocId); + } + finally + { + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } + + if (localClient != null) + { + localClient.Dispose(); + } + } } - [TestMethod] - [TestCategory("ThinClient")] - public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); - - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions - { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, - }); - - string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - - - string sprocId = "testSproc_" + Guid.NewGuid().ToString(); - string sprocBody = @"function(itemToCreate) { - var context = getContext(); - var collection = context.getCollection(); - var response = context.getResponse(); - - if (!itemToCreate) throw new Error('Item is undefined or null.'); - - // Create a document - var accepted = collection.createDocument( - collection.getSelfLink(), - itemToCreate, - function(err, newItem) { - if (err) throw err; - - // Query the created document - var query = 'SELECT * FROM c WHERE c.id = ""' + newItem.id + '""'; - var isAccepted = collection.queryDocuments( - collection.getSelfLink(), - query, - function(queryErr, documents) { - if (queryErr) throw queryErr; - response.setBody({ - created: newItem, - queried: documents[0] - }); - } - ); - if (!isAccepted) throw 'Query not accepted'; - }); - - if (!accepted) throw new Error('Create was not accepted.'); - }"; - - // Create stored procedure - Scripts.StoredProcedureResponse createResponse = await this.container.Scripts.CreateStoredProcedureAsync( - new Scripts.StoredProcedureProperties(sprocId, sprocBody)); - Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); - - // Execute stored procedure - string testPartitionId = Guid.NewGuid().ToString(); - TestObject testItem = new TestObject - { - Id = Guid.NewGuid().ToString(), - Pk = testPartitionId, - Other = "Created by Stored Procedure" - }; - - using (ResponseMessage executeResponse = - await this.container.Scripts.ExecuteStoredProcedureStreamAsync( - sprocId, - new PartitionKey(testPartitionId), - new dynamic[] { testItem })) - { - Assert.AreEqual(HttpStatusCode.OK, executeResponse.StatusCode); - Assert.IsNotNull(executeResponse.Content); - string diagnostics = executeResponse.Diagnostics.ToString(); - Assert.IsTrue(diagnostics.Contains("|F4"), "Diagnostics User Agent should contain '|F4' for ThinClient"); - } - - // Delete stored procedure - await this.container.Scripts.DeleteStoredProcedureAsync(sprocId); + CosmosClient localClient = null; + Database localDatabase = null; + + try + { + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); + + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + + localClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = localSerializer, + }); + + string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + + + string sprocId = "testSproc_" + Guid.NewGuid().ToString(); + string sprocBody = @"function(itemToCreate) { + var context = getContext(); + var collection = context.getCollection(); + var response = context.getResponse(); + + if (!itemToCreate) throw new Error('Item is undefined or null.'); + + // Create a document + var accepted = collection.createDocument( + collection.getSelfLink(), + itemToCreate, + function(err, newItem) { + if (err) throw err; + + // Query the created document + var query = 'SELECT * FROM c WHERE c.id = ""' + newItem.id + '""'; + var isAccepted = collection.queryDocuments( + collection.getSelfLink(), + query, + function(queryErr, documents) { + if (queryErr) throw queryErr; + response.setBody({ + created: newItem, + queried: documents[0] + }); + } + ); + if (!isAccepted) throw 'Query not accepted'; + }); + + if (!accepted) throw new Error('Create was not accepted.'); + }"; + + // Create stored procedure + Scripts.StoredProcedureResponse createResponse = await localContainer.Scripts.CreateStoredProcedureAsync( + new Scripts.StoredProcedureProperties(sprocId, sprocBody)); + Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); + + // Execute stored procedure + string testPartitionId = Guid.NewGuid().ToString(); + TestObject testItem = new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = testPartitionId, + Other = "Created by Stored Procedure" + }; + + using (ResponseMessage executeResponse = + await localContainer.Scripts.ExecuteStoredProcedureStreamAsync( + sprocId, + new PartitionKey(testPartitionId), + new dynamic[] { testItem })) + { + Assert.AreEqual(HttpStatusCode.OK, executeResponse.StatusCode); + Assert.IsNotNull(executeResponse.Content); + string diagnostics = executeResponse.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), "Diagnostics User Agent should contain '|F4' for ThinClient"); + } + + // Delete stored procedure + await localContainer.Scripts.DeleteStoredProcedureAsync(sprocId); + } + finally + { + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } + + if (localClient != null) + { + localClient.Dispose(); + } + } } [TestMethod] @@ -426,42 +470,65 @@ public async Task CreateItemsTest() [TestCategory("ThinClient")] public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); - string authKey = Utils.ConfigurationManager.AppSettings["MasterKey"]; - string endpoint = Utils.ConfigurationManager.AppSettings["GatewayEndpoint"]; - AzureKeyCredential masterKeyCredential = new AzureKeyCredential(authKey); + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + try { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - - this.client = new CosmosClient( - endpoint, - masterKeyCredential, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, - }); + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); + string authKey = Utils.ConfigurationManager.AppSettings["MasterKey"]; + string endpoint = Utils.ConfigurationManager.AppSettings["GatewayEndpoint"]; + AzureKeyCredential masterKeyCredential = new AzureKeyCredential(authKey); - string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - - string pk = "pk_create"; - IEnumerable items = this.GenerateItems(pk); - - foreach (TestObject item in items) + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + + localClient = new CosmosClient( + endpoint, + masterKeyCredential, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = localSerializer, + }); + + string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + + string pk = "pk_create"; + IEnumerable items = this.GenerateItems(pk); + + foreach (TestObject item in items) + { + ItemResponse response = await localContainer.CreateItemAsync(item, new PartitionKey(item.Pk)); + Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); + string diagnostics = response.Diagnostics.ToString(); + Assert.IsFalse(diagnostics.Contains("|F4"), "Diagnostics User Agent should NOT contain '|F4' for Gateway"); + } + } + finally { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); - Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); - string diagnostics = response.Diagnostics.ToString(); - Assert.IsFalse(diagnostics.Contains("|F4"), "Diagnostics User Agent should NOT contain '|F4' for Gateway"); + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -469,42 +536,65 @@ public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() [TestCategory("ThinClient")] public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountEnabled() { - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; + + try { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Direct, - Serializer = this.cosmosSystemTextJsonSerializer, - }); + localClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Direct, + Serializer = localSerializer, + }); - string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - string pk = "pk_create"; - IEnumerable items = this.GenerateItems(pk); + string pk = "pk_create"; + IEnumerable items = this.GenerateItems(pk); - foreach (TestObject item in items) + foreach (TestObject item in items) + { + ItemResponse response = await localContainer.CreateItemAsync(item, new PartitionKey(item.Pk)); + Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); + JsonDocument doc = JsonDocument.Parse(response.Diagnostics.ToString()); + string connectionMode = doc.RootElement + .GetProperty("data") + .GetProperty("Client Configuration") + .GetProperty("ConnectionMode") + .GetString(); + + Assert.AreEqual("Direct", connectionMode, "Diagnostics should have ConnectionMode set to 'Direct'"); + } + } + finally { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); - Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); - JsonDocument doc = JsonDocument.Parse(response.Diagnostics.ToString()); - string connectionMode = doc.RootElement - .GetProperty("data") - .GetProperty("Client Configuration") - .GetProperty("ConnectionMode") - .GetString(); - - Assert.AreEqual("Direct", connectionMode, "Diagnostics should have ConnectionMode set to 'Direct'"); + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -512,38 +602,61 @@ public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountE [TestCategory("ThinClient")] public async Task CreateItemsTestWithThinClientFlagDisabledAccountEnabled() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "False"); + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + try { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "False"); - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, - }); + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - string uniqueDbName = "TestDbTCDisabled_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestContainerTCDisabled_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + localClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = localSerializer, + }); - string pk = "pk_create"; - IEnumerable items = this.GenerateItems(pk); + string uniqueDbName = "TestDbTCDisabled_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestContainerTCDisabled_" + Guid.NewGuid().ToString(); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - foreach (TestObject item in items) + string pk = "pk_create"; + IEnumerable items = this.GenerateItems(pk); + + foreach (TestObject item in items) + { + ItemResponse response = await localContainer.CreateItemAsync(item, new PartitionKey(item.Pk)); + Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); + string diagnostics = response.Diagnostics.ToString(); + Assert.IsFalse(diagnostics.Contains("|F4"), "Diagnostics User Agent should NOT contain '|F4' for Gateway"); + } + } + finally { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); - Assert.AreEqual(HttpStatusCode.Created, response.StatusCode); - string diagnostics = response.Diagnostics.ToString(); - Assert.IsFalse(diagnostics.Contains("|F4"), "Diagnostics User Agent should NOT contain '|F4' for Gateway"); + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -764,77 +877,150 @@ public async Task QueryItemsTest() [TestCategory("ThinClient")] public async Task QueryItemsTestWithStrongConsistency() { - string connectionString = ConfigurationManager.GetEnvironmentVariable("COSMOSDB_THINCLIENTSTRONG", string.Empty); - if (string.IsNullOrEmpty(connectionString)) + CosmosClient localClient = null; + Database localDatabase = null; + + try { - Assert.Fail("Set environment variable COSMOSDB_THINCLIENTSTRONG to run the tests"); - } - this.client = new CosmosClient( - connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - RequestTimeout = TimeSpan.FromSeconds(60), - ConsistencyLevel = Microsoft.Azure.Cosmos.ConsistencyLevel.Strong - }); - - string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + string connectionString = ConfigurationManager.GetEnvironmentVariable("COSMOSDB_THINCLIENTSTRONG", string.Empty); + if (string.IsNullOrEmpty(connectionString)) + { + Assert.Fail("Set environment variable COSMOSDB_THINCLIENTSTRONG to run the tests"); + } - string pk = "pk_query"; - List items = this.GenerateItems(pk).ToList(); + localClient = new CosmosClient( + connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + RequestTimeout = TimeSpan.FromSeconds(60), + ConsistencyLevel = Microsoft.Azure.Cosmos.ConsistencyLevel.Strong + }); + + string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + + string pk = "pk_query"; + List items = this.GenerateItems(pk).ToList(); + + List itemsCreated = new List(); + foreach (TestObject item in items) + { + try + { + ItemResponse response = await localContainer.CreateItemAsync(item, new PartitionKey(item.Pk)); + if (response.StatusCode == HttpStatusCode.Created) + { + itemsCreated.Add(item); + } + } + catch (CosmosException) + { + } + } - List createdItems = await this.CreateItemsSafeAsync(items); + string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; + FeedIterator iterator = localContainer.GetItemQueryIterator(query); - string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + int count = 0; + while (iterator.HasMoreResults) + { + FeedResponse response = await iterator.ReadNextAsync(); + count += response.Count; + } - int count = 0; - while (iterator.HasMoreResults) - { - FeedResponse response = await iterator.ReadNextAsync(); - count += response.Count; + Assert.AreEqual(itemsCreated.Count, count); } + finally + { + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } - Assert.AreEqual(createdItems.Count, count); + if (localClient != null) + { + localClient.Dispose(); + } + } } [TestMethod] [TestCategory("ThinClient")] public async Task QueryItemsTestWithSessionConsistency() { - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Gateway, - RequestTimeout = TimeSpan.FromSeconds(60), - ConsistencyLevel = Microsoft.Azure.Cosmos.ConsistencyLevel.Session - }); - - string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + CosmosClient localClient = null; + Database localDatabase = null; - string pk = "pk_query"; - List items = this.GenerateItems(pk).ToList(); + try + { + localClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + RequestTimeout = TimeSpan.FromSeconds(60), + ConsistencyLevel = Microsoft.Azure.Cosmos.ConsistencyLevel.Session + }); + + string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); + string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + + string pk = "pk_query"; + List items = this.GenerateItems(pk).ToList(); + + List itemsCreated = new List(); + foreach (TestObject item in items) + { + try + { + ItemResponse response = await localContainer.CreateItemAsync(item, new PartitionKey(item.Pk)); + if (response.StatusCode == HttpStatusCode.Created) + { + itemsCreated.Add(item); + } + } + catch (CosmosException) + { + } + } - List createdItems = await this.CreateItemsSafeAsync(items); + string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; + FeedIterator iterator = localContainer.GetItemQueryIterator(query); - string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + int count = 0; + while (iterator.HasMoreResults) + { + FeedResponse response = await iterator.ReadNextAsync(); + count += response.Count; + } - int count = 0; - while (iterator.HasMoreResults) - { - FeedResponse response = await iterator.ReadNextAsync(); - count += response.Count; + Assert.AreEqual(itemsCreated.Count, count); } + finally + { + if (localDatabase != null) + { + try + { + await localDatabase.DeleteAsync(); + } + catch { } + } - Assert.AreEqual(createdItems.Count, count); + if (localClient != null) + { + localClient.Dispose(); + } + } } [TestMethod] @@ -927,87 +1113,247 @@ public async Task TransactionalBatchCreateItemsTest() } } - [TestMethod] - [TestCategory("ThinClient")] - public async Task RegionalFailoverWithHttpRequestException_EnsuresThinClientHeaderInRefreshRequest() - { - // Arrange - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); - - bool headerFoundInRefreshRequest = false; - int accountRefreshCount = 0; - bool hasThrown = false; - - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions - { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - CosmosSystemTextJsonSerializer serializer = new CosmosSystemTextJsonSerializer(jsonSerializerOptions); - - FaultInjectionDelegatingHandler faultHandler = new FaultInjectionDelegatingHandler( - (request) => - { - // Check for account refresh requests (GET to "/" with HTTP/1.1) - if (request.Method == HttpMethod.Get && - request.RequestUri.AbsolutePath == "/" && - request.Version == new Version(1, 1)) - { - accountRefreshCount++; - - // Only check header after we've thrown the exception - if (hasThrown) - { - if (request.Headers.TryGetValues( - ThinClientConstants.EnableThinClientEndpointDiscoveryHeaderName, - out IEnumerable headerValues)) - { - if (headerValues.Contains("True")) - { - headerFoundInRefreshRequest = true; - } - } - } - } - - // Throw HttpRequestException only ONCE on ThinClient POST requests - if (!hasThrown && - request.Method == HttpMethod.Post && - request.Version == new Version(2, 0)) - { - hasThrown = true; - throw new HttpRequestException("Simulated endpoint failure"); - } - }); - - CosmosClientBuilder builder = new CosmosClientBuilder(this.connectionString) - .WithConnectionModeGateway() - .WithCustomSerializer(serializer) - .WithHttpClientFactory(() => new HttpClient(faultHandler)); - - using CosmosClient client = builder.Build(); - - string uniqueDbName = "TestFailoverDb_" + Guid.NewGuid().ToString(); - Database database = await client.CreateDatabaseIfNotExistsAsync(uniqueDbName); - string uniqueContainerName = "TestFailoverContainer_" + Guid.NewGuid().ToString(); - Container container = await database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); - - string pk = "pk_failover_test"; - TestObject testItem = this.GenerateItems(pk).First(); - - // Act - CreateItemAsync will fail once, then SDK retries and succeeds - ItemResponse response = await container.CreateItemAsync(testItem, new PartitionKey(testItem.Pk)); - - // Assert - Assert.AreEqual(HttpStatusCode.Created, response.StatusCode, "Request should succeed after retry"); - Assert.IsTrue(hasThrown, "Exception should have been thrown once"); - Assert.IsTrue(headerFoundInRefreshRequest, "Account refresh after HttpRequestException should contain thin client header"); - - // Cleanup - await database.DeleteAsync(); - } - + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientQueryPlanWithOrderBy() + { + List items = new List(); + string commonPk = "pk_orderby_test_" + Guid.NewGuid().ToString(); + + try + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, Boolean.TrueString); + + for (int i = 0; i < 5; i++) + { + items.Add(new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = commonPk, + Other = $"Item_{i:D3}", + }); + } + + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(5, createdItems.Count, "All items should be created"); + + // Execute ORDER BY query - this requires QueryPlan and EPK range conversion + string query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.other DESC"; + QueryDefinition queryDef = new QueryDefinition(query).WithParameter("@pk", commonPk); + + FeedIterator iterator = this.container.GetItemQueryIterator(queryDef); + + List results = new List(); + int pageCount = 0; + + while (iterator.HasMoreResults) + { + FeedResponse response = await iterator.ReadNextAsync(); + results.AddRange(response); + pageCount++; + + string diagnostics = response.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), $"Page {pageCount}: Should use ThinClient"); + } + + Assert.AreEqual(5, results.Count, "Should return all 5 items"); + + } + finally + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, null); + + foreach (TestObject item in items) + { + try + { + await this.container.DeleteItemAsync(item.Id, new PartitionKey(item.Pk)); + } + catch { } + } + } + } + + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientQueryPlanCrossPartitionWithFilter() + { + List items = new List(); + string baseGuid = Guid.NewGuid().ToString(); + + try + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, "True"); + + string[] partitionKeys = { + $"pk_filter_1_{baseGuid}", + $"pk_filter_2_{baseGuid}", + $"pk_filter_3_{baseGuid}" + }; + + for (int pkIndex = 0; pkIndex < partitionKeys.Length; pkIndex++) + { + for (int i = 0; i < 3; i++) + { + items.Add(new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = partitionKeys[pkIndex], + Other = $"Value_{i}", + }); + } + } + + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(9, createdItems.Count, "All 9 items should be created"); + + string query = "SELECT * FROM c ORDER BY c._ts"; + + FeedIterator iterator = this.container.GetItemQueryIterator(query); + + List results = new List(); + int pageCount = 0; + + while (iterator.HasMoreResults) + { + FeedResponse response = await iterator.ReadNextAsync(); + results.AddRange(response); + pageCount++; + + string diagnostics = response.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), $"Page {pageCount}: Should use ThinClient"); + } + + Assert.IsTrue(results.Count >= 9, + $"Should return at least 9 items, got {results.Count}"); + + int foundCount = 0; + foreach (TestObject item in createdItems) + { + if (results.Any(r => r.Id == item.Id)) + { + foundCount++; + } + } + + Assert.IsTrue(foundCount >= 9, + $"Should find all 9 test items in results, found {foundCount}"); + + } + finally + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, null); + + foreach (TestObject item in items) + { + try + { + await this.container.DeleteItemAsync(item.Id, new PartitionKey(item.Pk)); + } + catch { } + } + } + } + + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientQueryPlanMultiPartitionFanout() + { + List items = new List(); + string baseGuid = Guid.NewGuid().ToString(); + + try + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, Boolean.TrueString); + + // Create items across many distinct partition keys to ensure multi-partition fanout + int partitionCount = 10; + int itemsPerPartition = 3; + + for (int pkIndex = 0; pkIndex < partitionCount; pkIndex++) + { + string pk = $"pk_fanout_{pkIndex}_{baseGuid}"; + for (int i = 0; i < itemsPerPartition; i++) + { + items.Add(new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = pk, + Other = $"Partition_{pkIndex}_Item_{i}", + }); + } + } + + int totalExpected = partitionCount * itemsPerPartition; + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(totalExpected, createdItems.Count, $"All {totalExpected} items should be created"); + + // Execute a cross-partition ORDER BY query (requires QueryPlan + fanout) + string query = "SELECT * FROM c WHERE STARTSWITH(c.other, 'Partition_') ORDER BY c.other ASC"; + + // Run query via ThinClient mode + FeedIterator thinClientIterator = this.container.GetItemQueryIterator(query); + + List thinClientResults = new List(); + while (thinClientIterator.HasMoreResults) + { + FeedResponse response = await thinClientIterator.ReadNextAsync(); + thinClientResults.AddRange(response); + + string diagnostics = response.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), "Should use ThinClient mode"); + } + + // Verify all items are returned + int foundCount = createdItems.Count(created => + thinClientResults.Any(r => r.Id == created.Id)); + Assert.AreEqual(totalExpected, foundCount, + $"Should find all {totalExpected} test items in fanout results, found {foundCount}"); + + // Compare with Gateway mode results to verify correctness + using CosmosClient gatewayClient = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = this.cosmosSystemTextJsonSerializer, + }); + + Container gatewayContainer = gatewayClient.GetContainer(this.database.Id, this.container.Id); + FeedIterator gatewayIterator = gatewayContainer.GetItemQueryIterator(query); + + List gatewayResults = new List(); + while (gatewayIterator.HasMoreResults) + { + FeedResponse response = await gatewayIterator.ReadNextAsync(); + gatewayResults.AddRange(response); + } + + // ThinClient and Gateway should return the same item count + Assert.AreEqual(gatewayResults.Count, thinClientResults.Count, + $"ThinClient ({thinClientResults.Count}) and Gateway ({gatewayResults.Count}) should return the same number of items."); + + // Verify both results contain the same item IDs + HashSet thinClientIds = new HashSet(thinClientResults.Select(r => r.Id)); + HashSet gatewayIds = new HashSet(gatewayResults.Select(r => r.Id)); + Assert.IsTrue(thinClientIds.SetEquals(gatewayIds), + "ThinClient and Gateway should return the same set of items."); + } + finally + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, null); + + foreach (TestObject item in items) + { + try + { + await this.container.DeleteItemAsync(item.Id, new PartitionKey(item.Pk)); + } + catch { } + } + } + } + /// /// DelegatingHandler that intercepts HTTP requests and can inject faults /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index ba3eaaeb81..38ba61c481 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -1333,7 +1333,7 @@ public async Task ThinClient_ProcessMessageAsync_WithUnsupportedOperations_Shoul .ReturnsAsync(successResponse); DocumentServiceRequest request = DocumentServiceRequest.Create( - operationType: OperationType.QueryPlan, + operationType: OperationType.ReadFeed, resourceType: ResourceType.Document, resourceId: "NH1uAJ6ANm0=", body: null, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientStoreClientTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientStoreClientTests.cs index 7e400450a8..d97bb331ae 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientStoreClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientStoreClientTests.cs @@ -4,21 +4,25 @@ namespace Microsoft.Azure.Cosmos.Tests { - using System; - using System.Collections.ObjectModel; - using System.IO; - using System.Linq; + using System; + using System.Collections.Generic; + using System.Collections.ObjectModel; + using System.IO; + using System.Linq; using System.Net; - using System.Net.Http; + using System.Net.Http; using System.Text; using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Routing; - using Microsoft.Azure.Cosmos.Telemetry; - using Microsoft.Azure.Cosmos.Tracing; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.Telemetry; + using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; - using Microsoft.VisualStudio.TestTools.UnitTesting; - using Moq; + using Microsoft.Azure.Documents.Routing; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + using Newtonsoft.Json; [TestClass] public class ThinClientStoreClientTests @@ -391,6 +395,133 @@ public void Constructor_ShouldThrowArgumentNullException_WhenUserAgentContainerI StringAssert.Contains(ex.Message, "UserAgentContainer cannot be null"); } + #region ThinClientQueryPlanHelper Tests + + private static readonly PartitionKeyDefinition HashPartitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new Collection() { "/id" }, + Kind = PartitionKind.Hash, + }; + + [TestMethod] + [DynamicData(nameof(GetQueryPlanJsonTestCases), DynamicDataSourceType.Method)] + public void DeserializeQueryPlanResponse_ConsistentWithQueryPartitionProvider(string queryPlanJson, string description) + { + // Deserialize via ThinClientQueryPlanHelper (stream-based, as used in thin client mode) + PartitionedQueryExecutionInfo thinClientResult; + using (Stream stream = new MemoryStream(Encoding.UTF8.GetBytes(queryPlanJson))) + { + thinClientResult = ThinClientQueryPlanHelper.DeserializeQueryPlanResponse( + stream, + HashPartitionKeyDefinition); + } + + // Deserialize via QueryPartitionProvider (string-based, as used in gateway/service-interop mode) + QueryPartitionProvider queryPartitionProvider = new QueryPartitionProvider( + new Dictionary() { { "maxSqlQueryInputLength", 524288 } }); + + PartitionedQueryExecutionInfoInternal queryInfoInternal = + JsonConvert.DeserializeObject( + queryPlanJson, + new JsonSerializerSettings { DateParseHandling = DateParseHandling.None, MaxDepth = 64 }); + + PartitionedQueryExecutionInfo providerResult = queryPartitionProvider.ConvertPartitionedQueryExecutionInfo( + queryInfoInternal, + HashPartitionKeyDefinition); + + // Assert: Both paths must produce identical EPK ranges + Assert.AreEqual(providerResult.QueryRanges.Count, thinClientResult.QueryRanges.Count, description); + for (int i = 0; i < providerResult.QueryRanges.Count; i++) + { + Assert.AreEqual(providerResult.QueryRanges[i].Min, thinClientResult.QueryRanges[i].Min, $"{description} - range[{i}].Min"); + Assert.AreEqual(providerResult.QueryRanges[i].Max, thinClientResult.QueryRanges[i].Max, $"{description} - range[{i}].Max"); + Assert.AreEqual(providerResult.QueryRanges[i].IsMinInclusive, thinClientResult.QueryRanges[i].IsMinInclusive, $"{description} - range[{i}].IsMinInclusive"); + Assert.AreEqual(providerResult.QueryRanges[i].IsMaxInclusive, thinClientResult.QueryRanges[i].IsMaxInclusive, $"{description} - range[{i}].IsMaxInclusive"); + } + } + + private static IEnumerable GetQueryPlanJsonTestCases() + { + // Full range (cross-partition query) + yield return new object[] + { + @"{""queryInfo"":{""distinctType"":""None"",""top"":null,""offset"":null,""limit"":null,""orderBy"":[],""orderByExpressions"":[],""groupByExpressions"":[],""groupByAliases"":[],""aggregates"":[""CountIf""],""groupByAliasToAggregateType"":{},""rewrittenQuery"":""SELECT VALUE [{\""item\"": COUNTIF(c.valid)}]\nFROM c"",""hasSelectValue"":true,""dCountInfo"":null,""hasNonStreamingOrderBy"":false},""queryRanges"":[{""min"":[],""max"":""Infinity"",""isMinInclusive"":true,""isMaxInclusive"":false}]}", + "Full range with aggregate" + }; + + // Point query (single partition key) + yield return new object[] + { + @"{""queryInfo"":{""distinctType"":""None"",""top"":null,""offset"":null,""limit"":null,""orderBy"":[""Descending""],""orderByExpressions"":[],""groupByExpressions"":[],""groupByAliases"":[],""aggregates"":[],""groupByAliasToAggregateType"":{},""rewrittenQuery"":"""",""hasSelectValue"":false,""dCountInfo"":null,""hasNonStreamingOrderBy"":false},""queryRanges"":[{""min"":[""testValue""],""max"":[""testValue""],""isMinInclusive"":true,""isMaxInclusive"":true}]}", + "Point query with ORDER BY" + }; + + // HybridSearchQueryInfo + yield return new object[] + { + @"{""hybridSearchQueryInfo"":{""globalStatisticsQuery"":""SELECT COUNT(1) AS documentCount, [] AS fullTextStatistics\nFROM c"",""componentQueryInfos"":[],""componentWithoutPayloadQueryInfos"":[],""projectionQueryInfo"":null,""componentWeights"":null,""skip"":null,""take"":10,""requiresGlobalStatistics"":false},""queryRanges"":[{""min"":[],""max"":""Infinity"",""isMinInclusive"":true,""isMaxInclusive"":false}]}", + "HybridSearchQueryInfo" + }; + } + + [TestMethod] + public void DeserializeQueryPlanResponse_MultipleRanges_SortsOutput() + { + // Multiple ranges in deliberate reverse order to verify sorting + string queryPlanJson = @"{""queryInfo"":{""distinctType"":""None"",""top"":null,""offset"":null,""limit"":null,""orderBy"":[],""orderByExpressions"":[],""groupByExpressions"":[],""groupByAliases"":[],""aggregates"":[],""groupByAliasToAggregateType"":{},""rewrittenQuery"":"""",""hasSelectValue"":false,""dCountInfo"":null,""hasNonStreamingOrderBy"":false},""queryRanges"":[{""min"":[""zzz""],""max"":[""zzz""],""isMinInclusive"":true,""isMaxInclusive"":true},{""min"":[""aaa""],""max"":[""aaa""],""isMinInclusive"":true,""isMaxInclusive"":true},{""min"":[""mmm""],""max"":[""mmm""],""isMinInclusive"":true,""isMaxInclusive"":true}]}"; + + using Stream stream = new MemoryStream(Encoding.UTF8.GetBytes(queryPlanJson)); + + PartitionedQueryExecutionInfo result = ThinClientQueryPlanHelper.DeserializeQueryPlanResponse( + stream, + HashPartitionKeyDefinition); + + Assert.AreEqual(3, result.QueryRanges.Count); + for (int i = 0; i < result.QueryRanges.Count - 1; i++) + { + Assert.IsTrue( + string.Compare(result.QueryRanges[i].Min, result.QueryRanges[i + 1].Min, StringComparison.Ordinal) <= 0, + $"Ranges should be sorted: range[{i}].Min='{result.QueryRanges[i].Min}' should be <= range[{i + 1}].Min='{result.QueryRanges[i + 1].Min}'"); + } + } + + [TestMethod] + public void DeserializeQueryPlanResponse_InvalidInputs_FailsFast() + { + Assert.ThrowsException( + () => ThinClientQueryPlanHelper.DeserializeQueryPlanResponse(null, HashPartitionKeyDefinition), + "Null stream should throw"); + + using (Stream validStream = new MemoryStream(Encoding.UTF8.GetBytes("{}"))) + { + Assert.ThrowsException( + () => ThinClientQueryPlanHelper.DeserializeQueryPlanResponse(validStream, null), + "Null partitionKeyDefinition should throw"); + } + + using (Stream badJson = new MemoryStream(Encoding.UTF8.GetBytes("not valid json {{{"))) + { + try + { + ThinClientQueryPlanHelper.DeserializeQueryPlanResponse(badJson, HashPartitionKeyDefinition); + Assert.Fail("Malformed JSON should throw"); + } + catch (System.Text.Json.JsonException) + { + // Expected - System.Text.Json throws JsonException or a derived type for malformed JSON + } + } + + using (Stream nullJson = new MemoryStream(Encoding.UTF8.GetBytes("null"))) + { + Assert.ThrowsException( + () => ThinClientQueryPlanHelper.DeserializeQueryPlanResponse(nullJson, HashPartitionKeyDefinition), + "JSON null should throw FormatException"); + } + } + + #endregion + private ContainerProperties GetMockContainerProperties() { ContainerProperties containerProperties = new ContainerProperties