Fixed bug where query client would use old diagnostics which would ca…

…use an exception.
Azure · Mar 19, 2020 · f0c2d9e · f0c2d9e
1 parent 509a986
commit f0c2d9e
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 11 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs
@@ -49,7 +49,7 @@ internal abstract Task<QueryResponseCore> ExecuteItemQueryAsync(
             Documents.ResourceType resourceType,
             Documents.OperationType operationType,
             QueryRequestOptions requestOptions,
-            CosmosDiagnosticsContext diagnosticsContext,
+            Action<QueryPageDiagnostics> queryPageDiagnostics,
             SqlQuerySpec sqlQuerySpec,
             string continuationToken,
             Documents.PartitionKeyRangeIdentity partitionKeyRange,

diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs
@@ -117,7 +117,7 @@ internal override async Task<QueryResponseCore> ExecuteItemQueryAsync(
             ResourceType resourceType,
             OperationType operationType,
             QueryRequestOptions requestOptions,
-            CosmosDiagnosticsContext diagnosticsContext,
+            Action<QueryPageDiagnostics> queryPageDiagnostics,
             SqlQuerySpec sqlQuerySpec,
             string continuationToken,
             PartitionKeyRangeIdentity partitionKeyRange,
@@ -155,7 +155,7 @@ internal override async Task<QueryResponseCore> ExecuteItemQueryAsync(
                 resourceType,
                 message,
                 partitionKeyRange,
-                diagnosticsContext);
+                queryPageDiagnostics);
         }
 
         internal override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync(
@@ -265,7 +265,7 @@ private QueryResponseCore GetCosmosElementResponse(
             ResourceType resourceType,
             ResponseMessage cosmosResponseMessage,
             PartitionKeyRangeIdentity partitionKeyRangeIdentity,
-            CosmosDiagnosticsContext diagnosticsContext)
+            Action<QueryPageDiagnostics> queryPageDiagnostics)
         {
             using (cosmosResponseMessage)
             {
@@ -274,8 +274,8 @@ private QueryResponseCore GetCosmosElementResponse(
                     queryMetricText: cosmosResponseMessage.Headers.QueryMetricsText,
                     indexUtilizationText: cosmosResponseMessage.Headers[HttpConstants.HttpHeaders.IndexUtilization],
                     diagnosticsContext: cosmosResponseMessage.DiagnosticsContext);
+                queryPageDiagnostics(queryPage);
 
-                diagnosticsContext.AddDiagnosticsInternal(queryPage);
                 if (!cosmosResponseMessage.IsSuccessStatusCode)
                 {
                     return QueryResponseCore.CreateFailure(

diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs
@@ -82,7 +82,7 @@ internal override Task<QueryResponseCore> ExecuteQueryAsync(
                            partitionKeyRange: partitionKeyRange,
                            isContinuationExpected: isContinuationExpected,
                            pageSize: pageSize,
-                           diagnosticsContext: this.diagnosticsContext,
+                           queryPageDiagnostics: this.AddQueryPageDiagnostic,
                            cancellationToken: cancellationToken);
         }
 
@@ -105,5 +105,10 @@ internal override Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsy
                 this.diagnosticsContext,
                 cancellationToken);
         }
+
+        private void AddQueryPageDiagnostic(QueryPageDiagnostics queryPageDiagnostics)
+        {
+            this.diagnosticsContext.AddDiagnosticsInternal(queryPageDiagnostics);
+        }
     }
 }
diff --git a/...ft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs b/...ft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs
@@ -66,7 +66,7 @@ internal override Task<QueryResponseCore> ExecuteItemQueryAsync(
             ResourceType resourceType,
             OperationType operationType,
             QueryRequestOptions requestOptions,
-            CosmosDiagnosticsContext diagnosticsContext,
+            Action<QueryPageDiagnostics> queryPageDiagnostics,
             SqlQuerySpec sqlQuerySpec,
             string continuationToken,
             PartitionKeyRangeIdentity partitionKeyRange,
@@ -82,7 +82,7 @@ internal override Task<QueryResponseCore> ExecuteItemQueryAsync(
                 resourceType: resourceType,
                 operationType: operationType,
                 requestOptions: requestOptions,
-                diagnosticsContext: diagnosticsContext,
+                queryPageDiagnostics: queryPageDiagnostics,
                 sqlQuerySpec: sqlQuerySpec,
                 continuationToken: continuationToken,
                 partitionKeyRange: partitionKeyRange,

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockItemProducerFactory.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockItemProducerFactory.cs
@@ -234,7 +234,7 @@ public static List<ToDoItem> MockSinglePartitionKeyRangeContext(
                         ResourceType.Document,
                         OperationType.Query,
                         It.IsAny<QueryRequestOptions>(),
-                        It.IsAny<CosmosDiagnosticsContext>(),
+                        It.IsAny<Action<QueryPageDiagnostics>>(),
                         It.Is<SqlQuerySpec>(specInput => IsSqlQuerySpecEqual(sqlQuerySpec, specInput)),
                         previousContinuationToken,
                         It.Is<PartitionKeyRangeIdentity>(rangeId => string.Equals(rangeId.PartitionKeyRangeId, partitionKeyRange.Id) && string.Equals(rangeId.CollectionRid, collectionRid)),

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs
@@ -151,7 +151,7 @@ private static IList<ToDoItem> GenerateAndMockResponseHelper(
                           ResourceType.Document,
                           OperationType.Query,
                           It.IsAny<QueryRequestOptions>(),
-                          It.IsAny<CosmosDiagnosticsContext>(),
+                          It.IsAny<Action<QueryPageDiagnostics>>(),
                           It.Is<SqlQuerySpec>(specInput => MockItemProducerFactory.IsSqlQuerySpecEqual(sqlQuerySpec, specInput)),
                           previousContinuationToken,
                           It.Is<PartitionKeyRangeIdentity>(rangeId => string.Equals(rangeId.PartitionKeyRangeId, partitionKeyRange.Id) && string.Equals(rangeId.CollectionRid, containerRid)),
@@ -179,7 +179,7 @@ private static IList<ToDoItem> GenerateAndMockResponseHelper(
                          ResourceType.Document,
                          OperationType.Query,
                          It.IsAny<QueryRequestOptions>(),
-                         It.IsAny<CosmosDiagnosticsContext>(),
+                         It.IsAny<Action<QueryPageDiagnostics>>(),
                          It.Is<SqlQuerySpec>(specInput => MockItemProducerFactory.IsSqlQuerySpecEqual(sqlQuerySpec, specInput)),
                          previousContinuationToken,
                          It.Is<PartitionKeyRangeIdentity>(rangeId => string.Equals(rangeId.PartitionKeyRangeId, partitionKeyRange.Id) && string.Equals(rangeId.CollectionRid, containerRid)),