Add support for strong consistency in query.

Author: aavasthy

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

@@ -24,7 +24,8 @@ namespace Microsoft.Azure.Cosmos
     // Marking it as non-sealed in order to unit test it using Moq framework
     internal class GatewayStoreModel : IStoreModelExtension, IDisposable
     {
-        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isThinClientEnabled;
         private static readonly string sessionConsistencyAsString = ConsistencyLevel.Session.ToString();
         private readonly GlobalPartitionEndpointManager globalPartitionEndpointManager;
 
@@ -64,6 +65,7 @@ public GatewayStoreModel(
                 this.eventSource,
                 serializerSettings,
                 isPartitionLevelFailoverEnabled);
+            this.isThinClientEnabled = isThinClientEnabled;
 
             if (isThinClientEnabled)
             {
@@ -99,7 +101,7 @@ await GatewayStoreModel.ApplySessionTokenAsync(
                 }
 
                 // This is applicable for both per partition automatic failover and per partition circuit breaker.
-                if (this.isPartitionLevelFailoverEnabled
+                if ((this.isPartitionLevelFailoverEnabled || this.isThinClientEnabled)
                     && !ReplicatedResourceClient.IsMasterResource(request.ResourceType)
                     && request.ResourceType.IsPartitioned())
                 {
@@ -111,7 +113,11 @@ await GatewayStoreModel.ApplySessionTokenAsync(
                         refreshCache: false);
 
                     request.RequestContext.ResolvedPartitionKeyRange = partitionKeyRange;
-                    this.globalPartitionEndpointManager.TryAddPartitionLevelLocationOverride(request);
+
+                    if (this.isPartitionLevelFailoverEnabled)
+                    {
+                        this.globalPartitionEndpointManager.TryAddPartitionLevelLocationOverride(request);
+                    }
                 }
 
                 bool canUseThinClient =

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemThinClientTests.cs

@@ -13,6 +13,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     using System.Text.Json;
     using System.Text.Json.Serialization;
     using System.Threading.Tasks;
+    using global::Azure.Core;
     using Microsoft.Azure.Cosmos.Fluent;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using static Microsoft.Azure.Cosmos.SDK.EmulatorTests.MultiRegionSetupHelpers;
@@ -512,6 +513,83 @@ public async Task QueryItemsTest()
             Assert.AreEqual(createdItems.Count, count);
         }
 
+        [TestMethod]
+        [TestCategory("ThinClient")]
+        public async Task QueryItemsTestWithStrongConsistency()
+        {
+            string connectionString = ConfigurationManager.GetEnvironmentVariable<string>("COSMOSDB_TC_STRONG", string.Empty);
+            if (string.IsNullOrEmpty(connectionString))
+            {
+                Assert.Fail("Set environment variable COSMOSDB_TC_STRONG 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<TestObject> items = this.GenerateItems(pk).ToList();
+
+            List<TestObject> createdItems = await this.CreateItemsSafeAsync(items);
+
+            string query = $"SELECT * FROM c WHERE c.pk = '{pk}'";
+            FeedIterator<TestObject> iterator = this.container.GetItemQueryIterator<TestObject>(query);
+
+            int count = 0;
+            while (iterator.HasMoreResults)
+            {
+                FeedResponse<TestObject> response = await iterator.ReadNextAsync();
+                count += response.Count;
+            }
+
+            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");
+
+            string pk = "pk_query";
+            List<TestObject> items = this.GenerateItems(pk).ToList();
+
+            List<TestObject> createdItems = await this.CreateItemsSafeAsync(items);
+
+            string query = $"SELECT * FROM c WHERE c.pk = '{pk}'";
+            FeedIterator<TestObject> iterator = this.container.GetItemQueryIterator<TestObject>(query);
+
+            int count = 0;
+            while (iterator.HasMoreResults)
+            {
+                FeedResponse<TestObject> response = await iterator.ReadNextAsync();
+                count += response.Count;
+            }
+
+            Assert.AreEqual(createdItems.Count, count);
+        }
+
         [TestMethod]
         [TestCategory("ThinClient")]
         public async Task QueryItemsStreamTest()