Add ThrowOnPartialError to Azure.Monitor.Query (Azure#23988)

Co-authored-by: Christopher Scott <chriscott@hotmail.com>

Author: pakrym

sdk/monitor/Azure.Monitor.Query/README.md

@@ -212,7 +212,7 @@ string topQueryId = batch.AddQuery(
     "AzureActivity | summarize Count = count() by ResourceGroup | top 10 by Count",
     new DateTimeRange(TimeSpan.FromDays(1)));
 
-Response<LogsBatchQueryResults> response = await client.QueryBatchAsync(batch);
+Response<LogsBatchQueryResultCollection> response = await client.QueryBatchAsync(batch);
 
 var count = response.Value.GetResult<int>(countQueryId).Single();
 var topEntries = response.Value.GetResult<MyLogEntryModel>(topQueryId);

sdk/monitor/Azure.Monitor.Query/api/Azure.Monitor.Query.netstandard2.0.cs

@@ -17,8 +17,8 @@ public LogsQueryClient(System.Uri endpoint, Azure.Core.TokenCredential credentia
         public virtual Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult> Query(string workspaceId, string query, Azure.Core.DateTimeRange timeRange, Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult>> QueryAsync(string workspaceId, string query, Azure.Core.DateTimeRange timeRange, Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual System.Threading.Tasks.Task<Azure.Response<System.Collections.Generic.IReadOnlyList<T>>> QueryAsync<T>(string workspaceId, string query, Azure.Core.DateTimeRange timeRange, Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
-        public virtual Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResults> QueryBatch(Azure.Monitor.Query.LogsBatchQuery batch, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
-        public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResults>> QueryBatchAsync(Azure.Monitor.Query.LogsBatchQuery batch, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
+        public virtual Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResultCollection> QueryBatch(Azure.Monitor.Query.LogsBatchQuery batch, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
+        public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResultCollection>> QueryBatchAsync(Azure.Monitor.Query.LogsBatchQuery batch, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual Azure.Response<System.Collections.Generic.IReadOnlyList<T>> Query<T>(string workspaceId, string query, Azure.Core.DateTimeRange timeRange, Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
     public partial class LogsQueryClientOptions : Azure.Core.ClientOptions
@@ -37,6 +37,7 @@ public LogsQueryOptions() { }
         public bool IncludeStatistics { get { throw null; } set { } }
         public bool IncludeVisualization { get { throw null; } set { } }
         public System.TimeSpan? ServerTimeout { get { throw null; } set { } }
+        public bool ThrowOnPartialErrors { get { throw null; } set { } }
     }
     public partial class MetricsQueryClient
     {
@@ -76,16 +77,21 @@ namespace Azure.Monitor.Query.Models
     public partial class LogsBatchQueryResult : Azure.Monitor.Query.Models.LogsQueryResult
     {
         internal LogsBatchQueryResult() { }
-        public bool HasFailed { get { throw null; } }
         public string Id { get { throw null; } }
+        public Azure.Monitor.Query.Models.LogsBatchQueryResultStatus Status { get { throw null; } }
     }
-    public partial class LogsBatchQueryResults
+    public partial class LogsBatchQueryResultCollection : System.Collections.ObjectModel.ReadOnlyCollection<Azure.Monitor.Query.Models.LogsBatchQueryResult>
     {
-        internal LogsBatchQueryResults() { }
-        public System.Collections.Generic.IReadOnlyList<Azure.Monitor.Query.Models.LogsBatchQueryResult> Results { get { throw null; } }
+        internal LogsBatchQueryResultCollection() : base (default(System.Collections.Generic.IList<Azure.Monitor.Query.Models.LogsBatchQueryResult>)) { }
         public Azure.Monitor.Query.Models.LogsBatchQueryResult GetResult(string queryId) { throw null; }
         public System.Collections.Generic.IReadOnlyList<T> GetResult<T>(string queryId) { throw null; }
     }
+    public enum LogsBatchQueryResultStatus
+    {
+        Success = 0,
+        PartialFailure = 1,
+        Failure = 2,
+    }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
     public readonly partial struct LogsColumnType : System.IEquatable<Azure.Monitor.Query.Models.LogsColumnType>
     {

sdk/monitor/Azure.Monitor.Query/src/Generated/Models/BatchQueryResponse.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Globalization;
+using System.Linq;
 using System.Text;
 using System.Text.Json;
 using System.Threading;
@@ -219,7 +220,7 @@ public virtual async Task<Response<LogsQueryResult>> QueryAsync(string workspace
         ///     &quot;AzureActivity | summarize Count = count() by ResourceGroup | top 10 by Count&quot;,
         ///     new DateTimeRange(TimeSpan.FromDays(1)));
         ///
-        /// Response&lt;LogsBatchQueryResults&gt; response = await client.QueryBatchAsync(batch);
+        /// Response&lt;LogsBatchQueryResultCollection&gt; response = await client.QueryBatchAsync(batch);
         ///
         /// var count = response.Value.GetResult&lt;int&gt;(countQueryId).Single();
         /// var topEntries = response.Value.GetResult&lt;MyLogEntryModel&gt;(topQueryId);
@@ -233,8 +234,8 @@ public virtual async Task<Response<LogsQueryResult>> QueryAsync(string workspace
         /// </summary>
         /// <param name="batch">The batch of queries to send.</param>
         /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
-        /// <returns>The <see cref="LogsBatchQueryResults"/> containing the query identifier that has to be passed into <see cref="LogsBatchQueryResults.GetResult"/> to get the result.</returns>
-        public virtual Response<LogsBatchQueryResults> QueryBatch(LogsBatchQuery batch, CancellationToken cancellationToken = default)
+        /// <returns>The <see cref="LogsBatchQueryResultCollection"/> containing the query identifier that has to be passed into <see cref="LogsBatchQueryResultCollection.GetResult"/> to get the result.</returns>
+        public virtual Response<LogsBatchQueryResultCollection> QueryBatch(LogsBatchQuery batch, CancellationToken cancellationToken = default)
         {
             Argument.AssertNotNull(batch, nameof(batch));
 
@@ -243,7 +244,7 @@ public virtual Response<LogsBatchQueryResults> QueryBatch(LogsBatchQuery batch,
             try
             {
                 var response = _queryClient.Batch(new BatchRequest(batch.Requests), cancellationToken);
-                return response;
+                return ConvertBatchResponse(batch, response);
             }
             catch (Exception e)
             {
@@ -272,7 +273,7 @@ public virtual Response<LogsBatchQueryResults> QueryBatch(LogsBatchQuery batch,
         ///     &quot;AzureActivity | summarize Count = count() by ResourceGroup | top 10 by Count&quot;,
         ///     new DateTimeRange(TimeSpan.FromDays(1)));
         ///
-        /// Response&lt;LogsBatchQueryResults&gt; response = await client.QueryBatchAsync(batch);
+        /// Response&lt;LogsBatchQueryResultCollection&gt; response = await client.QueryBatchAsync(batch);
         ///
         /// var count = response.Value.GetResult&lt;int&gt;(countQueryId).Single();
         /// var topEntries = response.Value.GetResult&lt;MyLogEntryModel&gt;(topQueryId);
@@ -286,8 +287,8 @@ public virtual Response<LogsBatchQueryResults> QueryBatch(LogsBatchQuery batch,
         /// </summary>
         /// <param name="batch">The batch of Kusto queries to send.</param>
         /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
-        /// <returns>The <see cref="LogsBatchQueryResults"/> that allows retrieving query results.</returns>
-        public virtual async Task<Response<LogsBatchQueryResults>> QueryBatchAsync(LogsBatchQuery batch, CancellationToken cancellationToken = default)
+        /// <returns>The <see cref="LogsBatchQueryResultCollection"/> that allows retrieving query results.</returns>
+        public virtual async Task<Response<LogsBatchQueryResultCollection>> QueryBatchAsync(LogsBatchQuery batch, CancellationToken cancellationToken = default)
         {
             Argument.AssertNotNull(batch, nameof(batch));
 
@@ -296,7 +297,7 @@ public virtual async Task<Response<LogsBatchQueryResults>> QueryBatchAsync(LogsB
             try
             {
                 var response = await _queryClient.BatchAsync(new BatchRequest(batch.Requests), cancellationToken).ConfigureAwait(false);
-                return response;
+                return ConvertBatchResponse(batch, response);
             }
             catch (Exception e)
             {
@@ -305,6 +306,27 @@ public virtual async Task<Response<LogsBatchQueryResults>> QueryBatchAsync(LogsB
             }
         }
 
+        private static Response<LogsBatchQueryResultCollection> ConvertBatchResponse(LogsBatchQuery batch, Response<BatchResponse> response)
+        {
+            List<LogsBatchQueryResult> batchResponses = new List<LogsBatchQueryResult>();
+            foreach (var innerResponse in response.Value.Responses)
+            {
+                var body = innerResponse.Body;
+                body.Status = innerResponse.Status switch
+                {
+                    >= 400 => LogsBatchQueryResultStatus.Failure,
+                    _ when body.Error != null => LogsBatchQueryResultStatus.PartialFailure,
+                    _ => LogsBatchQueryResultStatus.Success
+                };
+                body.Id = innerResponse.Id;
+                batchResponses.Add(body);
+            }
+
+            return Response.FromValue(
+                new LogsBatchQueryResultCollection(batchResponses, batch),
+                response.GetRawResponse());
+        }
+
         /// <summary>
         /// Create a Kusto query from an interpolated string. The interpolated values will be quoted and escaped as necessary.
         /// </summary>
@@ -499,6 +521,11 @@ await JsonDocument.ParseAsync(message.Response.ContentStream, default, cancellat
                         JsonDocument.Parse(message.Response.ContentStream, default);
 
                     LogsQueryResult value = LogsQueryResult.DeserializeLogsQueryResult(document.RootElement);
+                    var responseError = value.Error;
+                    if (responseError != null && options?.ThrowOnPartialErrors != false)
+                    {
+                        throw value.CreateExceptionForErrorResponse(message.Response.Status);
+                    }
                     return Response.FromValue(value, message.Response);
                 }
                 default:

…s/LogsBatchQueryResults.Serialization.cs …ed/Models/BatchResponse.Serialization.cssdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsBatchQueryResults.Serialization.cs renamed to sdk/monitor/Azure.Monitor.Query/src/Generated/Models/BatchResponse.Serialization.cs sdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsBatchQueryResults.Serialization.cs renamed to sdk/monitor/Azure.Monitor.Query/src/Generated/Models/BatchResponse.Serialization.cs

@@ -36,5 +36,10 @@ public class LogsQueryOptions
         /// Gets a list of additional workspaces names to include in the query.
         /// </summary>
         public IList<string> AdditionalWorkspaces { get; } = new ChangeTrackingList<string>();
+
+        /// <summary>
+        /// Gets or sets the value indicating whether to throw when a partial error is returned with the logs response.
+        /// </summary>
+        public bool ThrowOnPartialErrors { get; set; } = true;
     }
 }

…enerated/Models/LogsBatchQueryResults.cs …ry/src/Generated/Models/BatchResponse.cssdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsBatchQueryResults.cs renamed to sdk/monitor/Azure.Monitor.Query/src/Generated/Models/BatchResponse.cs sdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsBatchQueryResults.cs renamed to sdk/monitor/Azure.Monitor.Query/src/Generated/Models/BatchResponse.cs

@@ -5,5 +5,6 @@ namespace Azure.Monitor.Query.Models
 {
     internal partial class BatchQueryRequest
     {
+        internal LogsQueryOptions Options { get; set; }
     }
 }

sdk/monitor/Azure.Monitor.Query/src/Generated/QueryRestClient.cs

@@ -1,24 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using Azure.Core;
-
 namespace Azure.Monitor.Query.Models
 {
     internal partial class BatchQueryResponse
     {
-        /// <summary> Contains the tables, columns &amp; rows resulting from the query or the error details if the query failed. </summary>
-        [CodeGenMember("Body")]
-        private LogsBatchQueryResult _body;
-
-        public LogsBatchQueryResult Body
-        {
-            get
-            {
-                _body.HasFailed = Status >= 400;
-                _body.Id = Id;
-                return _body;
-            }
-        }
     }
 }