Skip to content

Commit 63b90ff

Browse files
pakrympakrym
authored
Add logs query options (Azure#20598)
1 parent a73238a commit 63b90ff

23 files changed

+2300
-29
lines changed

‎sdk/core/Azure.Core.TestFramework/src/MockTransport.cs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ namespace Azure.Core.TestFramework
1313
public class MockTransport : HttpPipelineTransport
1414
{
1515
private readonly object _syncObj = new object();
16-
private readonly Func<MockRequest, CancellationToken, MockResponse> _responseFactory;
16+
private readonly Func<HttpMessage, MockResponse> _responseFactory;
1717

1818
public AsyncGate<MockRequest, MockResponse> RequestGate { get; }
1919

@@ -29,7 +29,7 @@ public MockTransport()
2929
public MockTransport(params MockResponse[] responses)
3030
{
3131
var requestIndex = 0;
32-
_responseFactory = (_, _) =>
32+
_responseFactory = _ =>
3333
{
3434
lock (_syncObj)
3535
{
@@ -38,15 +38,17 @@ public MockTransport(params MockResponse[] responses)
3838
};
3939
}
4040

41-
public MockTransport(Func<MockRequest, MockResponse> responseFactory): this((req, _) => responseFactory(req))
41+
public MockTransport(Func<MockRequest, MockResponse> responseFactory): this(req => responseFactory((MockRequest)req.Request))
4242
{
4343
}
4444

45-
public MockTransport(Func<MockRequest, CancellationToken, MockResponse> responseFactory)
45+
private MockTransport(Func<HttpMessage, MockResponse> responseFactory)
4646
{
4747
_responseFactory = responseFactory;
4848
}
4949

50+
public static MockTransport FromMessageCallback(Func<HttpMessage, MockResponse> responseFactory) => new MockTransport(responseFactory);
51+
5052
public override Request CreateRequest()
5153
=> new MockRequest();
5254

@@ -86,7 +88,7 @@ private async Task ProcessCore(HttpMessage message)
8688
}
8789
else
8890
{
89-
message.Response = _responseFactory(request, message.CancellationToken);
91+
message.Response = _responseFactory(message);
9092
}
9193

9294
message.Response.ClientRequestId = request.ClientRequestId;

‎sdk/core/Azure.Core/tests/ResponseBodyPolicyTests.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -174,9 +174,9 @@ public void CanOverrideDefaultNetworkTimeout()
174174
{
175175
var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
176176

177-
MockTransport mockTransport = new MockTransport((_, ct) =>
177+
MockTransport mockTransport = MockTransport.FromMessageCallback(message =>
178178
{
179-
tcs.Task.Wait(ct);
179+
tcs.Task.Wait(message.CancellationToken);
180180
return null;
181181
});
182182

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

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ namespace Azure.Monitor.Query
33
public partial class LogsBatchQuery
44
{
55
protected LogsBatchQuery() { }
6-
public virtual string AddQuery(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?)) { throw null; }
6+
public virtual string AddQuery(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), Azure.Monitor.Query.LogsQueryOptions options = null) { throw null; }
77
public virtual Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResult> Submit(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
88
public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsBatchQueryResult>> SubmitAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
99
}
@@ -13,10 +13,10 @@ protected LogsClient() { }
1313
public LogsClient(Azure.Core.TokenCredential credential) { }
1414
public LogsClient(Azure.Core.TokenCredential credential, Azure.Monitor.Query.LogsClientOptions options) { }
1515
public virtual Azure.Monitor.Query.LogsBatchQuery CreateBatchQuery() { throw null; }
16-
public virtual Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult> Query(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
17-
public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult>> QueryAsync(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
18-
public virtual System.Threading.Tasks.Task<Azure.Response<System.Collections.Generic.IReadOnlyList<T>>> QueryAsync<T>(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
19-
public virtual Azure.Response<System.Collections.Generic.IReadOnlyList<T>> Query<T>(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
16+
public virtual Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult> Query(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
17+
public virtual System.Threading.Tasks.Task<Azure.Response<Azure.Monitor.Query.Models.LogsQueryResult>> QueryAsync(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
18+
public virtual System.Threading.Tasks.Task<Azure.Response<System.Collections.Generic.IReadOnlyList<T>>> QueryAsync<T>(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
19+
public virtual Azure.Response<System.Collections.Generic.IReadOnlyList<T>> Query<T>(string workspaceId, string query, System.TimeSpan? timeSpan = default(System.TimeSpan?), Azure.Monitor.Query.LogsQueryOptions options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
2020
}
2121
public partial class LogsClientOptions : Azure.Core.ClientOptions
2222
{
@@ -26,6 +26,12 @@ public enum ServiceVersion
2626
V1 = 0,
2727
}
2828
}
29+
public partial class LogsQueryOptions
30+
{
31+
public LogsQueryOptions() { }
32+
public bool IncludeStatistics { get { throw null; } set { } }
33+
public System.TimeSpan? Timeout { get { throw null; } set { } }
34+
}
2935
public partial class MetricsClient
3036
{
3137
protected MetricsClient() { }
@@ -59,6 +65,7 @@ public partial class LogsQueryResult
5965
{
6066
internal LogsQueryResult() { }
6167
public Azure.Monitor.Query.Models.LogsQueryResultTable PrimaryTable { get { throw null; } }
68+
public System.Text.Json.JsonElement Statistics { get { throw null; } }
6269
public System.Collections.Generic.IReadOnlyList<Azure.Monitor.Query.Models.LogsQueryResultTable> Tables { get { throw null; } }
6370
}
6471
public partial class LogsQueryResultColumn

‎sdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsQueryResult.Serialization.cs‎

Lines changed: 7 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎sdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsQueryResult.cs‎

Lines changed: 4 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎sdk/monitor/Azure.Monitor.Query/src/LogsClient.cs‎

Lines changed: 72 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
using System;
55
using System.Collections.Generic;
6+
using System.Text;
7+
using System.Text.Json;
68
using System.Threading;
79
using System.Threading.Tasks;
810
using Azure.Core;
@@ -38,27 +40,27 @@ protected LogsClient()
3840
{
3941
}
4042

41-
public virtual Response<IReadOnlyList<T>> Query<T>(string workspaceId, string query, TimeSpan? timeSpan = null, CancellationToken cancellationToken = default)
43+
public virtual Response<IReadOnlyList<T>> Query<T>(string workspaceId, string query, TimeSpan? timeSpan = null, LogsQueryOptions options = null, CancellationToken cancellationToken = default)
4244
{
43-
Response<LogsQueryResult> response = Query(workspaceId, query, timeSpan, cancellationToken);
45+
Response<LogsQueryResult> response = Query(workspaceId, query, timeSpan, options, cancellationToken);
4446

4547
return Response.FromValue(_rowBinder.BindResults<T>(response), response.GetRawResponse());
4648
}
4749

48-
public virtual async Task<Response<IReadOnlyList<T>>> QueryAsync<T>(string workspaceId, string query, TimeSpan? timeSpan = null, CancellationToken cancellationToken = default)
50+
public virtual async Task<Response<IReadOnlyList<T>>> QueryAsync<T>(string workspaceId, string query, TimeSpan? timeSpan = null, LogsQueryOptions options = null, CancellationToken cancellationToken = default)
4951
{
50-
Response<LogsQueryResult> response = await QueryAsync(workspaceId, query, timeSpan, cancellationToken).ConfigureAwait(false);
52+
Response<LogsQueryResult> response = await QueryAsync(workspaceId, query, timeSpan, options, cancellationToken).ConfigureAwait(false);
5153

5254
return Response.FromValue(_rowBinder.BindResults<T>(response), response.GetRawResponse());
5355
}
5456

55-
public virtual Response<LogsQueryResult> Query(string workspaceId, string query, TimeSpan? timeSpan = null, CancellationToken cancellationToken = default)
57+
public virtual Response<LogsQueryResult> Query(string workspaceId, string query, TimeSpan? timeSpan = null, LogsQueryOptions options = null, CancellationToken cancellationToken = default)
5658
{
5759
using DiagnosticScope scope = _clientDiagnostics.CreateScope($"{nameof(LogsClient)}.{nameof(Query)}");
5860
scope.Start();
5961
try
6062
{
61-
return _queryClient.Execute(workspaceId, CreateQueryBody(query, timeSpan), null, cancellationToken);
63+
return ExecuteAsync(workspaceId, query, timeSpan, options, false, cancellationToken).EnsureCompleted();
6264
}
6365
catch (Exception e)
6466
{
@@ -67,13 +69,13 @@ public virtual Response<LogsQueryResult> Query(string workspaceId, string query,
6769
}
6870
}
6971

70-
public virtual async Task<Response<LogsQueryResult>> QueryAsync(string workspaceId, string query, TimeSpan? timeSpan = null, CancellationToken cancellationToken = default)
72+
public virtual async Task<Response<LogsQueryResult>> QueryAsync(string workspaceId, string query, TimeSpan? timeSpan = null, LogsQueryOptions options = null, CancellationToken cancellationToken = default)
7173
{
7274
using DiagnosticScope scope = _clientDiagnostics.CreateScope($"{nameof(LogsClient)}.{nameof(Query)}");
7375
scope.Start();
7476
try
7577
{
76-
return await _queryClient.ExecuteAsync(workspaceId, CreateQueryBody(query, timeSpan), null, cancellationToken).ConfigureAwait(false);
78+
return await ExecuteAsync(workspaceId, query, timeSpan, options, true, cancellationToken).ConfigureAwait(false);
7779
}
7880
catch (Exception e)
7981
{
@@ -87,15 +89,76 @@ public virtual LogsBatchQuery CreateBatchQuery()
8789
return new LogsBatchQuery(_clientDiagnostics, _queryClient, _rowBinder);
8890
}
8991

90-
internal static QueryBody CreateQueryBody(string query, TimeSpan? timeSpan)
92+
internal static QueryBody CreateQueryBody(string query, TimeSpan? timeSpan, LogsQueryOptions options, out string prefer)
9193
{
9294
var queryBody = new QueryBody(query);
9395
if (timeSpan != null)
9496
{
9597
queryBody.Timespan = TypeFormatters.ToString(timeSpan.Value, "P");
9698
}
9799

100+
prefer = null;
101+
102+
if (options?.Timeout is TimeSpan timeout)
103+
{
104+
prefer = "wait=" + (int) timeout.TotalSeconds;
105+
}
106+
107+
if (options?.IncludeStatistics == true)
108+
{
109+
prefer += " include-statistics=true";
110+
}
111+
98112
return queryBody;
99113
}
114+
115+
private async Task<Response<LogsQueryResult>> ExecuteAsync(string workspaceId, string query, TimeSpan? timeSpan, LogsQueryOptions options, bool async, CancellationToken cancellationToken = default)
116+
{
117+
if (workspaceId == null)
118+
{
119+
throw new ArgumentNullException(nameof(workspaceId));
120+
}
121+
122+
QueryBody queryBody = CreateQueryBody(query, timeSpan, options, out string prefer);
123+
using var message = _queryClient.CreateExecuteRequest(workspaceId, queryBody, prefer);
124+
125+
if (options?.Timeout != null)
126+
{
127+
message.SetProperty("NetworkTimeoutOverride", options.Timeout);
128+
}
129+
130+
if (async)
131+
{
132+
await _pipeline.SendAsync(message, cancellationToken).ConfigureAwait(false);
133+
}
134+
else
135+
{
136+
_pipeline.Send(message, cancellationToken);
137+
}
138+
139+
switch (message.Response.Status)
140+
{
141+
case 200:
142+
{
143+
using var document = async ?
144+
await JsonDocument.ParseAsync(message.Response.ContentStream, default, cancellationToken).ConfigureAwait(false) :
145+
JsonDocument.Parse(message.Response.ContentStream, default);
146+
147+
LogsQueryResult value = LogsQueryResult.DeserializeLogsQueryResult(document.RootElement);
148+
return Response.FromValue(value, message.Response);
149+
}
150+
default:
151+
{
152+
if (async)
153+
{
154+
throw await _clientDiagnostics.CreateRequestFailedExceptionAsync(message.Response).ConfigureAwait(false);
155+
}
156+
else
157+
{
158+
throw _clientDiagnostics.CreateRequestFailedException(message.Response);
159+
}
160+
}
161+
}
162+
}
100163
}
101164
}

‎sdk/monitor/Azure.Monitor.Query/src/LogsQueryOptions.cs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
using System;
5+
6+
namespace Azure.Monitor.Query
7+
{
8+
public class LogsQueryOptions
9+
{
10+
public TimeSpan? Timeout { get; set; }
11+
public bool IncludeStatistics { get; set; }
12+
}
13+
}

‎sdk/monitor/Azure.Monitor.Query/src/Models/LogsBatchQuery.cs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,21 @@ protected LogsBatchQuery()
3030
{
3131
}
3232

33-
public virtual string AddQuery(string workspaceId, string query, TimeSpan? timeSpan = null)
33+
public virtual string AddQuery(string workspaceId, string query, TimeSpan? timeSpan = null, LogsQueryOptions options = null)
3434
{
3535
var id = _counter.ToString("G", CultureInfo.InvariantCulture);
3636
_counter++;
37-
_batch.Requests.Add(new LogQueryRequest()
37+
var logQueryRequest = new LogQueryRequest()
3838
{
3939
Id = id,
40-
Body = LogsClient.CreateQueryBody(query, timeSpan),
40+
Body = LogsClient.CreateQueryBody(query, timeSpan, options, out string prefer),
4141
Workspace = workspaceId
42-
});
42+
};
43+
if (prefer != null)
44+
{
45+
logQueryRequest.Headers.Add("prefer", prefer);
46+
}
47+
_batch.Requests.Add(logQueryRequest);
4348
return id;
4449
}
4550

‎sdk/monitor/Azure.Monitor.Query/src/Models/LogsQueryResult.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// Licensed under the MIT License.
33

44
using System.Linq;
5+
using System.Text.Json;
56
using Azure.Core;
67

78
namespace Azure.Monitor.Query.Models
@@ -12,6 +13,10 @@ public partial class LogsQueryResult
1213
// TODO: Handle not found
1314
public LogsQueryResultTable PrimaryTable => Tables.Single(t => t.Name == "PrimaryResult");
1415

16+
#pragma warning disable AZC0014
17+
public JsonElement Statistics { get; }
18+
#pragma warning restore AZC0014
19+
1520
internal ErrorDetails Error { get; }
1621
}
1722
}

‎sdk/monitor/Azure.Monitor.Query/src/autorest.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,14 @@ directive:
2929
}
3030
```
3131
32-
### Rename errors property
32+
### Rename errors property and add statistics
3333
3434
``` yaml
3535
directive:
3636
- from: swagger-document
3737
where: $.definitions.queryResults
3838
transform: >
39+
$.properties["statistics"] = { "type": "object" };
3940
$.properties["error"] = $.properties["errors"];
4041
delete $.properties["errors"];
41-
4242
```

0 commit comments

Comments
 (0)