From 63323e54bc8be54f0d8d9296dc9ff90773d71f63 Mon Sep 17 00:00:00 2001
From: Asket Agarwal <asket.agarwal96@gmail.com>
Date: Thu, 10 Jun 2021 23:55:00 +0530
Subject: [PATCH] Fixing Container Recreate scenarios for Gateway Mode

---
 .../src/GatewayStoreModel.cs                  | 15 ++++++----
 .../CosmosItemTests.cs                        | 28 +++++++++++++++++++
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs
index 6cc73ab350..c3e2f723d8 100644
--- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs
+++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs
@@ -358,12 +358,15 @@ private static async Task<Tuple<bool, PartitionKeyRange>> TryResolvePartitionKey
                         NoOpTrace.Singleton);
                 }
 
-                partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(
-                    request: request,
-                    partitionKeyString: partitionKeyString,
-                    collectionCacheUptoDate: false,
-                    collection: collection,
-                    routingMap: collectionRoutingMap);
+                if (collectionRoutingMap != null)
+                {
+                    partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(
+                        request: request,
+                        partitionKeyString: partitionKeyString,
+                        collectionCacheUptoDate: false,
+                        collection: collection,
+                        routingMap: collectionRoutingMap);
+                }
             }
             else if (request.PartitionKeyRangeIdentity != null)
             {
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs
index 94f807e95a..ebb8ba4269 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs
@@ -1873,6 +1873,34 @@ public async Task PatchItemStreamTest()
             }
         }
 
+        [TestMethod]
+        public async Task ContainerRecreateScenarioGatewayTest()
+        {
+            ContainerResponse response = await this.database.CreateContainerAsync(
+                        new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: "/pk"));
+
+            Container createdContainer = (ContainerInlineCore)response;
+
+            CosmosClient client1 = TestCommon.CreateCosmosClient(useGateway: true);
+            CosmosClient client2 = TestCommon.CreateCosmosClient(useGateway: true);
+
+            Container container1 = client1.GetContainer(this.database.Id, createdContainer.Id);
+            Container container2 = client2.GetContainer(this.database.Id, createdContainer.Id);
+            Cosmos.Database database2 = client2.GetDatabase(this.database.Id);
+
+            ToDoActivity item = ToDoActivity.CreateRandomToDoActivity("pk2002", "id2002");
+            await container1.CreateItemAsync<ToDoActivity>(item);
+
+            await container2.DeleteContainerAsync();
+            await database2.CreateContainerAsync(createdContainer.Id, "/pk");
+
+            container2 = database2.GetContainer(this.Container.Id);
+            await container2.CreateItemAsync<ToDoActivity>(item);
+
+            // should not throw exception
+            await this.Container.ReadItemAsync<ToDoActivity>("id2002", new Cosmos.PartitionKey("pk2002"));
+        }
+
         [Ignore]
         [TestMethod]
         public async Task BatchPatchConditionTest()