Adding additional diagnostics context

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

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

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

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Query
 {
     using System.Collections.Generic;
+    using System.Text;
     using Newtonsoft.Json;
 
     internal class CosmosDiagnosticsAggregate : CosmosDiagnostics
@@ -13,7 +14,18 @@ internal class CosmosDiagnosticsAggregate : CosmosDiagnostics
 
         public override string ToString()
         {
-            return JsonConvert.SerializeObject(this);
+            if (this.Diagnostics.Count == 0)
+            {
+                return string.Empty;
+            }
+
+            StringBuilder stringBuilder = new StringBuilder();
+            foreach (CosmosDiagnostics diagnostics in this.Diagnostics)
+            {
+                stringBuilder.Append(diagnostics.ToString());
+            }
+
+            return stringBuilder.ToString();
         }
     }
 }

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

@@ -28,12 +28,14 @@ internal ContainerResponse(
             HttpStatusCode httpStatusCode,
             Headers headers,
             ContainerProperties containerProperties,
-            Container container)
+            Container container,
+            CosmosDiagnostics diagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = containerProperties;
             this.Container = container;
+            this.Diagnostics = diagnostics;
         }
 
         /// <summary>
@@ -51,6 +53,9 @@ internal ContainerResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;
 

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

@@ -45,6 +45,9 @@ internal ItemResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;
 

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

@@ -44,6 +44,7 @@ internal CosmosException(
                 this.RequestCharge = this.Headers.RequestCharge;
                 this.RetryAfter = this.Headers.RetryAfter;
                 this.SubStatusCode = (int)this.Headers.SubStatusCode;
+                this.Diagnostics = cosmosResponseMessage.Diagnostics;
                 if (this.Headers.ContentLengthAsLong > 0)
                 {
                     using (StreamReader responseReader = new StreamReader(cosmosResponseMessage.Content))
@@ -122,6 +123,11 @@ public CosmosException(
         /// </summary>
         public virtual Headers Headers { get; }
 
+        /// <summary>
+        /// Gets the diagnostics for the request
+        /// </summary>
+        public virtual CosmosDiagnostics Diagnostics { get; }
+
         /// <summary>
         /// Gets the internal error object
         /// </summary>
@@ -150,7 +156,8 @@ public virtual bool TryGetHeader(string headerName, out string value)
         /// <returns>A string representation of the exception.</returns>
         public override string ToString()
         {
-            return $"{nameof(CosmosException)};StatusCode={this.StatusCode};SubStatusCode={this.SubStatusCode};ActivityId={this.ActivityId ?? string.Empty};RequestCharge={this.RequestCharge};Message={this.Message};";
+            string diagnostics = this.Diagnostics != null ? this.Diagnostics.ToString() : string.Empty;
+            return $"{nameof(CosmosException)};StatusCode={this.StatusCode};SubStatusCode={this.SubStatusCode};ActivityId={this.ActivityId ?? string.Empty};RequestCharge={this.RequestCharge};Message={this.Message};Diagnostics{diagnostics}";
         }
 
         internal ResponseMessage ToCosmosResponseMessage(RequestMessage request)

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

@@ -93,7 +93,8 @@ internal Task<ContainerResponse> CreateContainerResponseAsync(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
                     containerProperties,
-                    container);
+                    container,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -108,7 +109,8 @@ internal Task<UserResponse> CreateUserResponseAsync(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
                     userProperties,
-                    user);
+                    user,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -123,7 +125,8 @@ internal Task<PermissionResponse> CreatePermissionResponseAsync(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
                     permissionProperties,
-                    permission);
+                    permission,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -141,7 +144,8 @@ internal Task<DatabaseResponse> CreateDatabaseResponseAsync(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
                     databaseProperties,
-                    database);
+                    database,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -154,7 +158,8 @@ internal Task<ThroughputResponse> CreateThroughputResponseAsync(
                 return new ThroughputResponse(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    throughputProperties);
+                    throughputProperties,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -166,7 +171,8 @@ internal Task<StoredProcedureExecuteResponse<T>> CreateStoredProcedureExecuteRes
                 return new StoredProcedureExecuteResponse<T>(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    item);
+                    item,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -178,7 +184,8 @@ internal Task<StoredProcedureResponse> CreateStoredProcedureResponseAsync(Task<R
                 return new StoredProcedureResponse(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    cosmosStoredProcedure);
+                    cosmosStoredProcedure,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -190,7 +197,8 @@ internal Task<TriggerResponse> CreateTriggerResponseAsync(Task<ResponseMessage>
                 return new TriggerResponse(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    triggerProperties);
+                    triggerProperties,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 
@@ -202,7 +210,8 @@ internal Task<UserDefinedFunctionResponse> CreateUserDefinedFunctionResponseAsyn
                 return new UserDefinedFunctionResponse(
                     cosmosResponseMessage.StatusCode,
                     cosmosResponseMessage.Headers,
-                    settings);
+                    settings,
+                    cosmosResponseMessage.Diagnostics);
             });
         }
 

Microsoft.Azure.Cosmos/src/Resource/Database/DatabaseResponse.cs

@@ -28,12 +28,14 @@ internal DatabaseResponse(
             HttpStatusCode httpStatusCode,
             Headers headers,
             DatabaseProperties databaseProperties,
-            Database database)
+            Database database,
+            CosmosDiagnostics diagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = databaseProperties;
             this.Database = database;
+            this.Diagnostics = diagnostics;
         }
 
         /// <summary>
@@ -51,6 +53,9 @@ internal DatabaseResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;
 

Microsoft.Azure.Cosmos/src/Resource/Offer/CosmosOffers.cs

@@ -52,7 +52,8 @@ internal async Task<ThroughputResponse> ReadThroughputIfExistsAsync(
                 return new ThroughputResponse(
                     HttpStatusCode.NotFound,
                     headers: null,
-                    throughputProperties: null);
+                    throughputProperties: null,
+                    diagnostics: null);
             }
 
             return await this.GetThroughputResponseAsync(
@@ -107,15 +108,17 @@ internal async Task<ThroughputResponse> ReplaceThroughputIfExistsAsync(
                 return new ThroughputResponse(
                     responseMessage.StatusCode,
                     headers: responseMessage.Headers,
-                    throughputProperties: null);
+                    throughputProperties: null,
+                    diagnostics: responseMessage.Diagnostics);
             }
             catch (AggregateException ex)
             {
                 ResponseMessage responseMessage = TransportHandler.AggregateExceptionConverter(ex, null);
                 return new ThroughputResponse(
                     responseMessage.StatusCode,
                     headers: responseMessage.Headers,
-                    throughputProperties: null);
+                    throughputProperties: null,
+                    diagnostics: responseMessage.Diagnostics);
             }
         }
 

Microsoft.Azure.Cosmos/src/Resource/Permission/PermissionResponse.cs

@@ -28,12 +28,14 @@ internal PermissionResponse(
             HttpStatusCode httpStatusCode,
             Headers headers,
             PermissionProperties permissionProperties,
-            Permission permission)
+            Permission permission,
+            CosmosDiagnostics diagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = permissionProperties;
             this.Permission = permission;
+            this.Diagnostics = diagnostics;
         }
 
         /// <summary>
@@ -51,6 +53,9 @@ internal PermissionResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;
 

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

@@ -172,6 +172,8 @@ private QueryResponse(
 
         public override HttpStatusCode StatusCode { get; }
 
+        public override CosmosDiagnostics Diagnostics { get; }
+
         public override int Count => this.cosmosElements?.Count() ?? 0;
 
         internal CosmosQueryResponseMessageHeaders QueryHeaders { get; }

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

@@ -11,12 +11,14 @@ internal class ReadFeedResponse<T> : FeedResponse<T>
         protected ReadFeedResponse(
             HttpStatusCode httpStatusCode,
             ICollection<T> resource,
-            Headers responseMessageHeaders)
+            Headers responseMessageHeaders,
+            CosmosDiagnostics diagnostics)
         {
             this.Count = resource.Count;
             this.Headers = responseMessageHeaders;
             this.Resource = resource;
             this.StatusCode = httpStatusCode;
+            this.Diagnostics = diagnostics;
         }
 
         public override int Count { get; }
@@ -29,6 +31,8 @@ protected ReadFeedResponse(
 
         public override HttpStatusCode StatusCode { get; }
 
+        public override CosmosDiagnostics Diagnostics { get; }
+
         public override IEnumerator<T> GetEnumerator()
         {
             return this.Resource.GetEnumerator();
@@ -50,7 +54,8 @@ internal static ReadFeedResponse<TInput> CreateResponse<TInput>(
                 ReadFeedResponse<TInput> readFeedResponse = new ReadFeedResponse<TInput>(
                     httpStatusCode: responseMessage.StatusCode,
                     resource: resources,
-                    responseMessageHeaders: responseMessage.Headers);
+                    responseMessageHeaders: responseMessage.Headers,
+                    diagnostics: responseMessage.Diagnostics);
 
                 return readFeedResponse;
             }

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

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

Microsoft.Azure.Cosmos/src/Resource/Scripts/StoredProcedureExecuteResponse.cs

@@ -26,13 +26,15 @@ protected StoredProcedureExecuteResponse()
         /// This will prevent memory leaks when handling the HttpResponseMessage
         /// </summary>
         internal StoredProcedureExecuteResponse(
-           HttpStatusCode httpStatusCode,
-           Headers headers,
-           T response)
+            HttpStatusCode httpStatusCode,
+            Headers headers,
+            T response,
+            CosmosDiagnostics diagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = response;
+            this.Diagnostics = diagnostics;
         }
 
         /// <inheritdoc/>
@@ -44,6 +46,9 @@ internal StoredProcedureExecuteResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;
 

Microsoft.Azure.Cosmos/src/Resource/Scripts/StoredProcedureResponse.cs

@@ -25,13 +25,15 @@ protected StoredProcedureResponse()
         /// This will prevent memory leaks when handling the HttpResponseMessage
         /// </summary>
         internal StoredProcedureResponse(
-           HttpStatusCode httpStatusCode,
-           Headers headers,
-           StoredProcedureProperties storedProcedureProperties)
+            HttpStatusCode httpStatusCode,
+            Headers headers,
+            StoredProcedureProperties storedProcedureProperties,
+            CosmosDiagnostics diagnostics)
         {
             this.StatusCode = httpStatusCode;
             this.Headers = headers;
             this.Resource = storedProcedureProperties;
+            this.Diagnostics = diagnostics;
         }
 
         /// <inheritdoc/>
@@ -43,6 +45,9 @@ internal StoredProcedureResponse(
         /// <inheritdoc/>
         public override HttpStatusCode StatusCode { get; }
 
+        /// <inheritdoc/>
+        public override CosmosDiagnostics Diagnostics { get; }
+
         /// <inheritdoc/>
         public override double RequestCharge => this.Headers?.RequestCharge ?? 0;