Adding basic request diagnostics for V3

Microsoft.Azure.Cosmos/src/Handler/ResponseMessage.cs

@@ -108,6 +108,11 @@ public virtual Stream Content
         /// </summary>
         public virtual RequestMessage RequestMessage { get; internal set; }
 
+        /// <summary>
+        /// Gets the cosmos diagnostic information for the current request to Azure Cosmos DB service
+        /// </summary>
+        public CosmosDiagnostics cosmosDiagnostics { get; set; }
+
         /// <summary>
         /// Gets the internal error object.
         /// </summary>

Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs

@@ -75,6 +75,7 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties(
                 RetryAfterLiteral = this.RetryAfterLiteral,
                 SubStatusCodeLiteral = this.SubStatusCodeLiteral,
                 ContentType = this.ContentType,
+                QueryMetricsText = QueryMetricsText
             };
         }
 
@@ -106,6 +107,7 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
                 RetryAfterLiteral = sourceHeaders.RetryAfterLiteral,
                 SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral,
                 ContentType = sourceHeaders.ContentType,
+                QueryMetricsText = sourceHeaders.QueryMetricsText
             };
         }
     }

Microsoft.Azure.Cosmos/src/Headers/Headers.cs

@@ -167,6 +167,9 @@ internal string RetryAfterLiteral
         [CosmosKnownHeaderAttribute(HeaderName = HttpConstants.HttpHeaders.PageSize)]
         internal string PageSize { get; set; }
 
+        [CosmosKnownHeaderAttribute(HeaderName = HttpConstants.HttpHeaders.QueryMetrics)]
+        internal string QueryMetricsText { get; set; }
+
         /// <summary>
         /// Creates a new instance of <see cref="Headers"/>.
         /// </summary>

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

@@ -236,14 +236,6 @@ public CosmosQueryResponseMessageHeaders GetResponseHeaders()
 
             this.SetQueryMetrics();
 
-            IReadOnlyDictionary<string, QueryMetrics> groupedQueryMetrics = this.GetQueryMetrics();
-            if (groupedQueryMetrics != null && groupedQueryMetrics.Count != 0)
-            {
-                responseHeaders[HttpConstants.HttpHeaders.QueryMetrics] = QueryMetrics
-                .CreateFromIEnumerable(groupedQueryMetrics.Values)
-                .ToDelimitedString();
-            }
-
             responseHeaders.RequestCharge = this.requestChargeTracker.GetAndResetCharge();
 
             return responseHeaders;
@@ -392,6 +384,7 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
                 result: results,
                 count: results.Count,
                 responseHeaders: this.GetResponseHeaders(),
+                queryMetrics: this.GetQueryMetrics(),
                 responseLengthBytes: this.GetAndResetResponseLengthBytes());
         }
 

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

@@ -199,6 +199,11 @@ private async Task<ResponseMessage> ExecuteNextHelperAsync(CancellationToken can
                 }
             }
 
+            if (response?.queryMetrics != null && response?.queryMetrics.Count > 0)
+            {
+                response.cosmosDiagnostics = new QueryOperationStatistics(response.queryMetrics);
+            }
+
             return response;
         }
 
@@ -307,7 +312,7 @@ private async Task<CosmosQueryExecutionContext> CreateItemQueryExecutionContextA
                 rewrittenComosQueryContext = this.cosmosQueryContext;
             }
 
-            return await CreateSpecializedDocumentQueryExecutionContextAsync(
+            return await this.CreateSpecializedDocumentQueryExecutionContextAsync(
                 rewrittenComosQueryContext,
                 partitionedQueryExecutionInfo,
                 targetRanges,

Microsoft.Azure.Cosmos/src/Query/ExecutionComponent/AggregateDocumentQueryExecutionComponent.cs

@@ -106,7 +106,6 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
             PartitionedQueryMetrics partitionedQueryMetrics = new PartitionedQueryMetrics();
             ResourceType resourceType = ResourceType.Document;
             string containerRid = null;
-
             while (!this.IsDone)
             {
                 QueryResponse result = await base.DrainAsync(int.MaxValue, token);
@@ -150,13 +149,14 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
                 count: finalResult.Count,
                 responseLengthBytes: responseLengthBytes,
                 responseHeaders: new CosmosQueryResponseMessageHeaders(
-                    continauationToken: null,
-                    disallowContinuationTokenMessage: null,
-                    resourceType: resourceType,
+                    continauationToken: null, 
+                    disallowContinuationTokenMessage: null, 
+                    resourceType: resourceType, 
                     containerRid: containerRid)
                 {
                     RequestCharge = requestCharge
-                });
+                },
+                queryMetrics: this.GetQueryMetrics());
         }
 
         /// <summary>

Microsoft.Azure.Cosmos/src/Query/ExecutionComponent/DistinctDocumentQueryExecutionComponent.cs

@@ -133,7 +133,8 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
                 distinctResults,
                 distinctResults.Count,
                 cosmosQueryResponse.ResponseLengthBytes,
-                cosmosQueryResponse.QueryHeaders.CloneKnownProperties(updatedContinuationToken, disallowContinuationTokenMessage));
+                cosmosQueryResponse.QueryHeaders.CloneKnownProperties(updatedContinuationToken, disallowContinuationTokenMessage),
+                cosmosQueryResponse.queryMetrics);
         }
 
         /// <summary>

Microsoft.Azure.Cosmos/src/Query/ExecutionComponent/SkipDocumentQueryExecutionComponent.cs

@@ -88,6 +88,7 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
                     result: documentsAfterSkip,
                     count: documentsAfterSkip.Count(),
                     responseHeaders: sourcePage.QueryHeaders.CloneKnownProperties(updatedContinuationToken, sourcePage.QueryHeaders.DisallowContinuationTokenMessage),
+                    queryMetrics: sourcePage.queryMetrics,
                     responseLengthBytes: sourcePage.ResponseLengthBytes);
         }
 

Microsoft.Azure.Cosmos/src/Query/ExecutionComponent/TakeDocumentQueryExecutionComponent.cs

@@ -139,7 +139,8 @@ public override async Task<QueryResponse> DrainAsync(int maxElements, Cancellati
                 takedDocuments,
                 takedDocuments.Count,
                 results.ResponseLengthBytes,
-                results.QueryHeaders.CloneKnownProperties(updatedContinuationToken, results.QueryHeaders.DisallowContinuationTokenMessage));
+                results.QueryHeaders.CloneKnownProperties(updatedContinuationToken, results.QueryHeaders.DisallowContinuationTokenMessage),
+                results.queryMetrics);
         }
 
         private enum TakeEnum

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

@@ -426,7 +426,8 @@ public override async Task<QueryResponse> ExecuteNextAsync(CancellationToken tok
                     dynamics,
                     queryResponse.Count,
                     queryResponse.ResponseLengthBytes,
-                    queryResponse.QueryHeaders.CloneKnownProperties());
+                    queryResponse.QueryHeaders.CloneKnownProperties(),
+                    queryMetrics: queryResponse.queryMetrics);
             }
             catch (Exception)
             {

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

@@ -157,7 +157,7 @@ public long OutputDocumentCount
         /// <summary>
         /// Gets the size of documents outputted in bytes during query in the Azure DocumentDB database service.
         /// </summary>
-        internal long OutputDocumentSize
+        public long OutputDocumentSize
         {
             get
             {
@@ -168,7 +168,7 @@ internal long OutputDocumentSize
         /// <summary>
         /// Gets the total query time in the Azure DocumentDB database service.
         /// </summary>
-        internal TimeSpan TotalQueryExecutionTime
+        public TimeSpan TotalQueryExecutionTime
         {
             get
             {
@@ -212,7 +212,7 @@ public long Retries
         /// <summary>
         /// Gets the query index lookup time in the Azure DocumentDB database service.
         /// </summary>
-        internal TimeSpan IndexLookupTime
+        public TimeSpan IndexLookupTime
         {
             get
             {
@@ -223,7 +223,7 @@ internal TimeSpan IndexLookupTime
         /// <summary>
         /// Gets the document loading time during query in the Azure DocumentDB database service.
         /// </summary>
-        internal TimeSpan DocumentLoadTime
+        public TimeSpan DocumentLoadTime
         {
             get
             {
@@ -234,7 +234,7 @@ internal TimeSpan DocumentLoadTime
         /// <summary>
         /// Gets the query runtime execution times during query in the Azure DocumentDB database service.
         /// </summary>
-        internal RuntimeExecutionTimes RuntimeExecutionTimes
+        public RuntimeExecutionTimes RuntimeExecutionTimes
         {
             get
             {
@@ -245,7 +245,7 @@ internal RuntimeExecutionTimes RuntimeExecutionTimes
         /// <summary>
         /// Gets the output writing/serializing time during query in the Azure DocumentDB database service.
         /// </summary>
-        internal TimeSpan DocumentWriteTime
+        public TimeSpan DocumentWriteTime
         {
             get
             {
@@ -281,7 +281,7 @@ public double IndexHitRatio
         /// <summary>
         /// Gets the Index Hit Document Count.
         /// </summary>
-        internal long IndexHitDocumentCount
+        public long IndexHitDocumentCount
         {
             get
             {
@@ -292,7 +292,7 @@ internal long IndexHitDocumentCount
         /// <summary>
         /// Gets the VMExecution Time.
         /// </summary>
-        internal TimeSpan VMExecutionTime
+        public TimeSpan VMExecutionTime
         {
             get
             {
@@ -560,7 +560,7 @@ public TimeSpan WriteOutputTime
         /// <summary>
         /// Gets the query runtime execution times during query in the Azure DocumentDB database service.
         /// </summary>
-        internal RuntimeExecutionTimes RuntimeExecutionTimes
+        public RuntimeExecutionTimes RuntimeExecutionTimes
         {
             get
             {

Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs

@@ -198,6 +198,8 @@ internal override void PopulateRequestOptions(RequestMessage request)
                 request.Headers.Add(HttpConstants.HttpHeaders.ContentSerializationFormat, this.CosmosSerializationOptions.ContentSerializationFormat);
             }
 
+            request.Headers.Add(HttpConstants.HttpHeaders.PopulateQueryMetrics, bool.TrueString);
+
             base.PopulateRequestOptions(request);
         }
 
@@ -220,7 +222,7 @@ internal QueryRequestOptions Clone()
                 EnableCrossPartitionSkipTake = this.EnableCrossPartitionSkipTake,
                 EnableGroupBy = this.EnableGroupBy,
                 Properties = this.Properties,
-                IsEffectivePartitionKeyRouting = this.IsEffectivePartitionKeyRouting
+                IsEffectivePartitionKeyRouting = this.IsEffectivePartitionKeyRouting,
             };
 
             return queryRequestOptions;

Microsoft.Azure.Cosmos/src/Resource/Container/ItemResponse.cs

@@ -27,11 +27,13 @@ protected ItemResponse()
         internal ItemResponse(
             HttpStatusCode httpStatusCode,
             Headers headers,
-            T item)
+            T item,
+            CosmosDiagnostics cosmosDiagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = item;
+            this.cosmosDiagnostics = cosmosDiagnostics;
         }
 
         /// <inheritdoc/>

Microsoft.Azure.Cosmos/src/Resource/CosmosResponseFactory.cs

@@ -38,7 +38,7 @@ internal FeedResponse<T> CreateQueryFeedResponse<T>(
             if (queryResponse != null)
             {
                 return QueryResponse<T>.CreateResponse<T>(
-                    responseMessage: queryResponse,
+                    cosmosQueryResponse: queryResponse,
                     jsonSerializer: this.cosmosSerializer);
             }
 
@@ -56,7 +56,8 @@ internal Task<ItemResponse<T>> CreateItemResponseAsync<T>(
                 return new ItemResponse<T>(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    item);
+                    item,
+                    cosmosResponseMessage.cosmosDiagnostics);
             });
         }
 

Microsoft.Azure.Cosmos/src/Resource/QueryResponses/QueryResponse.cs

@@ -28,6 +28,7 @@ private QueryResponse(
             int count,
             long responseLengthBytes,
             CosmosQueryResponseMessageHeaders responseHeaders,
+            IReadOnlyDictionary<string, QueryMetrics> queryMetrics,
             HttpStatusCode statusCode,
             RequestMessage requestMessage,
             string errorMessage,
@@ -42,6 +43,7 @@ private QueryResponse(
             this.CosmosElements = result;
             this.Count = count;
             this.ResponseLengthBytes = responseLengthBytes;
+            this.queryMetrics = queryMetrics;
         }
 
         public int Count { get; }
@@ -72,6 +74,8 @@ private QueryResponse(
         /// </remarks>
         internal ClientSideRequestStatistics RequestStatistics { get; }
 
+        internal IReadOnlyDictionary<string, QueryMetrics> queryMetrics { get; set; }
+
         internal virtual CosmosSerializationFormatOptions CosmosSerializationOptions { get; set; }
 
         internal bool GetHasMoreResults()
@@ -83,7 +87,8 @@ internal static QueryResponse CreateSuccess(
             IEnumerable<CosmosElement> result,
             int count,
             long responseLengthBytes,
-            CosmosQueryResponseMessageHeaders responseHeaders)
+            CosmosQueryResponseMessageHeaders responseHeaders,
+            IReadOnlyDictionary<string, QueryMetrics> queryMetrics = null)
         {
             if (count < 0)
             {
@@ -100,6 +105,7 @@ internal static QueryResponse CreateSuccess(
                count: count,
                responseLengthBytes: responseLengthBytes,
                responseHeaders: responseHeaders,
+               queryMetrics: queryMetrics,
                statusCode: HttpStatusCode.OK,
                errorMessage: null,
                error: null,
@@ -120,6 +126,7 @@ internal static QueryResponse CreateFailure(
                 count: 0,
                 responseLengthBytes: 0,
                 responseHeaders: responseHeaders,
+                queryMetrics: null,
                 statusCode: statusCode,
                 errorMessage: errorMessage,
                 error: error,
@@ -144,11 +151,13 @@ private QueryResponse(
             HttpStatusCode httpStatusCode,
             IEnumerable<CosmosElement> cosmosElements,
             CosmosQueryResponseMessageHeaders responseMessageHeaders,
+            CosmosDiagnostics cosmosDiagnostics,
             CosmosSerializer jsonSerializer,
             CosmosSerializationFormatOptions serializationOptions)
         {
             this.cosmosElements = cosmosElements;
             this.QueryHeaders = responseMessageHeaders;
+            this.cosmosDiagnostics = cosmosDiagnostics;
             this.jsonSerializer = jsonSerializer;
             this.serializationOptions = serializationOptions;
             this.StatusCode = httpStatusCode;
@@ -189,21 +198,6 @@ public override IEnumerable<T> Resource
             }
         }
 
-        internal static QueryResponse<TInput> CreateResponse<TInput>(
-            ResponseMessage responseMessage,
-            CosmosSerializer jsonSerializer)
-        {
-            QueryResponse queryResponse = responseMessage as QueryResponse;
-            if (queryResponse == null)
-            {
-                throw new ArgumentException($"{nameof(responseMessage)} must be of type query response.");
-            }
-
-            return QueryResponse<TInput>.CreateResponse<TInput>(
-                queryResponse,
-                jsonSerializer);
-        }
-
         internal static QueryResponse<TInput> CreateResponse<TInput>(
             QueryResponse cosmosQueryResponse,
             CosmosSerializer jsonSerializer)
@@ -216,10 +210,10 @@ internal static QueryResponse<TInput> CreateResponse<TInput>(
                     httpStatusCode: cosmosQueryResponse.StatusCode,
                     cosmosElements: cosmosQueryResponse.CosmosElements,
                     responseMessageHeaders: cosmosQueryResponse.QueryHeaders,
+                    cosmosDiagnostics: cosmosQueryResponse.cosmosDiagnostics,
                     jsonSerializer: jsonSerializer,
                     serializationOptions: cosmosQueryResponse.CosmosSerializationOptions);
             }
-
             return queryResponse;
         }
     }

Microsoft.Azure.Cosmos/src/Resource/Response.cs

@@ -63,6 +63,11 @@ public static implicit operator T(Response<T> response)
         /// </remarks>
         public abstract string ETag { get; }
 
+        /// <summary>
+        /// Gets the cosmos diagnostics information for the current request to Azure Cosmos DB service
+        /// </summary>
+        public CosmosDiagnostics cosmosDiagnostics { get; set; }
+
         /// <summary>
         /// Gets the maximum size limit for this entity from the Azure Cosmos DB service.
         /// </summary>