Client Encryption: Adds fix for supporting Prefix Partition Key (Hierarchical partitioning) (#3979)

* Hirarchical pk bug fix

* Hirarchical pk bug fix

* Hirarchical pk bug fix

* Hirarchical pk bug fix

* Hirarchical pk bug fix

* testing new version

* adding more tests

* adding more tests

* adding more tests

* code review changes

* test fix

* test fix

* test fix

* test fix

---------

Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com>

Directory.Build.props

@@ -4,8 +4,8 @@
 		<ClientPreviewVersion>3.35.2</ClientPreviewVersion>
 		<ClientPreviewSuffixVersion>preview</ClientPreviewSuffixVersion>
 		<DirectVersion>3.31.3</DirectVersion>
-		<EncryptionOfficialVersion>2.0.2</EncryptionOfficialVersion>
+		<EncryptionOfficialVersion>2.0.3</EncryptionOfficialVersion>
-		<EncryptionPreviewVersion>2.0.2</EncryptionPreviewVersion>
+		<EncryptionPreviewVersion>2.0.3</EncryptionPreviewVersion>
 		<EncryptionPreviewSuffixVersion>preview</EncryptionPreviewSuffixVersion>
 		<CustomEncryptionVersion>1.0.0-preview06</CustomEncryptionVersion>
 		<HybridRowVersion>1.1.0-preview3</HybridRowVersion>

Microsoft.Azure.Cosmos.Encryption/changelog.md

@@ -3,6 +3,16 @@ Preview features are treated as a separate branch and will not be included in th
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+### <a name="2.0.3"/> [2.0.3](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.3) - 2023-07-12
+
+#### Added
+- [#3979](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3979) Adds fix for supporting Prefix Partition Key (Hierarchical partitioning).
+
+### <a name="2.0.3-preview"/> [2.0.3-preview](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.3-preview) - 2023-07-12
+
+#### Added
+- [#3979](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3979) Adds fix for supporting Prefix Partition Key (Hierarchical partitioning).
+
 ### <a name="2.0.2"/> [2.0.2](https://www.nuget.org/packages/Microsoft.Azure.Cosmos.Encryption/2.0.2) - 2023-06-01
 
 #### Added

Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs

@@ -957,15 +957,22 @@ namespace Microsoft.Azure.Cosmos.Encryption
             JArray jArray = JArray.Parse(partitionKey.ToString());
 
 #if ENCRYPTIONPREVIEW
-            if (jArray.Count > 1)
+            if (encryptionSettings.PartitionKeyPaths.Count > 1)
             {
-                int i = 0;
+                int counter = 0;
                 PartitionKeyBuilder partitionKeyBuilder = new PartitionKeyBuilder();
 
-                bool isPkEncrypted = false;
+                if (jArray.Count() > encryptionSettings.PartitionKeyPaths.Count())
-                // partitionKeyBuilder expects the paths and values to be in same order.
-                foreach (string path in encryptionSettings.PartitionKeyPaths)
                 {
+                    throw new NotSupportedException($"The number of partition keys passed in the query exceeds the number of keys initialized on the container. Container Id : {this.Id}");
+                }
+                bool isPkEncrypted = false;
+
+                // partitionKeyBuilder expects the paths and values to be in same order.
+                for(counter = 0; counter < jArray.Count(); counter++)
+                {
+                    string path = encryptionSettings.PartitionKeyPaths[counter];
+
                     // case: partition key path is /a/b/c and the client encryption policy has /a in path.
                     // hence encrypt the partition key value with using its top level path /a since /c would have been encrypted in the document using /a's policy.
                     string partitionKeyPath = path.Split('/')[1];
@@ -975,12 +982,12 @@ namespace Microsoft.Azure.Cosmos.Encryption
 
                     if (encryptionSettingForProperty == null)
                     {
-                        partitionKeyBuilder.Add(jArray[i++].ToString());
+                        partitionKeyBuilder.Add(jArray[counter].ToString());
                         continue;
                     }
 
                     isPkEncrypted = true;
-                    Stream valueStream = EncryptionProcessor.BaseSerializer.ToStream(jArray[i++]);
+                    Stream valueStream = EncryptionProcessor.BaseSerializer.ToStream(jArray[counter]);
 
                     Stream encryptedPartitionKey = await EncryptionProcessor.EncryptValueStreamAsync(
                         valueStreamToEncrypt: valueStream,

Microsoft.Azure.Cosmos.Encryption/src/Microsoft.Azure.Cosmos.Encryption.csproj

@@ -28,11 +28,11 @@
   </ItemGroup>
 	
   <ItemGroup Condition=" '$(SdkProjectRef)' != 'True' AND '$(IsPreview)' != 'True' ">
-    <PackageReference Include="Microsoft.Azure.Cosmos" Version="[3.31.0,3.34.0]" />
+    <PackageReference Include="Microsoft.Azure.Cosmos" Version="[3.31.0,3.35.2]" />
   </ItemGroup>
 
   <ItemGroup Condition=" '$(SdkProjectRef)' != 'True' AND '$(IsPreview)' == 'True' ">
-    <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.34.0-preview" />
+    <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.35.2-preview" />
   </ItemGroup>
 
   <ItemGroup Condition=" '$(SdkProjectRef)' == 'True' ">

Microsoft.Azure.Cosmos.Encryption/tests/EmulatorTests/MdeEncryptionTests.cs

@@ -189,6 +189,9 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests
         {
             // Reset static cache TTL
             Microsoft.Data.Encryption.Cryptography.ProtectedDataEncryptionKey.TimeToLive = TimeSpan.FromHours(2);
+            // flag to disable https://github.com/Azure/azure-cosmos-dotnet-v3/pull/3951
+            // need to be removed after the fix
+            Environment.SetEnvironmentVariable("AZURE_COSMOS_REPLICA_VALIDATION_ENABLED", "False");
         }
 
         private static async Task<ClientEncryptionKeyResponse> CreateClientEncryptionKeyAsync(string cekId, Cosmos.EncryptionKeyWrapMetadata encryptionKeyWrapMetadata)
@@ -2303,7 +2306,7 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests
             VerifyExpectedDocResponse(testDoc, readResponse.Resource);
 
 #if ENCRYPTIONTESTPREVIEW
-            // hierarchical
+            // hierarchical pk container test
             cepWithPKIdPath1 = new ClientEncryptionIncludedPath()
             {
                 Path = "/Sensitive_LongFormat",
@@ -2352,9 +2355,320 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests
 
             Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
             VerifyExpectedDocResponse(testDoc, readResponse.Resource);
+
+            // test to validate query with one partition key (topmost) in hierarchical pk container of 3 keys
+            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = new PartitionKeyBuilder().Add(testDoc.Sensitive_StringFormat).Build()
+            };
+
+            using FeedIterator<TestDoc> setIterator = encryptionContainer.GetItemQueryIterator<TestDoc>("select * from c", requestOptions: queryRequestOptions);
+
+            while (setIterator.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            // test to validate query with one partition key (topmost) in hierarchical pk container of 3 keys with where clause on topmost pk
+            QueryDefinition queryDefinition = encryptionContainer.CreateQueryDefinition("SELECT * FROM c WHERE c.Sensitive_StringFormat = @Sensitive_StringFormat");
+
+            await queryDefinition.AddParameterAsync("@Sensitive_StringFormat", testDoc.Sensitive_StringFormat, "/Sensitive_StringFormat");
+
+            FeedIterator<TestDoc> setIteratorWithFilter = encryptionContainer.GetItemQueryIterator<TestDoc>(queryDefinition, requestOptions: queryRequestOptions);
+
+            while (setIteratorWithFilter.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIteratorWithFilter.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            // test to validate query with one partition key (2nd topmost) in hierarchical pk container of 3 keys with where clause on topmost pk
+            // this shold give 0 items as PK is set wrongly
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = new PartitionKeyBuilder().Add(testDoc.Sensitive_NestedObjectFormatL1.Sensitive_NestedObjectFormatL2.Sensitive_StringFormatL2).Build()
+            };
+
+            setIteratorWithFilter = encryptionContainer.GetItemQueryIterator<TestDoc>(queryDefinition, requestOptions: queryRequestOptions);
+
+            while (setIteratorWithFilter.HasMoreResults)
+            {
+                FeedResponse<TestDoc> response = await setIteratorWithFilter.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                Assert.AreEqual(0, response.Count());
+            }
 #endif
         }
+#if ENCRYPTIONTESTPREVIEW
+        [TestMethod]
+        public async Task TestHirarchicalPkWithFullAndPartialKey() 
+        {
+            HirarchicalPkTestDoc testDoc = HirarchicalPkTestDoc.Create();
 
+            ClientEncryptionIncludedPath cepWithPKIdPath1 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/State",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath2 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/City",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath3 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/ZipCode",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            Collection<ClientEncryptionIncludedPath> paths = new Collection<ClientEncryptionIncludedPath> { cepWithPKIdPath1, cepWithPKIdPath2, cepWithPKIdPath3 };
+
+            ClientEncryptionPolicy clientEncryptionPolicy= new ClientEncryptionPolicy(paths, 2);
+
+            ContainerProperties containerProperties = new ContainerProperties()
+            {
+                Id = "HierarchicalPkContainerWith3Pk",
+                PartitionKeyPaths = new List<string> { "/State", "/City", "/ZipCode" },
+                ClientEncryptionPolicy = clientEncryptionPolicy
+            };
+
+            Container encryptionContainer = await database.CreateContainerAsync(containerProperties, 400);
+            await encryptionContainer.InitializeEncryptionAsync();
+
+            PartitionKey hirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Add(testDoc.City)
+                .Add(testDoc.ZipCode)
+                .Build();
+
+            ItemResponse<HirarchicalPkTestDoc> createResponse = await encryptionContainer.CreateItemAsync(
+                testDoc,
+                partitionKey: hirarchicalPk);
+            Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, createResponse.Resource);
+
+            // read back
+            ItemResponse<HirarchicalPkTestDoc> readResponse = await encryptionContainer.ReadItemAsync<HirarchicalPkTestDoc>(
+               testDoc.Id,
+               hirarchicalPk);
+
+            Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, readResponse.Resource);
+
+            PartitionKey partialHirarchicalPk = new PartitionKeyBuilder()
+               .Add(testDoc.State)
+               .Add(testDoc.City)
+               .Build();
+
+            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            using FeedIterator<HirarchicalPkTestDoc> setIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>("select * from c", requestOptions: queryRequestOptions);
+
+            while (setIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await setIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            QueryDefinition withEncryptedParameter = encryptionContainer.CreateQueryDefinition(
+                    "SELECT * FROM c WHERE c.City = @cityInput AND c.State = @stateInput");
+
+            await withEncryptedParameter.AddParameterAsync(
+                    "@cityInput",
+                    testDoc.City,
+                    "/City");
+
+            await withEncryptedParameter.AddParameterAsync(
+                    "@stateInput",
+                    testDoc.State,
+                    "/State");
+
+            // query with partial HirarchicalPk  state and city
+            FeedIterator<HirarchicalPkTestDoc> queryResponseIterator;
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+
+            while (queryResponseIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await queryResponseIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Build();
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            // query with partial HirarchicalPk  state
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+
+            while (queryResponseIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await queryResponseIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.ZipCode)
+                .Build();
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            // query with partial HirarchicalPk  zipCode.
+            // Since zipCode is 3rd in HirarchicalPk set. Query will get 0 response.
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+
+            while (queryResponseIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await queryResponseIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                Assert.AreEqual(0, response.Count());
+            }
+
+            // query with no HirarchicalPk set.
+            queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter);
+
+            while (queryResponseIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await queryResponseIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+
+            partialHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Add(testDoc.City)
+                .Add(testDoc.ZipCode)
+                .Add("Extra Value")
+                .Build();
+
+            queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = partialHirarchicalPk
+            };
+
+            // query with more PKs greater than number of PK feilds set in the container settings.
+            try
+            {
+                queryResponseIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>(withEncryptedParameter, requestOptions: queryRequestOptions);
+                while (queryResponseIterator.HasMoreResults)
+                {
+                    FeedResponse<HirarchicalPkTestDoc> response = await queryResponseIterator.ReadNextAsync().ConfigureAwait(false);
+                    Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                    Assert.AreEqual(0, response.Count());
+                }
+            }
+            catch (Exception ex)
+            {
+                Assert.IsTrue(ex is NotSupportedException);
+                if (ex is NotSupportedException notSupportedException)
+                    Assert.IsTrue(notSupportedException.Message.Contains("The number of partition keys passed in the query exceeds the number of keys initialized on the container"));
+            }
+        }
+
+        [TestMethod]
+        public async Task TestHirarchicalPkWithOnlyOneKey()
+        {
+            HirarchicalPkTestDoc testDoc = HirarchicalPkTestDoc.Create();
+
+            ClientEncryptionIncludedPath cepWithPKIdPath1 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/State",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath2 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/City",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            ClientEncryptionIncludedPath cepWithPKIdPath3 = new ClientEncryptionIncludedPath()
+            {
+                Path = "/ZipCode",
+                ClientEncryptionKeyId = "key1",
+                EncryptionType = "Deterministic",
+                EncryptionAlgorithm = "AEAD_AES_256_CBC_HMAC_SHA256",
+            };
+
+            Collection<ClientEncryptionIncludedPath> paths = new Collection<ClientEncryptionIncludedPath> { cepWithPKIdPath1, cepWithPKIdPath2, cepWithPKIdPath3 };
+
+            ClientEncryptionPolicy clientEncryptionPolicy = new ClientEncryptionPolicy(paths, 2);
+
+            ContainerProperties containerProperties = new ContainerProperties()
+            {
+                Id = "HierarchicalPkContainerWithOnePk",
+                PartitionKeyPaths = new List<string> { "/State" },
+                ClientEncryptionPolicy = clientEncryptionPolicy
+            };
+
+            Container encryptionContainer = await database.CreateContainerAsync(containerProperties, 400);
+            await encryptionContainer.InitializeEncryptionAsync();
+
+            PartitionKey hirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Build();
+
+            ItemResponse<HirarchicalPkTestDoc> createResponse = await encryptionContainer.CreateItemAsync(
+                testDoc,
+                partitionKey: hirarchicalPk);
+            Assert.AreEqual(HttpStatusCode.Created, createResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, createResponse.Resource);
+
+            // read back
+            ItemResponse<HirarchicalPkTestDoc> readResponse = await encryptionContainer.ReadItemAsync<HirarchicalPkTestDoc>(
+               testDoc.Id,
+               hirarchicalPk);
+
+            Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
+            VerifyExpectedDocResponse(testDoc, readResponse.Resource);
+
+            PartitionKey fullHirarchicalPk = new PartitionKeyBuilder()
+                .Add(testDoc.State)
+                .Build();
+
+            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
+            {
+                PartitionKey = fullHirarchicalPk
+            };
+
+            using FeedIterator<HirarchicalPkTestDoc> setIterator = encryptionContainer.GetItemQueryIterator<HirarchicalPkTestDoc>("select * from c", requestOptions: queryRequestOptions);
+
+            while (setIterator.HasMoreResults)
+            {
+                FeedResponse<HirarchicalPkTestDoc> response = await setIterator.ReadNextAsync().ConfigureAwait(false);
+                Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                VerifyExpectedDocResponse(testDoc, response.First());
+            }
+        }
+#endif
         [TestMethod]
         public async Task EncryptionStreamIteratorValidation()
         {
@@ -3204,6 +3518,14 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests
             return deleteResponse;
         }
 
+        private static void VerifyExpectedDocResponse(HirarchicalPkTestDoc expectedDoc, HirarchicalPkTestDoc verifyDoc)
+        {
+            Assert.AreEqual(expectedDoc.Id, verifyDoc.Id);
+            Assert.AreEqual(expectedDoc.State, verifyDoc.State);
+            Assert.AreEqual(expectedDoc.City, verifyDoc.City);
+            Assert.AreEqual(expectedDoc.ZipCode, verifyDoc.ZipCode);
+        }
+
         private static void VerifyExpectedDocResponse(TestDoc expectedDoc, TestDoc verifyDoc)
         {
             Assert.AreEqual(expectedDoc.Id, verifyDoc.Id);
@@ -3716,6 +4038,70 @@ namespace Microsoft.Azure.Cosmos.Encryption.EmulatorTests
             }
         }
 
+        public class HirarchicalPkTestDoc
+        {
+            [JsonProperty("id")]
+            public string Id { get; set; }
+
+            public string PK { get; set; }
+
+            public string State { get; set; }
+
+            public string City { get; set; }
+
+            public string ZipCode { get; set; }
+
+            public HirarchicalPkTestDoc()
+            {
+            }
+            public HirarchicalPkTestDoc(HirarchicalPkTestDoc other)
+            {
+                this.Id = other.Id;
+                this.PK = other.PK;
+                this.State = other.State;
+                this.City = other.City;
+                this.ZipCode = other.ZipCode;
+            }
+
+            public override bool Equals(object obj)
+            {
+                return obj is HirarchicalPkTestDoc doc
+                       && this.Id == doc.Id
+                       && this.PK == doc.PK
+                       && this.State == doc.State
+                       && this.City == doc.City
+                       && this.ZipCode == doc.ZipCode;
+            }
+
+            public override int GetHashCode()
+            {
+                int hashCode = 1652434776;
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.Id);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.PK);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.State);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.City);
+                hashCode = (hashCode * -1521134295) + EqualityComparer<string>.Default.GetHashCode(this.ZipCode);
+                return hashCode;
+            }
+
+            public static HirarchicalPkTestDoc Create(string partitionKey = null)
+            {
+                return new HirarchicalPkTestDoc()
+                {
+                    Id = Guid.NewGuid().ToString(),
+                    PK = partitionKey ?? Guid.NewGuid().ToString(),
+                    State = Guid.NewGuid().ToString(),
+                    City = Guid.NewGuid().ToString(),
+                    ZipCode = Guid.NewGuid().ToString()
+                };
+            }
+
+            public Stream ToStream()
+            {
+                return TestCommon.ToStream(this);
+            }
+        }
+
         internal class TestKeyEncryptionKey : IKeyEncryptionKey
         {
             private static readonly Dictionary<string, int> keyinfo = new Dictionary<string, int>