From 96413b74d77f39f5135384bc79b103a4c066b0bc Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 12 Feb 2026 20:43:15 -0800 Subject: [PATCH 01/10] Add QueryPlan support --- .../src/GatewayStoreModel.cs | 13 +- .../PartitionedQueryExecutionInfo.cs | 138 ++++++++++++++++-- .../Query/v3Query/CosmosQueryClientCore.cs | 97 +++++++----- 3 files changed, 194 insertions(+), 54 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 4c18b4a937..5244e1fea2 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -101,11 +101,15 @@ await GatewayStoreModel.ApplySessionTokenAsync( request.RequestContext.RegionName = regionName; } + bool isQueryPlanInThinClientMode = this.isThinClientEnabled + && request.OperationType == OperationType.QueryPlan + && request.ResourceType == ResourceType.Document; + 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)) + if ((isPPAFEnabled || this.isThinClientEnabled) + && !ReplicatedResourceClient.IsMasterResource(request.ResourceType) + && (request.ResourceType.IsPartitioned() || request.ResourceType == ResourceType.StoredProcedure || isQueryPlanInThinClientMode)) { (bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync( request: request, @@ -562,7 +566,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/Core/QueryPlan/PartitionedQueryExecutionInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs index 14ae921d8f..39bd0e0bf1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs @@ -7,10 +7,15 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan using System; using System.Collections.Generic; using Newtonsoft.Json; + using Newtonsoft.Json.Linq; using Constants = Documents.Constants; + using PartitionKeyDefinition = Documents.PartitionKeyDefinition; + using PartitionKeyInternal = Documents.Routing.PartitionKeyInternal; internal sealed class PartitionedQueryExecutionInfo { + private List> queryRanges; + public PartitionedQueryExecutionInfo() { this.Version = Constants.PartitionedQueryExecutionInfo.CurrentVersion; @@ -30,22 +35,72 @@ public QueryInfo QueryInfo set; } + /// + /// Gets or sets the query ranges. In thin client mode, this property + /// lazily converts PartitionKeyInternal ranges to EPK hex string ranges. + /// + [JsonIgnore] + public List> QueryRanges + { + get + { + if (this.queryRanges != null) + { + return this.queryRanges; + } + + if (this.RawQueryRanges != null) + { + if (this.UseThinClientMode && this.PartitionKeyDefinition != null) + { + this.queryRanges = this.ParseQueryRangesForThinClient(); + } + else + { + // Non-thin client: deserialize directly as string ranges + this.queryRanges = this.RawQueryRanges.ToObject>>(); + } + } + + return this.queryRanges; + } + set => this.queryRanges = value; + } + + /// + /// Raw query ranges from JSON deserialization. Used for thin client mode parsing. + /// In non-thin client mode, this is deserialized directly to QueryRanges. + /// [JsonProperty(Constants.Properties.QueryRanges)] - public List> QueryRanges + internal JArray RawQueryRanges { get; set; - } - - // Change to the below after Direct package upgrade - // [JsonProperty(Constants.Properties.HybridSearchQueryInfo)] - [JsonProperty("hybridSearchQueryInfo")] - public HybridSearchQueryInfo HybridSearchQueryInfo - { - get; - set; } + // Change to the below after Direct package upgrade + // [JsonProperty(Constants.Properties.HybridSearchQueryInfo)] + [JsonProperty("hybridSearchQueryInfo")] + public HybridSearchQueryInfo HybridSearchQueryInfo + { + get; + set; + } + + /// + /// Flag indicating if thin client mode is enabled. + /// Must be set before accessing QueryRanges property. + /// + [JsonIgnore] + internal bool UseThinClientMode { get; set; } + + /// + /// Partition key definition used for converting PartitionKeyInternal to EPK strings. + /// Must be set before accessing QueryRanges property in thin client mode. + /// + [JsonIgnore] + internal PartitionKeyDefinition PartitionKeyDefinition { get; set; } + public override string ToString() { return JsonConvert.SerializeObject(this); @@ -69,5 +124,68 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec return false; } } + + /// + /// Parses query ranges for thin client mode where the proxy returns ranges + /// in PartitionKeyInternal format (e.g., {"min": [[""]], "max": [["Infinity"]]}) + /// and converts them to EPK hex string ranges. + /// + private List> ParseQueryRangesForThinClient() + { + if (this.RawQueryRanges == null || this.PartitionKeyDefinition == null) + { + return null; + } + + List> epkRanges = new List>(this.RawQueryRanges.Count); + + foreach (JToken rangeToken in this.RawQueryRanges) + { + if (!(rangeToken is JObject rangeObject)) + { + continue; + } + + // Parse min and max as PartitionKeyInternal + JToken minToken = rangeObject["min"]; + JToken maxToken = rangeObject["max"]; + + PartitionKeyInternal minPk = this.ParsePartitionKeyInternal(minToken); + PartitionKeyInternal maxPk = this.ParsePartitionKeyInternal(maxToken); + + // Convert to EPK hex strings + string minEpk = minPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); + string maxEpk = maxPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); + + // Parse isMinInclusive and isMaxInclusive (defaults: min=true, max=false) + bool isMinInclusive = rangeObject["isMinInclusive"]?.Value() ?? true; + bool isMaxInclusive = rangeObject["isMaxInclusive"]?.Value() ?? false; + + epkRanges.Add(new Documents.Routing.Range(minEpk, maxEpk, isMinInclusive, isMaxInclusive)); + } + + return epkRanges; + } + + /// + /// Parses a JSON token representing a PartitionKeyInternal. + /// Handles formats like [[""]] (empty), [["Infinity"]] (infinity), or actual partition key values. + /// + private PartitionKeyInternal ParsePartitionKeyInternal(JToken token) + { + if (token == null || token.Type == JTokenType.Null) + { + return PartitionKeyInternal.Empty; + } + + try + { + return token.ToObject(); + } + catch (JsonException) + { + return PartitionKeyInternal.Empty; + } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 8f3c7ca366..d0a61c9f33 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -173,46 +173,63 @@ public override async Task> ExecuteItemQueryAsync( resourceType, message, trace); - } - - public override async Task ExecuteQueryPlanRequestAsync( - string resourceUri, - ResourceType resourceType, - OperationType operationType, - SqlQuerySpec sqlQuerySpec, - PartitionKey? partitionKey, - string supportedQueryFeatures, - Guid clientQueryCorrelationId, - ITrace trace, - CancellationToken cancellationToken) - { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; - using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( - resourceUri: resourceUri, - resourceType: resourceType, - operationType: operationType, - requestOptions: null, - feedRange: partitionKey.HasValue ? new FeedRangePartitionKey(partitionKey.Value) : null, - cosmosContainerCore: this.cosmosContainerCore, - streamPayload: this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType), - requestEnricher: (requestMessage) => - { - requestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, RuntimeConstants.MediaTypes.QueryJson); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQueryPlanRequest, bool.TrueString); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.SupportedQueryFeatures, supportedQueryFeatures); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.QueryVersion, new Version(major: 1, minor: 0).ToString()); - requestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, clientQueryCorrelationId.ToString()); - requestMessage.UseGatewayMode = true; - }, - trace: trace, - cancellationToken: cancellationToken)) - { - // Syntax exception are argument exceptions and thrown to the user. - message.EnsureSuccessStatusCode(); - partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); - } - - return partitionedQueryExecutionInfo; + } + + public override async Task ExecuteQueryPlanRequestAsync( + string resourceUri, + ResourceType resourceType, + OperationType operationType, + SqlQuerySpec sqlQuerySpec, + PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + ITrace trace, + CancellationToken cancellationToken) + { + // Determine if thin client mode is enabled + bool isThinClientEnabled = this.clientContext.ClientOptions.ConnectionMode == ConnectionMode.Gateway + && ConfigurationManager.GetEnvironmentVariable( + ConfigurationManager.ThinClientModeEnabled, + defaultValue: false); + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( + resourceUri: resourceUri, + resourceType: resourceType, + operationType: operationType, + requestOptions: null, + feedRange: partitionKey.HasValue ? new FeedRangePartitionKey(partitionKey.Value) : null, + cosmosContainerCore: this.cosmosContainerCore, + streamPayload: this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType), + requestEnricher: (requestMessage) => + { + requestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, RuntimeConstants.MediaTypes.QueryJson); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQueryPlanRequest, bool.TrueString); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.SupportedQueryFeatures, supportedQueryFeatures); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.QueryVersion, new Version(major: 1, minor: 0).ToString()); + requestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, clientQueryCorrelationId.ToString()); + if (!isThinClientEnabled) + { + requestMessage.UseGatewayMode = true; + } + }, + trace: trace, + cancellationToken: cancellationToken)) + { + // Syntax exception are argument exceptions and thrown to the user. + message.EnsureSuccessStatusCode(); + partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); + + if (isThinClientEnabled) + { + partitionedQueryExecutionInfo.UseThinClientMode = true; + ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( + resourceUri, trace, cancellationToken); + partitionedQueryExecutionInfo.PartitionKeyDefinition = containerProperties.PartitionKey; + } + } + + return partitionedQueryExecutionInfo; } public override async Task GetClientDisableOptimisticDirectExecutionAsync() From 25cb4bb423ae1f980c29a59c084596dbe33d9278 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Mon, 16 Feb 2026 22:09:52 -0800 Subject: [PATCH 02/10] Update tests for QueryPlan thinclient --- .../PartitionedQueryExecutionInfo.cs | 84 +- .../Query/v3Query/CosmosQueryClientCore.cs | 98 +- .../CosmosItemThinClientTests.cs | 878 +++++++++++------- .../Utils/MultiRegionSetupHelpers.cs | 3 + 4 files changed, 605 insertions(+), 458 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs index 39bd0e0bf1..4a33d48413 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs @@ -21,7 +21,7 @@ public PartitionedQueryExecutionInfo() this.Version = Constants.PartitionedQueryExecutionInfo.CurrentVersion; } - [JsonProperty(Constants.Properties.PartitionedQueryExecutionInfoVersion)] + [JsonProperty(Constants.Properties.PartitionedQueryExecutionInfoVersion)] public int Version { get; @@ -40,31 +40,31 @@ public QueryInfo QueryInfo /// lazily converts PartitionKeyInternal ranges to EPK hex string ranges. /// [JsonIgnore] - public List> QueryRanges - { - get - { - if (this.queryRanges != null) - { - return this.queryRanges; - } - - if (this.RawQueryRanges != null) - { - if (this.UseThinClientMode && this.PartitionKeyDefinition != null) - { - this.queryRanges = this.ParseQueryRangesForThinClient(); - } - else - { - // Non-thin client: deserialize directly as string ranges - this.queryRanges = this.RawQueryRanges.ToObject>>(); - } - } - - return this.queryRanges; - } - set => this.queryRanges = value; + public List> QueryRanges + { + get + { + if (this.queryRanges != null) + { + return this.queryRanges; + } + + if (this.RawQueryRanges != null) + { + if (this.PartitionKeyDefinition != null) + { + // convert PartitionKeyInternal format to EPK strings + this.queryRanges = this.ParseQueryRangesWithPartitionKeyDefinition(); + } + else + { + this.queryRanges = this.RawQueryRanges.ToObject>>(); + } + } + + return this.queryRanges; + } + set => this.queryRanges = value; } /// @@ -86,14 +86,6 @@ public HybridSearchQueryInfo HybridSearchQueryInfo get; set; } - - /// - /// Flag indicating if thin client mode is enabled. - /// Must be set before accessing QueryRanges property. - /// - [JsonIgnore] - internal bool UseThinClientMode { get; set; } - /// /// Partition key definition used for converting PartitionKeyInternal to EPK strings. /// Must be set before accessing QueryRanges property in thin client mode. @@ -123,20 +115,15 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec partitionedQueryExecutionInfo = default; return false; } - } - + } + /// /// Parses query ranges for thin client mode where the proxy returns ranges /// in PartitionKeyInternal format (e.g., {"min": [[""]], "max": [["Infinity"]]}) /// and converts them to EPK hex string ranges. /// - private List> ParseQueryRangesForThinClient() + private List> ParseQueryRangesWithPartitionKeyDefinition() { - if (this.RawQueryRanges == null || this.PartitionKeyDefinition == null) - { - return null; - } - List> epkRanges = new List>(this.RawQueryRanges.Count); foreach (JToken rangeToken in this.RawQueryRanges) @@ -146,18 +133,15 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec continue; } - // Parse min and max as PartitionKeyInternal JToken minToken = rangeObject["min"]; JToken maxToken = rangeObject["max"]; - PartitionKeyInternal minPk = this.ParsePartitionKeyInternal(minToken); - PartitionKeyInternal maxPk = this.ParsePartitionKeyInternal(maxToken); + PartitionKeyInternal minPk = PartitionedQueryExecutionInfo.ParsePartitionKeyInternal(minToken); + PartitionKeyInternal maxPk = PartitionedQueryExecutionInfo.ParsePartitionKeyInternal(maxToken); - // Convert to EPK hex strings string minEpk = minPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); string maxEpk = maxPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); - // Parse isMinInclusive and isMaxInclusive (defaults: min=true, max=false) bool isMinInclusive = rangeObject["isMinInclusive"]?.Value() ?? true; bool isMaxInclusive = rangeObject["isMaxInclusive"]?.Value() ?? false; @@ -165,13 +149,13 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec } return epkRanges; - } - + } + /// /// Parses a JSON token representing a PartitionKeyInternal. /// Handles formats like [[""]] (empty), [["Infinity"]] (infinity), or actual partition key values. /// - private PartitionKeyInternal ParsePartitionKeyInternal(JToken token) + private static PartitionKeyInternal ParsePartitionKeyInternal(JToken token) { if (token == null || token.Type == JTokenType.Null) { diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index d0a61c9f33..c4ecbdfa85 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -173,63 +173,53 @@ public override async Task> ExecuteItemQueryAsync( resourceType, message, trace); - } - - public override async Task ExecuteQueryPlanRequestAsync( - string resourceUri, - ResourceType resourceType, - OperationType operationType, - SqlQuerySpec sqlQuerySpec, - PartitionKey? partitionKey, - string supportedQueryFeatures, - Guid clientQueryCorrelationId, - ITrace trace, - CancellationToken cancellationToken) - { - // Determine if thin client mode is enabled - bool isThinClientEnabled = this.clientContext.ClientOptions.ConnectionMode == ConnectionMode.Gateway - && ConfigurationManager.GetEnvironmentVariable( - ConfigurationManager.ThinClientModeEnabled, - defaultValue: false); - - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; - using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( - resourceUri: resourceUri, - resourceType: resourceType, - operationType: operationType, - requestOptions: null, - feedRange: partitionKey.HasValue ? new FeedRangePartitionKey(partitionKey.Value) : null, - cosmosContainerCore: this.cosmosContainerCore, - streamPayload: this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType), - requestEnricher: (requestMessage) => - { - requestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, RuntimeConstants.MediaTypes.QueryJson); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQueryPlanRequest, bool.TrueString); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.SupportedQueryFeatures, supportedQueryFeatures); - requestMessage.Headers.Add(HttpConstants.HttpHeaders.QueryVersion, new Version(major: 1, minor: 0).ToString()); - requestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, clientQueryCorrelationId.ToString()); - if (!isThinClientEnabled) - { - requestMessage.UseGatewayMode = true; - } - }, - trace: trace, - cancellationToken: cancellationToken)) - { - // Syntax exception are argument exceptions and thrown to the user. - message.EnsureSuccessStatusCode(); + } + + public override async Task ExecuteQueryPlanRequestAsync( + string resourceUri, + ResourceType resourceType, + OperationType operationType, + SqlQuerySpec sqlQuerySpec, + PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + ITrace trace, + CancellationToken cancellationToken) + { + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( + resourceUri: resourceUri, + resourceType: resourceType, + operationType: operationType, + requestOptions: null, + feedRange: partitionKey.HasValue ? new FeedRangePartitionKey(partitionKey.Value) : null, + cosmosContainerCore: this.cosmosContainerCore, + streamPayload: this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType), + requestEnricher: (requestMessage) => + { + requestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, RuntimeConstants.MediaTypes.QueryJson); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQueryPlanRequest, bool.TrueString); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.SupportedQueryFeatures, supportedQueryFeatures); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.QueryVersion, new Version(major: 1, minor: 0).ToString()); + requestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, clientQueryCorrelationId.ToString()); + requestMessage.UseGatewayMode = true; + }, + trace: trace, + cancellationToken: cancellationToken)) + { + // Syntax exception are argument exceptions and thrown to the user. + message.EnsureSuccessStatusCode(); partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); - if (isThinClientEnabled) + if (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) { - partitionedQueryExecutionInfo.UseThinClientMode = true; ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( resourceUri, trace, cancellationToken); partitionedQueryExecutionInfo.PartitionKeyDefinition = containerProperties.PartitionKey; - } - } - - return partitionedQueryExecutionInfo; + } + } + + return partitionedQueryExecutionInfo; } public override async Task GetClientDisableOptimisticDirectExecutionAsync() @@ -307,7 +297,11 @@ public override async Task> GetTargetPartitionKeyRangesA } public override bool BypassQueryParsing() - { + { + if (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) + { + return true; + } return QueryPlanRetriever.BypassQueryParsing(); } 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..6309cba90b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs @@ -188,183 +188,212 @@ public async Task RegionalDatabaseAccountNameIsEmptyInPayload() [TestCategory("ThinClient")] public async Task TestThinClientWithExecuteStoredProcedureAsync() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); - - 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 + }; + this.cosmosSystemTextJsonSerializer = 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"); + 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(); + + 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 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" - }; + // 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 this.container.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 this.container.Scripts.DeleteStoredProcedureAsync(sprocId); + finally + { + if (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } + } } [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(); + try + { + 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.'); + 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.'); - }"; - - // 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"); + 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); + } + finally + { + if (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } - - // Delete stored procedure - await this.container.Scripts.DeleteStoredProcedureAsync(sprocId); } [TestMethod] @@ -426,42 +455,56 @@ 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); - - 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, - }); - - 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); + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); + string authKey = Utils.ConfigurationManager.AppSettings["MasterKey"]; + string endpoint = Utils.ConfigurationManager.AppSettings["GatewayEndpoint"]; + AzureKeyCredential masterKeyCredential = new AzureKeyCredential(authKey); - foreach (TestObject item in items) + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + 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, + }); + + 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) + { + 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"); + } + } + 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 (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } } @@ -469,42 +512,56 @@ public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() [TestCategory("ThinClient")] public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountEnabled() { - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + 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 + }; + this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( - this.connectionString, - new CosmosClientOptions() - { - ConnectionMode = ConnectionMode.Direct, - Serializer = this.cosmosSystemTextJsonSerializer, - }); + this.client = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Direct, + Serializer = this.cosmosSystemTextJsonSerializer, + }); - 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(); + 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); + string pk = "pk_create"; + IEnumerable items = this.GenerateItems(pk); - foreach (TestObject item in items) + foreach (TestObject item in items) + { + 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'"); + } + } + 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 (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } } @@ -512,38 +569,53 @@ public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountE [TestCategory("ThinClient")] public async Task CreateItemsTestWithThinClientFlagDisabledAccountEnabled() { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "False"); - - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + try { - 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, - }); + Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "False"); - 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"); + JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = null, + PropertyNameCaseInsensitive = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull + }; + this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - string pk = "pk_create"; - IEnumerable items = this.GenerateItems(pk); + this.client = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + Serializer = this.cosmosSystemTextJsonSerializer, + }); - foreach (TestObject item in items) + 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"); + + string pk = "pk_create"; + IEnumerable items = this.GenerateItems(pk); + + foreach (TestObject item in items) + { + 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"); + } + } + 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 (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } } @@ -764,77 +836,106 @@ public async Task QueryItemsTest() [TestCategory("ThinClient")] public async Task QueryItemsTestWithStrongConsistency() { - string connectionString = ConfigurationManager.GetEnvironmentVariable("COSMOSDB_THINCLIENTSTRONG", string.Empty); - if (string.IsNullOrEmpty(connectionString)) + 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 pk = "pk_query"; - List items = this.GenerateItems(pk).ToList(); - - List createdItems = await this.CreateItemsSafeAsync(items); - - string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + string connectionString = ConfigurationManager.GetEnvironmentVariable("COSMOSDB_THINCLIENTSTRONG", string.Empty); + if (string.IsNullOrEmpty(connectionString)) + { + 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 pk = "pk_query"; + List items = this.GenerateItems(pk).ToList(); + + List createdItems = await this.CreateItemsSafeAsync(items); + + 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) + Assert.AreEqual(createdItems.Count, count); + } + finally { - FeedResponse response = await iterator.ReadNextAsync(); - count += response.Count; + if (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } - - Assert.AreEqual(createdItems.Count, count); } [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"); + try + { - string pk = "pk_query"; - List items = this.GenerateItems(pk).ToList(); + this.client = new CosmosClient( + this.connectionString, + new CosmosClientOptions() + { + ConnectionMode = ConnectionMode.Gateway, + RequestTimeout = TimeSpan.FromSeconds(60), + ConsistencyLevel = Microsoft.Azure.Cosmos.ConsistencyLevel.Session + }); - List createdItems = await this.CreateItemsSafeAsync(items); + 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 query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + string pk = "pk_query"; + List items = this.GenerateItems(pk).ToList(); - int count = 0; - while (iterator.HasMoreResults) + List createdItems = await this.CreateItemsSafeAsync(items); + + 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; + } + + Assert.AreEqual(createdItems.Count, count); + } + finally { - FeedResponse response = await iterator.ReadNextAsync(); - count += response.Count; + if (this.database != null) + { + try + { + await this.database.DeleteAsync(); + } + catch { } + } } - - Assert.AreEqual(createdItems.Count, count); } [TestMethod] @@ -929,83 +1030,148 @@ public async Task TransactionalBatchCreateItemsTest() [TestMethod] [TestCategory("ThinClient")] - public async Task RegionalFailoverWithHttpRequestException_EnsuresThinClientHeaderInRefreshRequest() + public async Task TestThinClientQueryPlanWithOrderBy() { - // Arrange - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); - - bool headerFoundInRefreshRequest = false; - int accountRefreshCount = 0; - bool hasThrown = false; + List items = new List(); + string commonPk = "pk_orderby_test_" + Guid.NewGuid().ToString(); - JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions + try { - PropertyNamingPolicy = null, - PropertyNameCaseInsensitive = true, - DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull - }; - CosmosSystemTextJsonSerializer serializer = new CosmosSystemTextJsonSerializer(jsonSerializerOptions); - - FaultInjectionDelegatingHandler faultHandler = new FaultInjectionDelegatingHandler( - (request) => + for (int i = 0; i < 5; i++) { - // 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)) + items.Add(new TestObject { - accountRefreshCount++; + Id = Guid.NewGuid().ToString(), + Pk = commonPk, + Other = $"Item_{i:D3}", + SortField = i + }); + } - // 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; - } - } - } + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(5, createdItems.Count, "All items should be created"); + + await Task.Delay(1000); + + // Execute ORDER BY query - this requires QueryPlan and EPK range conversion + string query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.SortField 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"); + + for (int i = 0; i < results.Count; i++) + { + int expectedSortField = 4 - i; // Descending: 4, 3, 2, 1, 0 + Assert.AreEqual(expectedSortField, results[i].SortField, + $"Item at position {i} should have SortField={expectedSortField}"); + } + + } + finally + { + 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(); - // Throw HttpRequestException only ONCE on ThinClient POST requests - if (!hasThrown && - request.Method == HttpMethod.Post && - request.Version == new Version(2, 0)) + try + { + 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++) { - hasThrown = true; - throw new HttpRequestException("Simulated endpoint failure"); + items.Add(new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = partitionKeys[pkIndex], + Other = $"Value_{i}", + SortField = i + }); } - }); + } - CosmosClientBuilder builder = new CosmosClientBuilder(this.connectionString) - .WithConnectionModeGateway() - .WithCustomSerializer(serializer) - .WithHttpClientFactory(() => new HttpClient(faultHandler)); + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(9, createdItems.Count, "All 9 items should be created"); - using CosmosClient client = builder.Build(); + await Task.Delay(2000); + string query = "SELECT * FROM c ORDER BY c._ts"; - 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"); + FeedIterator iterator = this.container.GetItemQueryIterator(query); - string pk = "pk_failover_test"; - TestObject testItem = this.GenerateItems(pk).First(); + List results = new List(); + int pageCount = 0; - // Act - CreateItemAsync will fail once, then SDK retries and succeeds - ItemResponse response = await container.CreateItemAsync(testItem, new PartitionKey(testItem.Pk)); + while (iterator.HasMoreResults) + { + FeedResponse response = await iterator.ReadNextAsync(); + results.AddRange(response); + pageCount++; - // 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"); + string diagnostics = response.Diagnostics.ToString(); + Assert.IsTrue(diagnostics.Contains("|F4"), $"Page {pageCount}: Should use ThinClient"); + } - // Cleanup - await database.DeleteAsync(); + 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 + { + foreach (TestObject item in items) + { + try + { + await this.container.DeleteItemAsync(item.Id, new PartitionKey(item.Pk)); + } + catch { } + } + } } /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs index ad8cf2863e..2ea6cdb958 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs @@ -77,6 +77,9 @@ internal class CosmosIntegrationTestObject [JsonPropertyName("other")] public string Other { get; set; } + + [JsonPropertyName("sortField")] + public int SortField { get; set; } } internal class CosmosSystemTextJsonSerializer : CosmosSerializer From 4ee8687bb7863a0c9af5617d064d6fec803cba9e Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Wed, 18 Feb 2026 22:26:44 -0800 Subject: [PATCH 03/10] Fix tests for thinclient --- .../CosmosItemThinClientTests.cs | 281 ++++++++++++------ 1 file changed, 183 insertions(+), 98 deletions(-) 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 6309cba90b..493f18c78c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs @@ -188,6 +188,9 @@ public async Task RegionalDatabaseAccountNameIsEmptyInPayload() [TestCategory("ThinClient")] public async Task TestThinClientWithExecuteStoredProcedureAsync() { + CosmosClient localClient = null; + Database localDatabase = null; + try { Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); @@ -198,20 +201,20 @@ public async Task TestThinClientWithExecuteStoredProcedureAsync() PropertyNameCaseInsensitive = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( + localClient = new CosmosClient( this.connectionString, new CosmosClientOptions() { ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, + Serializer = localSerializer, }); string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string sprocId = "testSproc_" + Guid.NewGuid().ToString(); @@ -249,7 +252,7 @@ public async Task TestThinClientWithExecuteStoredProcedureAsync() }"; // Create stored procedure - Scripts.StoredProcedureResponse createResponse = await this.container.Scripts.CreateStoredProcedureAsync( + Scripts.StoredProcedureResponse createResponse = await localContainer.Scripts.CreateStoredProcedureAsync( new Scripts.StoredProcedureProperties(sprocId, sprocBody)); Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); @@ -263,7 +266,7 @@ public async Task TestThinClientWithExecuteStoredProcedureAsync() }; Scripts.StoredProcedureExecuteResponse executeResponse = - await this.container.Scripts.ExecuteStoredProcedureAsync( + await localContainer.Scripts.ExecuteStoredProcedureAsync( sprocId, new PartitionKey(testPartitionId), new dynamic[] { testItem }); @@ -274,26 +277,33 @@ await this.container.Scripts.ExecuteStoredProcedureAsync( Assert.IsTrue(diagnostics.Contains("|F4"), "Diagnostics User Agent should contain '|F4' for ThinClient"); // Delete stored procedure - await this.container.Scripts.DeleteStoredProcedureAsync(sprocId); + await localContainer.Scripts.DeleteStoredProcedureAsync(sprocId); } - finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } - [TestMethod] - [TestCategory("ThinClient")] - public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() { + CosmosClient localClient = null; + Database localDatabase = null; + try { Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "true"); @@ -304,58 +314,58 @@ public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() PropertyNameCaseInsensitive = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( + localClient = new CosmosClient( this.connectionString, new CosmosClientOptions() { ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, + Serializer = localSerializer, }); string uniqueDbName = "TestDbStoreProc_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestDbStoreProcContainer_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + 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.'); + 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( + Scripts.StoredProcedureResponse createResponse = await localContainer.Scripts.CreateStoredProcedureAsync( new Scripts.StoredProcedureProperties(sprocId, sprocBody)); Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode); @@ -369,7 +379,7 @@ public async Task TestThinClientWithExecuteStoredProcedureStreamAsync() }; using (ResponseMessage executeResponse = - await this.container.Scripts.ExecuteStoredProcedureStreamAsync( + await localContainer.Scripts.ExecuteStoredProcedureStreamAsync( sprocId, new PartitionKey(testPartitionId), new dynamic[] { testItem })) @@ -381,19 +391,24 @@ await this.container.Scripts.ExecuteStoredProcedureStreamAsync( } // Delete stored procedure - await this.container.Scripts.DeleteStoredProcedureAsync(sprocId); + await localContainer.Scripts.DeleteStoredProcedureAsync(sprocId); } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } - } + + if (localClient != null) + { + localClient.Dispose(); + } + } } [TestMethod] @@ -455,6 +470,10 @@ public async Task CreateItemsTest() [TestCategory("ThinClient")] public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() { + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; + try { Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True"); @@ -468,28 +487,28 @@ public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() PropertyNameCaseInsensitive = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( + localClient = new CosmosClient( endpoint, masterKeyCredential, new CosmosClientOptions() { ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, + Serializer = localSerializer, }); string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string pk = "pk_create"; IEnumerable items = this.GenerateItems(pk); foreach (TestObject item in items) { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); + 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"); @@ -497,14 +516,19 @@ public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -512,6 +536,10 @@ public async Task CreateItemsTestWithThinClientFlagEnabledAndAccountDisabled() [TestCategory("ThinClient")] public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountEnabled() { + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; + try { JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions @@ -520,27 +548,27 @@ public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountE PropertyNameCaseInsensitive = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( + localClient = new CosmosClient( this.connectionString, new CosmosClientOptions() { ConnectionMode = ConnectionMode.Direct, - Serializer = this.cosmosSystemTextJsonSerializer, + Serializer = localSerializer, }); string uniqueDbName = "TestDb2_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestContainer2_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string pk = "pk_create"; IEnumerable items = this.GenerateItems(pk); foreach (TestObject item in items) { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); + 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 @@ -554,14 +582,19 @@ public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountE } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -569,9 +602,12 @@ public async Task CreateItemsTestWithDirectMode_ThinClientFlagEnabledAndAccountE [TestCategory("ThinClient")] public async Task CreateItemsTestWithThinClientFlagDisabledAccountEnabled() { + CosmosClient localClient = null; + Database localDatabase = null; + Container localContainer = null; + try { - Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "False"); JsonSerializerOptions jsonSerializerOptions = new JsonSerializerOptions @@ -580,27 +616,27 @@ public async Task CreateItemsTestWithThinClientFlagDisabledAccountEnabled() PropertyNameCaseInsensitive = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull }; - this.cosmosSystemTextJsonSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); + CosmosSystemTextJsonSerializer localSerializer = new MultiRegionSetupHelpers.CosmosSystemTextJsonSerializer(jsonSerializerOptions); - this.client = new CosmosClient( + localClient = new CosmosClient( this.connectionString, new CosmosClientOptions() { ConnectionMode = ConnectionMode.Gateway, - Serializer = this.cosmosSystemTextJsonSerializer, + Serializer = localSerializer, }); string uniqueDbName = "TestDbTCDisabled_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestContainerTCDisabled_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string pk = "pk_create"; IEnumerable items = this.GenerateItems(pk); foreach (TestObject item in items) { - ItemResponse response = await this.container.CreateItemAsync(item, new PartitionKey(item.Pk)); + 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"); @@ -608,14 +644,19 @@ public async Task CreateItemsTestWithThinClientFlagDisabledAccountEnabled() } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -836,6 +877,9 @@ public async Task QueryItemsTest() [TestCategory("ThinClient")] public async Task QueryItemsTestWithStrongConsistency() { + CosmosClient localClient = null; + Database localDatabase = null; + try { string connectionString = ConfigurationManager.GetEnvironmentVariable("COSMOSDB_THINCLIENTSTRONG", string.Empty); @@ -843,7 +887,8 @@ public async Task QueryItemsTestWithStrongConsistency() { Assert.Fail("Set environment variable COSMOSDB_THINCLIENTSTRONG to run the tests"); } - this.client = new CosmosClient( + + localClient = new CosmosClient( connectionString, new CosmosClientOptions() { @@ -853,17 +898,31 @@ public async Task QueryItemsTestWithStrongConsistency() }); string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string pk = "pk_query"; List items = this.GenerateItems(pk).ToList(); - List createdItems = await this.CreateItemsSafeAsync(items); + 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) + { + } + } string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + FeedIterator iterator = localContainer.GetItemQueryIterator(query); int count = 0; while (iterator.HasMoreResults) @@ -872,18 +931,23 @@ public async Task QueryItemsTestWithStrongConsistency() count += response.Count; } - Assert.AreEqual(createdItems.Count, count); + Assert.AreEqual(itemsCreated.Count, count); } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } @@ -891,10 +955,12 @@ public async Task QueryItemsTestWithStrongConsistency() [TestCategory("ThinClient")] public async Task QueryItemsTestWithSessionConsistency() { + CosmosClient localClient = null; + Database localDatabase = null; + try { - - this.client = new CosmosClient( + localClient = new CosmosClient( this.connectionString, new CosmosClientOptions() { @@ -904,17 +970,31 @@ public async Task QueryItemsTestWithSessionConsistency() }); string uniqueDbName = "TestDbTC_" + Guid.NewGuid().ToString(); - this.database = await this.client.CreateDatabaseIfNotExistsAsync(uniqueDbName); + localDatabase = await localClient.CreateDatabaseIfNotExistsAsync(uniqueDbName); string uniqueContainerName = "TestContainerTC_" + Guid.NewGuid().ToString(); - this.container = await this.database.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); + Container localContainer = await localDatabase.CreateContainerIfNotExistsAsync(uniqueContainerName, "/pk"); string pk = "pk_query"; List items = this.GenerateItems(pk).ToList(); - List createdItems = await this.CreateItemsSafeAsync(items); + 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) + { + } + } string query = $"SELECT * FROM c WHERE c.pk = '{pk}'"; - FeedIterator iterator = this.container.GetItemQueryIterator(query); + FeedIterator iterator = localContainer.GetItemQueryIterator(query); int count = 0; while (iterator.HasMoreResults) @@ -923,18 +1003,23 @@ public async Task QueryItemsTestWithSessionConsistency() count += response.Count; } - Assert.AreEqual(createdItems.Count, count); + Assert.AreEqual(itemsCreated.Count, count); } finally { - if (this.database != null) + if (localDatabase != null) { try { - await this.database.DeleteAsync(); + await localDatabase.DeleteAsync(); } catch { } } + + if (localClient != null) + { + localClient.Dispose(); + } } } From 2bd0264eb05fc4c8b52b7c47ab94d5434cd050ee Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 19 Feb 2026 09:51:43 -0800 Subject: [PATCH 04/10] Fix tests for thinclient --- .../tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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, From c0b3c0e9808ec8de55ba0934f79a16e0e08c1cc1 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 26 Feb 2026 18:37:16 -0800 Subject: [PATCH 05/10] Update tests and logic for query by passing --- .../Query/v3Query/CosmosQueryClientCore.cs | 6 +- .../CosmosItemThinClientTests.cs | 288 +++++++++--------- 2 files changed, 149 insertions(+), 145 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index c4ecbdfa85..37c978a9c5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -297,11 +297,7 @@ public override async Task> GetTargetPartitionKeyRangesA } public override bool BypassQueryParsing() - { - if (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) - { - return true; - } + { return QueryPlanRetriever.BypassQueryParsing(); } 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 493f18c78c..a789d4b73f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs @@ -1113,150 +1113,158 @@ public async Task TransactionalBatchCreateItemsTest() } } - [TestMethod] - [TestCategory("ThinClient")] - public async Task TestThinClientQueryPlanWithOrderBy() - { - List items = new List(); - string commonPk = "pk_orderby_test_" + Guid.NewGuid().ToString(); - - try - { - for (int i = 0; i < 5; i++) - { - items.Add(new TestObject - { - Id = Guid.NewGuid().ToString(), - Pk = commonPk, - Other = $"Item_{i:D3}", - SortField = i - }); - } - - List createdItems = await this.CreateItemsSafeAsync(items); - Assert.AreEqual(5, createdItems.Count, "All items should be created"); - - await Task.Delay(1000); - - // Execute ORDER BY query - this requires QueryPlan and EPK range conversion - string query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.SortField 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"); - - for (int i = 0; i < results.Count; i++) - { - int expectedSortField = 4 - i; // Descending: 4, 3, 2, 1, 0 - Assert.AreEqual(expectedSortField, results[i].SortField, - $"Item at position {i} should have SortField={expectedSortField}"); - } - - } - finally - { - 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 - { + [TestMethod] + [TestCategory("ThinClient")] + public async Task TestThinClientQueryPlanWithOrderBy() + { + List items = new List(); + string commonPk = "pk_orderby_test_" + Guid.NewGuid().ToString(); + + try + { + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, "True"); + + for (int i = 0; i < 5; i++) + { + items.Add(new TestObject + { + Id = Guid.NewGuid().ToString(), + Pk = commonPk, + Other = $"Item_{i:D3}", + SortField = i + }); + } + + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(5, createdItems.Count, "All items should be created"); + + await Task.Delay(1000); + + // Execute ORDER BY query - this requires QueryPlan and EPK range conversion + string query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.SortField 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"); + + for (int i = 0; i < results.Count; i++) + { + int expectedSortField = 4 - i; // Descending: 4, 3, 2, 1, 0 + Assert.AreEqual(expectedSortField, results[i].SortField, + $"Item at position {i} should have SortField={expectedSortField}"); + } + + } + 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}", - SortField = i - }); - } - } - - List createdItems = await this.CreateItemsSafeAsync(items); - Assert.AreEqual(9, createdItems.Count, "All 9 items should be created"); - - await Task.Delay(2000); - 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 - { - foreach (TestObject item in items) - { - try - { - await this.container.DeleteItemAsync(item.Id, new PartitionKey(item.Pk)); - } - catch { } - } - } + }; + + 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}", + SortField = i + }); + } + } + + List createdItems = await this.CreateItemsSafeAsync(items); + Assert.AreEqual(9, createdItems.Count, "All 9 items should be created"); + + await Task.Delay(2000); + 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 { } + } + } } /// From fe47693e6ac00aaae64783cb5261ff13c3a3e262 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Wed, 4 Mar 2026 23:33:18 -0800 Subject: [PATCH 06/10] Add a separate class for thinclient parsing response --- .../Query/v3Query/CosmosQueryClientCore.cs | 42 ++++--- .../src/ThinClientQueryPlanHelper.cs | 110 ++++++++++++++++++ 2 files changed, 134 insertions(+), 18 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 37c978a9c5..98f48b7f9b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -173,17 +173,17 @@ public override async Task> ExecuteItemQueryAsync( resourceType, message, trace); - } - - public override async Task ExecuteQueryPlanRequestAsync( - string resourceUri, - ResourceType resourceType, - OperationType operationType, - SqlQuerySpec sqlQuerySpec, - PartitionKey? partitionKey, - string supportedQueryFeatures, - Guid clientQueryCorrelationId, - ITrace trace, + } + + public override async Task ExecuteQueryPlanRequestAsync( + string resourceUri, + ResourceType resourceType, + OperationType operationType, + SqlQuerySpec sqlQuerySpec, + PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + ITrace trace, CancellationToken cancellationToken) { PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; @@ -209,13 +209,19 @@ 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 (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) - { - ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( - resourceUri, trace, cancellationToken); - partitionedQueryExecutionInfo.PartitionKeyDefinition = containerProperties.PartitionKey; + + if (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) + { + ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( + resourceUri, trace, cancellationToken); + + partitionedQueryExecutionInfo = ThinClientQueryPlanHelper.DeserializeQueryPlanResponse( + message.Content, + containerProperties.PartitionKey); + } + else + { + partitionedQueryExecutionInfo = this.clientContext.SerializerCore.FromStream(message.Content); } } diff --git a/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs new file mode 100644 index 0000000000..a1ca962ba2 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs @@ -0,0 +1,110 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan +{ + using System; + using System.Collections.Generic; + using System.IO; + using Newtonsoft.Json; + using Newtonsoft.Json.Linq; + 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. + /// + internal static class ThinClientQueryPlanHelper + { + /// + /// Reads the raw query plan JSON from a stream, converts PartitionKeyInternal + /// ranges to EPK hex string ranges, and deserializes into a clean + /// DTO. + /// + /// The response stream containing the raw query plan JSON. + /// The partition key definition for the container. + /// with EPK string ranges. + public static PartitionedQueryExecutionInfo DeserializeQueryPlanResponse( + Stream stream, + PartitionKeyDefinition partitionKeyDefinition) + { + if (stream == null) + { + throw new ArgumentNullException(nameof(stream)); + } + + if (partitionKeyDefinition == null) + { + throw new ArgumentNullException(nameof(partitionKeyDefinition)); + } + + JObject queryPlanJson; + using (StreamReader reader = new StreamReader(stream)) + using (JsonTextReader jsonReader = new JsonTextReader(reader)) + { + queryPlanJson = JObject.Load(jsonReader); + } + + if (queryPlanJson[Documents.Constants.Properties.QueryRanges] is JArray rawQueryRanges) + { + List> epkRanges = ThinClientQueryPlanHelper.ConvertToEpkRanges( + rawQueryRanges, + partitionKeyDefinition); + + queryPlanJson[Documents.Constants.Properties.QueryRanges] = JToken.FromObject(epkRanges); + } + + return queryPlanJson.ToObject(); + } + + private static List> ConvertToEpkRanges( + JArray rawQueryRanges, + PartitionKeyDefinition partitionKeyDefinition) + { + List> epkRanges = new List>(rawQueryRanges.Count); + + foreach (JToken rangeToken in rawQueryRanges) + { + if (!(rangeToken is JObject rangeObject)) + { + continue; + } + + JToken minToken = rangeObject["min"]; + JToken maxToken = rangeObject["max"]; + + PartitionKeyInternal minPk = ThinClientQueryPlanHelper.ParsePartitionKeyInternal(minToken); + PartitionKeyInternal maxPk = ThinClientQueryPlanHelper.ParsePartitionKeyInternal(maxToken); + + string minEpk = minPk.GetEffectivePartitionKeyString(partitionKeyDefinition); + string maxEpk = maxPk.GetEffectivePartitionKeyString(partitionKeyDefinition); + + bool isMinInclusive = rangeObject["isMinInclusive"]?.Value() ?? true; + bool isMaxInclusive = rangeObject["isMaxInclusive"]?.Value() ?? false; + + epkRanges.Add(new Documents.Routing.Range(minEpk, maxEpk, isMinInclusive, isMaxInclusive)); + } + + return epkRanges; + } + + private static PartitionKeyInternal ParsePartitionKeyInternal(JToken token) + { + if (token == null || token.Type == JTokenType.Null) + { + return PartitionKeyInternal.Empty; + } + + try + { + return token.ToObject(); + } + catch (JsonException) + { + return PartitionKeyInternal.Empty; + } + } + } +} \ No newline at end of file From d83421f0868fb99b84a001b59f32f3d33e321f86 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Wed, 4 Mar 2026 23:38:28 -0800 Subject: [PATCH 07/10] Revert changes in PartitionedQueryExecutionInfo --- .../PartitionedQueryExecutionInfo.cs | 124 ++---------------- 1 file changed, 11 insertions(+), 113 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs index 4a33d48413..14ae921d8f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs @@ -7,21 +7,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan using System; using System.Collections.Generic; using Newtonsoft.Json; - using Newtonsoft.Json.Linq; using Constants = Documents.Constants; - using PartitionKeyDefinition = Documents.PartitionKeyDefinition; - using PartitionKeyInternal = Documents.Routing.PartitionKeyInternal; internal sealed class PartitionedQueryExecutionInfo { - private List> queryRanges; - public PartitionedQueryExecutionInfo() { this.Version = Constants.PartitionedQueryExecutionInfo.CurrentVersion; } - [JsonProperty(Constants.Properties.PartitionedQueryExecutionInfoVersion)] + [JsonProperty(Constants.Properties.PartitionedQueryExecutionInfoVersion)] public int Version { get; @@ -35,63 +30,21 @@ public QueryInfo QueryInfo set; } - /// - /// Gets or sets the query ranges. In thin client mode, this property - /// lazily converts PartitionKeyInternal ranges to EPK hex string ranges. - /// - [JsonIgnore] - public List> QueryRanges - { - get - { - if (this.queryRanges != null) - { - return this.queryRanges; - } - - if (this.RawQueryRanges != null) - { - if (this.PartitionKeyDefinition != null) - { - // convert PartitionKeyInternal format to EPK strings - this.queryRanges = this.ParseQueryRangesWithPartitionKeyDefinition(); - } - else - { - this.queryRanges = this.RawQueryRanges.ToObject>>(); - } - } - - return this.queryRanges; - } - set => this.queryRanges = value; - } - - /// - /// Raw query ranges from JSON deserialization. Used for thin client mode parsing. - /// In non-thin client mode, this is deserialized directly to QueryRanges. - /// [JsonProperty(Constants.Properties.QueryRanges)] - internal JArray RawQueryRanges - { - get; - set; - } - - // Change to the below after Direct package upgrade - // [JsonProperty(Constants.Properties.HybridSearchQueryInfo)] - [JsonProperty("hybridSearchQueryInfo")] - public HybridSearchQueryInfo HybridSearchQueryInfo + public List> QueryRanges { get; set; + } + + // Change to the below after Direct package upgrade + // [JsonProperty(Constants.Properties.HybridSearchQueryInfo)] + [JsonProperty("hybridSearchQueryInfo")] + public HybridSearchQueryInfo HybridSearchQueryInfo + { + get; + set; } - /// - /// Partition key definition used for converting PartitionKeyInternal to EPK strings. - /// Must be set before accessing QueryRanges property in thin client mode. - /// - [JsonIgnore] - internal PartitionKeyDefinition PartitionKeyDefinition { get; set; } public override string ToString() { @@ -115,61 +68,6 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec partitionedQueryExecutionInfo = default; return false; } - } - - /// - /// Parses query ranges for thin client mode where the proxy returns ranges - /// in PartitionKeyInternal format (e.g., {"min": [[""]], "max": [["Infinity"]]}) - /// and converts them to EPK hex string ranges. - /// - private List> ParseQueryRangesWithPartitionKeyDefinition() - { - List> epkRanges = new List>(this.RawQueryRanges.Count); - - foreach (JToken rangeToken in this.RawQueryRanges) - { - if (!(rangeToken is JObject rangeObject)) - { - continue; - } - - JToken minToken = rangeObject["min"]; - JToken maxToken = rangeObject["max"]; - - PartitionKeyInternal minPk = PartitionedQueryExecutionInfo.ParsePartitionKeyInternal(minToken); - PartitionKeyInternal maxPk = PartitionedQueryExecutionInfo.ParsePartitionKeyInternal(maxToken); - - string minEpk = minPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); - string maxEpk = maxPk.GetEffectivePartitionKeyString(this.PartitionKeyDefinition); - - bool isMinInclusive = rangeObject["isMinInclusive"]?.Value() ?? true; - bool isMaxInclusive = rangeObject["isMaxInclusive"]?.Value() ?? false; - - epkRanges.Add(new Documents.Routing.Range(minEpk, maxEpk, isMinInclusive, isMaxInclusive)); - } - - return epkRanges; - } - - /// - /// Parses a JSON token representing a PartitionKeyInternal. - /// Handles formats like [[""]] (empty), [["Infinity"]] (infinity), or actual partition key values. - /// - private static PartitionKeyInternal ParsePartitionKeyInternal(JToken token) - { - if (token == null || token.Type == JTokenType.Null) - { - return PartitionKeyInternal.Empty; - } - - try - { - return token.ToObject(); - } - catch (JsonException) - { - return PartitionKeyInternal.Empty; - } } } } From f223861edc07608b5503ba65159a2a2b61581c22 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Thu, 5 Mar 2026 16:21:47 -0800 Subject: [PATCH 08/10] Update flag check --- .../Query/v3Query/CosmosQueryClientCore.cs | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 98f48b7f9b..c8273575ca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -176,15 +176,15 @@ public override async Task> ExecuteItemQueryAsync( } public override async Task ExecuteQueryPlanRequestAsync( - string resourceUri, - ResourceType resourceType, - OperationType operationType, - SqlQuerySpec sqlQuerySpec, - PartitionKey? partitionKey, - string supportedQueryFeatures, - Guid clientQueryCorrelationId, - ITrace trace, - CancellationToken cancellationToken) + string resourceUri, + ResourceType resourceType, + OperationType operationType, + SqlQuerySpec sqlQuerySpec, + PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + ITrace trace, + CancellationToken cancellationToken) { PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( @@ -210,7 +210,7 @@ public override async Task ExecuteQueryPlanReques // Syntax exception are argument exceptions and thrown to the user. message.EnsureSuccessStatusCode(); - if (ConfigurationManager.IsThinClientEnabled(defaultValue: false)) + if (this.documentClient.isThinClientEnabled) { ContainerProperties containerProperties = await this.clientContext.GetCachedContainerPropertiesAsync( resourceUri, trace, cancellationToken); From f107e1cbed1a5055bbd8522108fc156402be4981 Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Sat, 7 Mar 2026 16:27:27 -0800 Subject: [PATCH 09/10] Test code cleanup --- .../CosmosItemThinClientTests.cs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 a789d4b73f..1aacf41ab1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs @@ -1122,7 +1122,7 @@ public async Task TestThinClientQueryPlanWithOrderBy() try { - Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, "True"); + Environment.SetEnvironmentVariable(ConfigurationManager.BypassQueryParsing, Boolean.TrueString); for (int i = 0; i < 5; i++) { @@ -1138,8 +1138,6 @@ public async Task TestThinClientQueryPlanWithOrderBy() List createdItems = await this.CreateItemsSafeAsync(items); Assert.AreEqual(5, createdItems.Count, "All items should be created"); - await Task.Delay(1000); - // Execute ORDER BY query - this requires QueryPlan and EPK range conversion string query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.SortField DESC"; QueryDefinition queryDef = new QueryDefinition(query).WithParameter("@pk", commonPk); From 9062eec78aa1bf8195f824e4b1d7aa4b7294767a Mon Sep 17 00:00:00 2001 From: Arooshi Avasthy Date: Fri, 13 Mar 2026 00:35:19 -0700 Subject: [PATCH 10/10] Update query plan serialization repsonse logic --- .../src/GatewayStoreModel.cs | 10 +- .../Query/v3Query/CosmosQueryClientCore.cs | 20 +-- .../src/ThinClientQueryPlanHelper.cs | 150 +++++++++++------ .../CosmosItemThinClientTests.cs | 125 +++++++++++--- .../Utils/MultiRegionSetupHelpers.cs | 3 - .../ThinClientStoreClientTests.cs | 153 ++++++++++++++++-- 6 files changed, 363 insertions(+), 98 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index d41ffeaa20..db8a549387 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -101,15 +101,11 @@ await GatewayStoreModel.ApplySessionTokenAsync( request.RequestContext.RegionName = regionName; } - bool isQueryPlanInThinClientMode = this.isThinClientEnabled - && request.OperationType == OperationType.QueryPlan - && request.ResourceType == ResourceType.Document; - - bool isPPAFEnabled = this.IsPartitionLevelFailoverEnabled(); - // This is applicable for both per partition automatic failover and per partition circuit breaker. + 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 || isQueryPlanInThinClientMode)) + && (request.ResourceType.IsPartitioned() || request.ResourceType == ResourceType.StoredProcedure)) { (bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync( request: request, diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index c8273575ca..d210dae18d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -175,16 +175,16 @@ public override async Task> ExecuteItemQueryAsync( trace); } - public override async Task ExecuteQueryPlanRequestAsync( - string resourceUri, - ResourceType resourceType, - OperationType operationType, - SqlQuerySpec sqlQuerySpec, - PartitionKey? partitionKey, - string supportedQueryFeatures, - Guid clientQueryCorrelationId, - ITrace trace, - CancellationToken cancellationToken) + public override async Task ExecuteQueryPlanRequestAsync( + string resourceUri, + ResourceType resourceType, + OperationType operationType, + SqlQuerySpec sqlQuerySpec, + PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + ITrace trace, + CancellationToken cancellationToken) { PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; using (ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync( diff --git a/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs index a1ca962ba2..6edb6c21a0 100644 --- a/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs +++ b/Microsoft.Azure.Cosmos/src/ThinClientQueryPlanHelper.cs @@ -7,25 +7,43 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan using System; using System.Collections.Generic; using System.IO; + using System.Text.Json; using Newtonsoft.Json; - using Newtonsoft.Json.Linq; 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, + }; + /// - /// Reads the raw query plan JSON from a stream, converts PartitionKeyInternal - /// ranges to EPK hex string ranges, and deserializes into a clean - /// DTO. + /// 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 EPK string ranges. + /// 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) @@ -40,71 +58,105 @@ public static PartitionedQueryExecutionInfo DeserializeQueryPlanResponse( throw new ArgumentNullException(nameof(partitionKeyDefinition)); } - JObject queryPlanJson; - using (StreamReader reader = new StreamReader(stream)) - using (JsonTextReader jsonReader = new JsonTextReader(reader)) + using JsonDocument doc = JsonDocument.Parse(stream); + JsonElement root = doc.RootElement; + + if (root.ValueKind != JsonValueKind.Object) { - queryPlanJson = JObject.Load(jsonReader); + throw new FormatException( + $"Thin client query plan response must be a JSON object, but was {root.ValueKind}."); } - if (queryPlanJson[Documents.Constants.Properties.QueryRanges] is JArray rawQueryRanges) + // Validate and extract queryRanges (required) + if (!root.TryGetProperty("queryRanges", out JsonElement queryRangesElement)) { - List> epkRanges = ThinClientQueryPlanHelper.ConvertToEpkRanges( - rawQueryRanges, - partitionKeyDefinition); + throw new FormatException( + "Thin client query plan response is missing the required 'queryRanges' property."); + } - queryPlanJson[Documents.Constants.Properties.QueryRanges] = JToken.FromObject(epkRanges); + if (queryRangesElement.ValueKind != JsonValueKind.Array) + { + throw new FormatException( + $"Expected 'queryRanges' to be a JSON array, but was {queryRangesElement.ValueKind}."); } - return queryPlanJson.ToObject(); - } + if (queryRangesElement.GetArrayLength() == 0) + { + throw new FormatException( + "Thin client query plan response 'queryRanges' array must not be empty."); + } - private static List> ConvertToEpkRanges( - JArray rawQueryRanges, - PartitionKeyDefinition partitionKeyDefinition) - { - List> epkRanges = new List>(rawQueryRanges.Count); + // 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); + } - foreach (JToken rangeToken in rawQueryRanges) + // 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 (!(rangeToken is JObject rangeObject)) + if (rangeElement.ValueKind != JsonValueKind.Object) { - continue; + throw new FormatException( + $"Each query range must be a JSON object, but was {rangeElement.ValueKind}."); } - JToken minToken = rangeObject["min"]; - JToken maxToken = rangeObject["max"]; + if (!rangeElement.TryGetProperty("min", out _)) + { + throw new FormatException( + "Query range is missing the required 'min' property."); + } - PartitionKeyInternal minPk = ThinClientQueryPlanHelper.ParsePartitionKeyInternal(minToken); - PartitionKeyInternal maxPk = ThinClientQueryPlanHelper.ParsePartitionKeyInternal(maxToken); + if (!rangeElement.TryGetProperty("max", out _)) + { + throw new FormatException( + "Query range is missing the required 'max' property."); + } - string minEpk = minPk.GetEffectivePartitionKeyString(partitionKeyDefinition); - string maxEpk = maxPk.GetEffectivePartitionKeyString(partitionKeyDefinition); + Documents.Routing.Range internalRange = + Newtonsoft.Json.JsonConvert.DeserializeObject>( + rangeElement.GetRawText(), + ThinClientQueryPlanHelper.NewtonsoftSettings); - bool isMinInclusive = rangeObject["isMinInclusive"]?.Value() ?? true; - bool isMaxInclusive = rangeObject["isMaxInclusive"]?.Value() ?? false; + if (internalRange == null) + { + throw new FormatException( + "Failed to deserialize query range from thin client response."); + } - epkRanges.Add(new Documents.Routing.Range(minEpk, maxEpk, isMinInclusive, isMaxInclusive)); + effectiveRanges.Add(PartitionKeyInternal.GetEffectivePartitionKeyRange( + partitionKeyDefinition, + internalRange)); } - return epkRanges; - } + effectiveRanges.Sort(Documents.Routing.Range.MinComparer.Instance); - private static PartitionKeyInternal ParsePartitionKeyInternal(JToken token) - { - if (token == null || token.Type == JTokenType.Null) - { - return PartitionKeyInternal.Empty; - } - - try + return new PartitionedQueryExecutionInfo() { - return token.ToObject(); - } - catch (JsonException) - { - return PartitionKeyInternal.Empty; - } + 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 1aacf41ab1..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"); @@ -1131,7 +1131,6 @@ public async Task TestThinClientQueryPlanWithOrderBy() Id = Guid.NewGuid().ToString(), Pk = commonPk, Other = $"Item_{i:D3}", - SortField = i }); } @@ -1139,7 +1138,7 @@ public async Task TestThinClientQueryPlanWithOrderBy() 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.SortField DESC"; + 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); @@ -1159,13 +1158,6 @@ public async Task TestThinClientQueryPlanWithOrderBy() Assert.AreEqual(5, results.Count, "Should return all 5 items"); - for (int i = 0; i < results.Count; i++) - { - int expectedSortField = 4 - i; // Descending: 4, 3, 2, 1, 0 - Assert.AreEqual(expectedSortField, results[i].SortField, - $"Item at position {i} should have SortField={expectedSortField}"); - } - } finally { @@ -1208,7 +1200,6 @@ public async Task TestThinClientQueryPlanCrossPartitionWithFilter() Id = Guid.NewGuid().ToString(), Pk = partitionKeys[pkIndex], Other = $"Value_{i}", - SortField = i }); } } @@ -1216,7 +1207,6 @@ public async Task TestThinClientQueryPlanCrossPartitionWithFilter() List createdItems = await this.CreateItemsSafeAsync(items); Assert.AreEqual(9, createdItems.Count, "All 9 items should be created"); - await Task.Delay(2000); string query = "SELECT * FROM c ORDER BY c._ts"; FeedIterator iterator = this.container.GetItemQueryIterator(query); @@ -1263,8 +1253,107 @@ public async Task TestThinClientQueryPlanCrossPartitionWithFilter() 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.EmulatorTests/Utils/MultiRegionSetupHelpers.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs index 2ea6cdb958..ad8cf2863e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/MultiRegionSetupHelpers.cs @@ -77,9 +77,6 @@ internal class CosmosIntegrationTestObject [JsonPropertyName("other")] public string Other { get; set; } - - [JsonPropertyName("sortField")] - public int SortField { get; set; } } internal class CosmosSystemTextJsonSerializer : CosmosSerializer 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