[Internal] Query: Adds IsPassThrough option to SqlQuerySpec (#3095)

* Draft PR for Pass Through Query Project * Updated code for Pass Through query on SDK side. * Fixed comments * Added Assert to check if queryPipelineStage is null or not * Added new baseline tests to cover Aggregare, Distinct and No Partition Key provided cases. * Fixed stylistic comments * Fixed comments pt. 2 * Added tests for SqlQueryOptions serialization and deserialization. * Added serialization test when options = false * Updated Linq baseline tests * Updated Linq Baseline Tests pt2 * Updated Linq Baseline Tests pt3 * Updated singleLogicalPartitionKeyQuery boolean * LinqGeneralBaselineTests * Removed error message from linq baseline tests * Added function to check if single partition. Updated test files. Removed SqlQueryOptions. Updated SqlQuerySpec * Added PassThrough check before getting QueryPlan * Added implemenation for ExecuteQueryPlanRequestAsync() * Updated the ExecuteQueryPlanRequestAsync() to include the new parameter * Updated variable names * Added tests for partition key value * Removed unused function * Removed null checker in ToJsonString() in PartitionKey.cs Co-authored-by: j82w <j82w@users.noreply.github.com>
Azure · Apr 29, 2022 · 3632f3a · 3632f3a
1 parent 9411d26
commit 3632f3a
Show file tree

Hide file tree

Showing 7 changed files with 538 additions and 9 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
@@ -131,6 +131,13 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                     cancellationToken);
                 cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId;
 
+                inputParameters.SqlQuerySpec.PassThrough = IsPassThroughCandidate(inputParameters, queryPlanFromContinuationToken);
+
+                if (inputParameters.SqlQuerySpec.PassThrough)
+                {
+                    //TODO: Add new pass through pipeline code here 
+                }
+
                 PartitionedQueryExecutionInfo partitionedQueryExecutionInfo;
                 if (inputParameters.ForcePassthrough)
                 {
@@ -232,7 +239,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                     }
                 }
 
-                return await TryCreateFromPartitionedQuerExecutionInfoAsync(
+                return await TryCreateFromPartitionedQueryExecutionInfoAsync(
                     documentContainer,
                     partitionedQueryExecutionInfo,
                     containerQueryProperties,
@@ -243,7 +250,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
             }
         }
 
-        public static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitionedQuerExecutionInfoAsync(
+        public static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitionedQueryExecutionInfoAsync(
             DocumentContainer documentContainer,
             PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
             ContainerQueryProperties containerQueryProperties,
@@ -265,7 +272,7 @@ public static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitioned
 
             bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue
                 || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1)
-                    && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue);
+                && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue);
             bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates
                 && !partitionedQueryExecutionInfo.QueryInfo.HasDistinct
                 && !partitionedQueryExecutionInfo.QueryInfo.HasGroupBy;
@@ -279,10 +286,19 @@ public static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitioned
                 && !partitionedQueryExecutionInfo.QueryInfo.HasOffset;
             bool streamingCrossContinuationQuery = !singleLogicalPartitionKeyQuery && clientStreamingQuery;
 
-            bool createPassthoughQuery = streamingSinglePartitionQuery || streamingCrossContinuationQuery;
-
+            bool createPassthroughQuery = streamingSinglePartitionQuery || streamingCrossContinuationQuery;
+            
             TryCatch<IQueryPipelineStage> tryCreatePipelineStage;
-            if (createPassthoughQuery)
+
+            // After getting the Query Plan if we find out that the query is single logical partition, then short circuit and send straight to Backend
+            inputParameters.SqlQuerySpec.PassThrough = IsPassThroughCandidate(inputParameters, partitionedQueryExecutionInfo);
+
+            if (inputParameters.SqlQuerySpec.PassThrough)
+            {
+                //TODO: Add new pass through pipeline code here 
+            }
+
+            if (createPassthroughQuery)
             {
                 TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
                 if (responseStats != null)
@@ -532,6 +548,29 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             return partitionKeyDefinition;
         }
 
+        private static bool IsPassThroughCandidate(InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo)
+        {
+            // case 1: Is query going to a single partition
+            bool hasPartitionKey = inputParameters.PartitionKey.HasValue
+                && inputParameters.PartitionKey != PartitionKey.Null
+                && inputParameters.PartitionKey != PartitionKey.None;
+
+            if (hasPartitionKey) return true;
+
+            // case 2: does query execution plan have a single query range
+            if (partitionedQueryExecutionInfo != null)
+            {
+                bool hasQueryRanges = (partitionedQueryExecutionInfo.QueryRanges.Count == 1)
+                && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue;
+
+                if (hasQueryRanges) return hasQueryRanges;
+            }
+
+            // TODO: case 3: does collection have only one physical partition
+
+            return false;
+        }
+
         public sealed class InputParameters
         {
             private const int DefaultMaxConcurrency = 0;

diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/SqlQuerySpec.cs b/Microsoft.Azure.Cosmos/src/Query/Core/SqlQuerySpec.cs
@@ -14,7 +14,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core
     internal sealed class SqlQuerySpec
     {
         private SqlParameterCollection parameters;
-
+        private bool passThrough;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="T:Microsoft.Azure.Documents.SqlQuerySpec"/> class for the Azure Cosmos DB service.</summary>
         /// <remarks> 
@@ -40,11 +41,23 @@ public SqlQuerySpec(string queryText)
         /// <param name="queryText">The text of the database query.</param>
         /// <param name="parameters">The <see cref="T:Microsoft.Azure.Documents.SqlParameterCollection"/> instance, which represents the collection of query parameters.</param>
         public SqlQuerySpec(string queryText, SqlParameterCollection parameters)
+            : this(queryText, parameters, false)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="T:Microsoft.Azure.Documents.SqlQuerySpec"/> class for the Azure Cosmos DB service.
+        /// </summary>
+        /// <param name="queryText">The text of the database query.</param>
+        /// <param name="parameters">The <see cref="T:Microsoft.Azure.Documents.SqlParameterCollection"/> instance, which represents the collection of query parameters.</param>
+        /// <param name="passThrough">A Boolean value that indicates whether the query should be executed as pass-through.</param>
+        public SqlQuerySpec(string queryText, SqlParameterCollection parameters, bool passThrough)
         {
             this.QueryText = queryText;
             this.parameters = parameters ?? throw new ArgumentNullException("parameters");
+            this.passThrough = passThrough;
         }
-
+        
         /// <summary>
         /// Gets or sets the text of the Azure Cosmos DB database query.
         /// </summary>
@@ -69,12 +82,37 @@ public SqlParameterCollection Parameters
             }
         }
 
+        /// <summary>
+        /// Gets or sets the boolean value that indicates whether the query should be executed as pass-through..
+        /// </summary>
+        /// <value>A boolean value that indicates whether the query should be executed as pass-through.</value>
+        [DataMember(Name = "passThrough")]        
+        public bool PassThrough
+        {
+            get
+            {
+                return this.passThrough;
+            }
+            set
+            {
+                this.passThrough = value;
+            }
+        }
+
         /// <summary>
         /// Returns a value that indicates whether the Azure Cosmos DB database <see cref="Parameters"/> property should be serialized.
         /// </summary>
         public bool ShouldSerializeParameters()
         {
             return this.parameters.Count > 0;
         }
+
+        /// <summary>
+        /// Returns a value that indicates whether the Azure Cosmos DB database <see cref="PassThrough"/> property should be serialized.
+        /// </summary>
+        public bool ShouldSerializePassThrough()
+        {
+            return this.passThrough;
+        }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientTests.cs
@@ -393,6 +393,24 @@ public void SqlQuerySpecSerializationTest()
                     QueryText = "SELECT 1",
                     Parameters = new SqlParameterCollection() { new SqlParameter("@p1", new JRaw("{\"a\":[1,2,3]}")) }
                 });
+            verifyJsonSerialization("{\"query\":\"SELECT 1\",\"parameters\":[" +
+                    "{\"name\":\"@p1\",\"value\":{\"a\":[1,2,3]}}" +
+                "]}",
+                new SqlQuerySpec()
+                {
+                    QueryText = "SELECT 1",
+                    Parameters = new SqlParameterCollection() { new SqlParameter("@p1", new JRaw("{\"a\":[1,2,3]}")) },
+                    PassThrough = false
+                });
+            verifyJsonSerialization("{\"query\":\"SELECT 1\",\"parameters\":[" +
+                    "{\"name\":\"@p1\",\"value\":{\"a\":[1,2,3]}}" + 
+                "]," + "\"passThrough\":true}",
+                new SqlQuerySpec()
+                {
+                    QueryText = "SELECT 1",
+                    Parameters = new SqlParameterCollection() { new SqlParameter("@p1", new JRaw("{\"a\":[1,2,3]}")) },
+                    PassThrough = true
+                });
 
             // Verify roundtrips
             verifyJsonSerializationText("{\"query\":null}");
@@ -416,7 +434,7 @@ public void SqlQuerySpecSerializationTest()
                     "\"query\":\"SELECT 1\"," +
                     "\"parameters\":[" +
                         "{\"name\":\"@p1\",\"value\":true}" +
-                    "]" +
+                    "]" + 
                 "}");
             verifyJsonSerializationText(
                 "{" +
@@ -446,6 +464,13 @@ public void SqlQuerySpecSerializationTest()
                         "{\"name\":\"@p1\",\"value\":{\"a\":[1,2,\"abc\"]}}" +
                     "]" +
                 "}");
+            verifyJsonSerializationText(
+                "{" +
+                    "\"query\":\"SELECT 1\"," +
+                    "\"parameters\":[" +
+                        "{\"name\":\"@p1\",\"value\":{\"a\":[1,2,\"abc\"]}}" +
+                    "]," + "\"passThrough\":true" + 
+                "}");
         }
 
         [TestMethod]

diff --git a/...sts/BaselineTest/TestBaseline/PassThroughQueryBaselineTests.NegativePassThroughOutput.xml b/...sts/BaselineTest/TestBaseline/PassThroughQueryBaselineTests.NegativePassThroughOutput.xml
@@ -0,0 +1,41 @@
+<Results>
+  <Result>
+    <Input>
+      <Description>Null Partition Key Value</Description>
+      <Query>SELECT * FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>false</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description>None Partition Key Value</Description>
+      <Query>SELECT * FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>false</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description>C# Null Partition Key Value</Description>
+      <Query>SELECT * FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>false</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+</Results>
diff --git a/...sts/BaselineTest/TestBaseline/PassThroughQueryBaselineTests.PositivePassThroughOutput.xml b/...sts/BaselineTest/TestBaseline/PassThroughQueryBaselineTests.PositivePassThroughOutput.xml
@@ -0,0 +1,41 @@
+<Results>
+  <Result>
+    <Input>
+      <Description>Partition Key + Value and Distinct</Description>
+      <Query>SELECT DISTINCT c.key FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>true</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description>Partition Key + Value and Min Aggregate</Description>
+      <Query>SELECT VALUE MIN(c.key) FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>true</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description>Partition Key + Value Fields</Description>
+      <Query>SELECT c.key FROM c</Query>
+      <PartitionKeys>
+        <Key>/pk</Key>
+      </PartitionKeys>
+      <PartitionKeyType>Hash</PartitionKeyType>
+    </Input>
+    <Output>
+      <ExecuteAsPassThrough>true</ExecuteAsPassThrough>
+    </Output>
+  </Result>
+</Results>
diff --git a/...osoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj b/...osoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj
@@ -228,6 +228,12 @@
     <None Update="BaselineTest\TestBaseline\QueryPlanBaselineTests.Top.xml">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </None>
+	<None Update="BaselineTest\TestBaseline\PassThroughQueryBaselineTests.PositivePassThroughOutput.xml">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+	</None>
+	<None Update="BaselineTest\TestBaseline\PassThroughQueryBaselineTests.NegativePassThroughOutput.xml">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+	</None>
     <None Update="Contracts\DotNetPreviewSDKAPI.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>