Adding a unit test to verify setting PartitionKeyDefinition works.

Microsoft.Azure.Cosmos/src/Query/CosmosQueryExecutionContextFactory.cs

@@ -145,7 +145,7 @@ private async Task<CosmosQueryExecutionContext> CreateItemQueryExecutionContextA
             //need to make it not rely on information from collection cache.
             PartitionKeyDefinition partitionKeyDefinition;
             object partitionKeyDefinitionObject;
-            if (this.cosmosQueryContext.QueryRequestOptions != null
+            if (this.cosmosQueryContext.QueryRequestOptions?.Properties != null
                 && this.cosmosQueryContext.QueryRequestOptions.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out partitionKeyDefinitionObject))
             {
                 if (partitionKeyDefinitionObject is PartitionKeyDefinition definition)
@@ -164,8 +164,8 @@ private async Task<CosmosQueryExecutionContext> CreateItemQueryExecutionContextA
                 partitionKeyDefinition = collection.PartitionKey;
             }
 
-            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync(
-                queryClient: this.cosmosQueryContext.QueryClient,
+            // $ISSUE-felixfan-2016-07-13: We should probably get PartitionedQueryExecutionInfo from Gateway in GatewayMode
+            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await this.cosmosQueryContext.QueryClient.GetPartitionedQueryExecutionInfoAsync(
                 sqlQuerySpec: this.cosmosQueryContext.SqlQuerySpec,
                 partitionKeyDefinition: partitionKeyDefinition,
                 requireFormattableOrderByQuery: true,
@@ -301,7 +301,7 @@ internal static async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRanges(
             return targetRanges;
         }
 
-        public static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInfoAsync(
+        public static Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInfoAsync(
             CosmosQueryClient queryClient,
             SqlQuerySpec sqlQuerySpec,
             PartitionKeyDefinition partitionKeyDefinition,
@@ -312,12 +312,13 @@ public static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecu
         {
             // $ISSUE-felixfan-2016-07-13: We should probably get PartitionedQueryExecutionInfo from Gateway in GatewayMode
 
-            QueryPartitionProvider queryPartitionProvider = await queryClient.GetQueryPartitionProviderAsync(cancellationToken);
-            return queryPartitionProvider.GetPartitionedQueryExecutionInfo(sqlQuerySpec, 
-                partitionKeyDefinition, 
-                requireFormattableOrderByQuery, 
-                isContinuationExpected, 
-                allowNonValueAggregateQuery);
+            return queryClient.GetPartitionedQueryExecutionInfoAsync(
+                sqlQuerySpec,
+                partitionKeyDefinition,
+                requireFormattableOrderByQuery,
+                isContinuationExpected,
+                allowNonValueAggregateQuery,
+                cancellationToken);
         }
 
 

Microsoft.Azure.Cosmos/src/Resource/Query/CosmosQueryClient.cs

@@ -23,7 +23,13 @@ internal abstract class CosmosQueryClient
 
         internal abstract Task<IRoutingMapProvider> GetRoutingMapProviderAsync();
 
-        internal abstract Task<QueryPartitionProvider> GetQueryPartitionProviderAsync(CancellationToken cancellationToken);
+        internal abstract Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInfoAsync(
+            SqlQuerySpec sqlQuerySpec,
+            PartitionKeyDefinition partitionKeyDefinition,
+            bool requireFormattableOrderByQuery,
+            bool isContinuationExpected,
+            bool allowNonValueAggregateQuery,
+            CancellationToken cancellationToken);
 
         internal abstract Task<CosmosQueryResponse> ExecuteItemQueryAsync(
             Uri resourceUri,

Microsoft.Azure.Cosmos/src/Resource/Query/CosmosQueryClientCore.cs

@@ -50,9 +50,21 @@ internal override Task<IRoutingMapProvider> GetRoutingMapProviderAsync()
             return this.DocumentQueryClient.GetRoutingMapProviderAsync();
         }
 
-        internal override Task<QueryPartitionProvider> GetQueryPartitionProviderAsync(CancellationToken cancellationToken)
+        internal override async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInfoAsync(
+            SqlQuerySpec sqlQuerySpec,
+            PartitionKeyDefinition partitionKeyDefinition,
+            bool requireFormattableOrderByQuery,
+            bool isContinuationExpected,
+            bool allowNonValueAggregateQuery,
+            CancellationToken cancellationToken)
         {
-            return this.DocumentQueryClient.GetQueryPartitionProviderAsync(cancellationToken);
+            QueryPartitionProvider queryPartitionProvider = await this.DocumentQueryClient.GetQueryPartitionProviderAsync(cancellationToken);
+            return queryPartitionProvider.GetPartitionedQueryExecutionInfo(
+                sqlQuerySpec,
+                partitionKeyDefinition,
+                requireFormattableOrderByQuery,
+                isContinuationExpected,
+                allowNonValueAggregateQuery);
         }
 
         internal override async Task<CosmosQueryResponse> ExecuteItemQueryAsync(

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs

@@ -8,8 +8,10 @@ namespace Microsoft.Azure.Cosmos.Tests
     using System.Collections.Generic;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Common;
     using Microsoft.Azure.Cosmos.Query;
     using Microsoft.Azure.Cosmos.Query.ExecutionComponent;
+    using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using Moq;
 
@@ -49,6 +51,52 @@ public async Task TestCosmosQueryExecutionComponentCancellation()
             }
         }
 
+        [TestMethod]
+        [ExpectedException(typeof(InvalidOperationException))]
+        public async Task TestCosmosQueryPartitionKeyDefinition()
+        {
+            PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition();
+            CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions();
+            queryRequestOptions.Properties = new Dictionary<string, object>()
+            {
+                {"x-ms-query-partitionkey-definition", partitionKeyDefinition }
+            };
+
+            SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(@"select * from t where t.something = 42 ");
+            bool allowNonValueAggregateQuery = true;
+            bool isContinuationExpected = true;
+            CancellationTokenSource cancellation = new CancellationTokenSource();
+            CancellationToken token = cancellation.Token;
+
+            Mock<CollectionCache> mockCollectionCache = new Mock<CollectionCache>();
+            mockCollectionCache.Setup(x => x.ResolveCollectionAsync(It.IsAny<DocumentServiceRequest>(), token)).Returns(Task.FromResult(new CosmosContainerSettings("mockContainer", "/pk")));
+
+            Mock<CosmosQueryClient> client = new Mock<CosmosQueryClient>();
+            client.Setup(x => x.GetCollectionCacheAsync()).Returns(Task.FromResult(mockCollectionCache.Object));
+            client.Setup(x => x.ByPassQueryParsing()).Returns(false);
+            client.Setup(x => x.GetPartitionedQueryExecutionInfoAsync(
+                sqlQuerySpec,
+                partitionKeyDefinition,
+                true,
+                isContinuationExpected,
+                allowNonValueAggregateQuery,
+                token)).Throws(new InvalidOperationException("Verified that the PartitionKeyDefinition was correctly set. Cancel the rest of the query"));
+
+            CosmosQueryExecutionContextFactory factory = new CosmosQueryExecutionContextFactory(
+                client: client.Object,
+                resourceTypeEnum: ResourceType.Document,
+                operationType: OperationType.Query,
+                resourceType: typeof(CosmosQueryResponse),
+                sqlQuerySpec: sqlQuerySpec,
+                queryRequestOptions: queryRequestOptions,
+                resourceLink: new Uri("dbs/mockdb/colls/mockColl", UriKind.Relative),
+                isContinuationExpected: isContinuationExpected,
+                allowNonValueAggregateQuery: allowNonValueAggregateQuery,
+                correlatedActivityId: new Guid("221FC86C-1825-4284-B10E-A6029652CCA6"));
+
+            await factory.ExecuteNextAsync(token);
+        }
+
         private async Task<(IList<DocumentQueryExecutionComponentBase> components, CosmosQueryResponse response)> GetAllExecutionComponents()
         {
             (Func<string, Task<IDocumentQueryExecutionComponent>> func, CosmosQueryResponse response) setupContext = this.SetupBaseContextToVerifyFailureScenario();