TryExecuteQuery (#837)

* allowed for retrieval of query plan without exception * wired query plan through the stack * wired a try execute function * added test * resolved iteration comments * removed contracts * added GROUP BY support * resolved iteration comments * added try execute query * added back methods that require private members
Azure · Oct 24, 2019 · a10e411 · a10e411
1 parent ab4f1d7
commit a10e411
Show file tree

Hide file tree

Showing 8 changed files with 438 additions and 176 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs
@@ -156,50 +156,70 @@ private async Task<CosmosQueryExecutionContext> CreateItemQueryExecutionContextA
             this.CosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId;
 
             PartitionedQueryExecutionInfo partitionedQueryExecutionInfo;
-            if (this.CosmosQueryContext.QueryClient.ByPassQueryParsing())
+            if (this.inputParameters.PartitionedQueryExecutionInfo != null)
             {
-                // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case.
-                // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this
-                partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync(
-                    this.CosmosQueryContext.QueryClient,
-                    this.inputParameters.SqlQuerySpec,
-                    this.CosmosQueryContext.ResourceLink,
-                    cancellationToken);
+                partitionedQueryExecutionInfo = this.inputParameters.PartitionedQueryExecutionInfo;
             }
             else
             {
-                //todo:elasticcollections this may rely on information from collection cache which is outdated
-                //if collection is deleted/created with same name.
-                //need to make it not rely on information from collection cache.
-                Documents.PartitionKeyDefinition partitionKeyDefinition;
-                object partitionKeyDefinitionObject;
-                if (this.inputParameters.Properties != null
-                    && this.inputParameters.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out partitionKeyDefinitionObject))
+                if (this.CosmosQueryContext.QueryClient.ByPassQueryParsing())
                 {
-                    if (partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition)
+                    // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case.
+                    // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this
+                    partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync(
+                        this.CosmosQueryContext.QueryClient,
+                        this.inputParameters.SqlQuerySpec,
+                        this.CosmosQueryContext.ResourceLink,
+                        cancellationToken);
+                }
+                else
+                {
+                    //todo:elasticcollections this may rely on information from collection cache which is outdated
+                    //if collection is deleted/created with same name.
+                    //need to make it not rely on information from collection cache.
+                    Documents.PartitionKeyDefinition partitionKeyDefinition;
+                    object partitionKeyDefinitionObject;
+                    if (this.inputParameters.Properties != null
+                        && this.inputParameters.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out partitionKeyDefinitionObject))
                     {
-                        partitionKeyDefinition = definition;
+                        if (partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition)
+                        {
+                            partitionKeyDefinition = definition;
+                        }
+                        else
+                        {
+                            throw new ArgumentException(
+                                "partitionkeydefinition has invalid type",
+                                nameof(partitionKeyDefinitionObject));
+                        }
                     }
                     else
                     {
-                        throw new ArgumentException(
-                            "partitionkeydefinition has invalid type",
-                            nameof(partitionKeyDefinitionObject));
+                        partitionKeyDefinition = containerQueryProperties.PartitionKeyDefinition;
                     }
-                }
-                else
-                {
-                    partitionKeyDefinition = containerQueryProperties.PartitionKeyDefinition;
-                }
 
-                partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
-                    this.CosmosQueryContext.QueryClient,
-                    this.inputParameters.SqlQuerySpec,
-                    partitionKeyDefinition,
-                    this.inputParameters.PartitionKey != null,
-                    cancellationToken);
+                    partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
+                        this.CosmosQueryContext.QueryClient,
+                        this.inputParameters.SqlQuerySpec,
+                        partitionKeyDefinition,
+                        this.inputParameters.PartitionKey != null,
+                        cancellationToken);
+                }
             }
 
+            return await this.CreateFromPartitionedQuerExecutionInfoAsync(
+                partitionedQueryExecutionInfo,
+                containerQueryProperties,
+                cancellationToken);
+        }
+
+        public async Task<CosmosQueryExecutionContext> CreateFromPartitionedQuerExecutionInfoAsync(
+            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
+            ContainerQueryProperties containerQueryProperties,
+            CancellationToken cancellationToken = default(CancellationToken))
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
             List<Documents.PartitionKeyRange> targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(
                    this.CosmosQueryContext.QueryClient,
                    this.CosmosQueryContext.ResourceLink.OriginalString,
@@ -415,6 +435,7 @@ public struct InputParameters
             internal int? MaxBufferedItemCount { get; set; }
             internal PartitionKey? PartitionKey { get; set; }
             internal IDictionary<string, object> Properties { get; set; }
+            internal PartitionedQueryExecutionInfo PartitionedQueryExecutionInfo { get; set; }
         }
     }
 }