Skip to content

Commit 2ff9d05

Browse files
kinelskikinelski
authored
[Core] Wiring up rehydration token logic in data-plane (#49224)
1 parent 23ebea4 commit 2ff9d05

File tree

41 files changed

+371
-6
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+371
-6
lines changed

‎sdk/communication/Azure.Communication.Email/src/EmailSendOperation.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,10 @@ async ValueTask<OperationState<EmailSendResult>> IOperation<EmailSendResult>.Upd
184184
return OperationState<EmailSendResult>.Pending(rawResponse);
185185
}
186186

187+
// This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
188+
RehydrationToken IOperation<EmailSendResult>.GetRehydrationToken() =>
189+
throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
190+
187191
private static IDictionary<string, string> CreateAdditionalInformation(ErrorDetail error)
188192
{
189193
if (string.IsNullOrEmpty(error.ToString()))

‎sdk/confidentialledger/Azure.Security.CodeTransparency/src/CreateEntryOperation.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,5 +125,9 @@ async ValueTask<OperationState> IOperation.UpdateStateAsync(bool async, Cancella
125125
}
126126
}
127127
}
128+
129+
// This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
130+
RehydrationToken IOperation.GetRehydrationToken() =>
131+
throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
128132
}
129133
}

‎sdk/confidentialledger/Azure.Security.ConfidentialLedger/src/PostLedgerEntryOperation.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ async ValueTask<OperationState> IOperation.UpdateStateAsync(bool async, Cancella
7777
return OperationState.Pending(statusResponse);
7878
}
7979

80+
// This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
81+
RehydrationToken IOperation.GetRehydrationToken() =>
82+
throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
83+
8084
/// <summary>
8185
/// The transactionId of the posted ledger entry.
8286
/// </summary>

‎sdk/core/Azure.Core/src/Shared/NextLinkOperationImplementation.cs

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -58,17 +58,21 @@ public static IOperation Create(
5858
apiVersionStr = !skipApiVersionOverride && TryGetApiVersion(startRequestUri, out ReadOnlySpan<char> apiVersion) ? apiVersion.ToString() : null;
5959
}
6060
var headerSource = GetHeaderSource(requestMethod, startRequestUri, response, apiVersionStr, out string nextRequestUri, out bool isNextRequestPolling);
61-
if (headerSource == HeaderSource.None && IsFinalState(response, headerSource, out var failureState, out _))
62-
{
63-
return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response));
64-
}
6561

6662
string? lastKnownLocation;
6763
if (!response.Headers.TryGetValue("Location", out lastKnownLocation))
6864
{
6965
lastKnownLocation = null;
7066
}
71-
return new NextLinkOperationImplementation(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, apiVersionStr, isNextRequestPolling : isNextRequestPolling);
67+
68+
NextLinkOperationImplementation operation = new(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, apiVersionStr, isNextRequestPolling: isNextRequestPolling);
69+
70+
if (headerSource == HeaderSource.None && IsFinalState(response, headerSource, out var failureState, out _))
71+
{
72+
return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response), operation);
73+
}
74+
75+
return operation;
7276
}
7377

7478
public static IOperation<T> Create<T>(
@@ -647,12 +651,17 @@ private class CompletedOperation : IOperation
647651
{
648652
private readonly OperationState _operationState;
649653

650-
public CompletedOperation(OperationState operationState)
654+
private readonly NextLinkOperationImplementation _operation;
655+
656+
public CompletedOperation(OperationState operationState, NextLinkOperationImplementation operation)
651657
{
652658
_operationState = operationState;
659+
_operation = operation;
653660
}
654661

655662
public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(_operationState);
663+
664+
public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
656665
}
657666

658667
private sealed class OperationToOperationOfT<T> : IOperation<T>
@@ -685,6 +694,8 @@ public async ValueTask<OperationState<T>> UpdateStateAsync(bool async, Cancellat
685694

686695
return OperationState<T>.Pending(state.RawResponse);
687696
}
697+
698+
public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
688699
}
689700
}
690701
}

‎sdk/core/Azure.Core/src/Shared/OperationInternal.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ public OperationToOperationOfTProxy(IOperation operation)
119119
_operation = operation;
120120
}
121121

122+
public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
123+
122124
public async ValueTask<OperationState<VoidValue>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
123125
{
124126
var state = await _operation.UpdateStateAsync(async, cancellationToken).ConfigureAwait(false);
@@ -172,6 +174,11 @@ internal interface IOperation
172174
/// </list>
173175
/// </returns>
174176
ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken);
177+
178+
/// <summary>
179+
/// Get a token that can be used to rehydrate the operation.
180+
/// </summary>
181+
RehydrationToken GetRehydrationToken();
175182
}
176183

177184
/// <summary>

‎sdk/core/Azure.Core/src/Shared/OperationInternalOfT.cs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,10 @@ private class FinalOperation : IOperation<T>
288288
{
289289
public ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
290290
=> throw new NotSupportedException("The operation has already completed");
291+
292+
// Unreachable path. _operation.GetRehydrationToken() is never invoked.
293+
public RehydrationToken GetRehydrationToken()
294+
=> throw new NotSupportedException($"Getting the rehydration token of a {nameof(FinalOperation)} is not supported");
291295
}
292296
}
293297

@@ -328,6 +332,11 @@ internal interface IOperation<T>
328332
/// </list>
329333
/// </returns>
330334
ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken);
335+
336+
/// <summary>
337+
/// Get a token that can be used to rehydrate the operation.
338+
/// </summary>
339+
RehydrationToken GetRehydrationToken();
331340
}
332341

333342
/// <summary>

‎sdk/core/Azure.Core/src/Shared/ProtocolOperation.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,11 @@ internal ProtocolOperation(ClientDiagnostics clientDiagnostics, HttpPipeline pip
3030
public override string Id => throw new NotSupportedException();
3131
#pragma warning restore CA1822
3232

33+
/// <inheritdoc />
34+
public override RehydrationToken? GetRehydrationToken() => ((IOperation<T>)this).GetRehydrationToken();
35+
36+
RehydrationToken IOperation<T>.GetRehydrationToken() => _nextLinkOperation.GetRehydrationToken();
37+
3338
/// <inheritdoc />
3439
public override T Value => _operation.Value;
3540

‎sdk/core/Azure.Core/src/Shared/ProtocolOperationHelpers.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ public ConvertOperation(Operation<TFrom> operation, ClientDiagnostics diagnostic
7676
public override TTo Value => GetOrCreateValue();
7777
public override bool HasValue => _operation.HasValue;
7878
public override bool HasCompleted => _operation.HasCompleted;
79+
public override RehydrationToken? GetRehydrationToken() => _operation.GetRehydrationToken();
7980
public override Response GetRawResponse() => _operation.GetRawResponse();
8081

8182
public override Response UpdateStatus(CancellationToken cancellationToken = default)

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -518,6 +518,9 @@ public TestOperation(Func<bool, CancellationToken, ValueTask<OperationState<int>
518518
}
519519

520520
public ValueTask<OperationState<int>> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
521+
522+
public RehydrationToken GetRehydrationToken() =>
523+
throw new NotImplementedException();
521524
}
522525

523526
private class CallCountStrategy : DelayStrategy

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,9 @@ public TestOperation(Func<bool, CancellationToken, ValueTask<OperationState>> up
438438
}
439439

440440
public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
441+
442+
public RehydrationToken GetRehydrationToken() =>
443+
throw new NotImplementedException();
441444
}
442445

443446
private class CallCountStrategy : DelayStrategy

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ public void WaitForCompletionResponseCancelled()
5757
private class EndlessOperation : IOperation
5858
{
5959
public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(OperationState.Pending(new MockResponse(200)));
60+
61+
public RehydrationToken GetRehydrationToken() =>
62+
throw new NotImplementedException();
6063
}
6164

6265
private class TestDelayStrategy : DelayStrategy

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

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
using System;
5+
using System.IO;
6+
using System.Text;
7+
using System.Threading.Tasks;
8+
using Azure.Core.Pipeline;
9+
using Azure.Core.TestFramework;
10+
using NUnit.Framework;
11+
12+
namespace Azure.Core.Tests
13+
{
14+
public class ProtocolOperationTests
15+
{
16+
[Test]
17+
public void GetRehydrationToken()
18+
{
19+
const string OperationId = "oper-ati-oni-d00";
20+
const string RequestUri = "https://www.example.com/request";
21+
const string LocationUri = $"https://www.example.com/{OperationId}";
22+
23+
using var request = new MockRequest() { Method = RequestMethod.Post };
24+
using var response = new MockResponse(200);
25+
var transport = new MockTransport();
26+
var pipeline = new HttpPipeline(transport);
27+
28+
request.Uri.Reset(new Uri(RequestUri));
29+
response.AddHeader("Location", LocationUri);
30+
31+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
32+
var token = operation.GetRehydrationToken();
33+
34+
Assert.True(token.HasValue);
35+
Assert.AreEqual(OperationId, token.Value.Id);
36+
Assert.AreEqual(NextLinkOperationImplementation.RehydrationTokenVersion, token.Value.Version);
37+
Assert.AreEqual("Location", token.Value.HeaderSource);
38+
Assert.AreEqual(LocationUri, token.Value.NextRequestUri);
39+
Assert.AreEqual(RequestUri, token.Value.InitialUri);
40+
Assert.AreEqual(RequestMethod.Post, token.Value.RequestMethod);
41+
Assert.AreEqual(LocationUri, token.Value.LastKnownLocation);
42+
Assert.AreEqual("Location", token.Value.FinalStateVia);
43+
}
44+
45+
[Test]
46+
public void GetRehydrationTokenWithCompletedOperation()
47+
{
48+
const string RequestUri = "https://www.example.com/request";
49+
50+
using var request = new MockRequest() { Method = RequestMethod.Post };
51+
using var response = new MockResponse(200);
52+
var transport = new MockTransport();
53+
var pipeline = new HttpPipeline(transport);
54+
55+
request.Uri.Reset(new Uri(RequestUri));
56+
57+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
58+
var token = operation.GetRehydrationToken();
59+
60+
Assert.True(token.HasValue);
61+
Assert.AreEqual("NOT_SET", token.Value.Id);
62+
Assert.AreEqual(NextLinkOperationImplementation.RehydrationTokenVersion, token.Value.Version);
63+
Assert.AreEqual("None", token.Value.HeaderSource);
64+
Assert.AreEqual(RequestUri, token.Value.NextRequestUri);
65+
Assert.AreEqual(RequestUri, token.Value.InitialUri);
66+
Assert.AreEqual(RequestMethod.Post, token.Value.RequestMethod);
67+
Assert.Null(token.Value.LastKnownLocation);
68+
Assert.AreEqual("Location", token.Value.FinalStateVia);
69+
}
70+
71+
[Test]
72+
[TestCase(true)]
73+
[TestCase(false)]
74+
public async Task RehydrationOperationCanPoll(bool isAsync)
75+
{
76+
const string RequestUri = "https://www.example.com/request";
77+
const string LocationUri = "https://www.example.com/oper-ati-oni-d00";
78+
79+
using var request = new MockRequest() { Method = RequestMethod.Post };
80+
using var response = new MockResponse(200);
81+
var transport = new MockTransport(response);
82+
var pipeline = new HttpPipeline(transport);
83+
84+
request.Uri.Reset(new Uri(RequestUri));
85+
response.AddHeader("Location", LocationUri);
86+
87+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
88+
var token = operation.GetRehydrationToken();
89+
_ = isAsync
90+
? await Operation.RehydrateAsync(pipeline, token.Value)
91+
: Operation.Rehydrate(pipeline, token.Value);
92+
var rehydrationUpdateRequest = transport.SingleRequest;
93+
94+
Assert.AreEqual(LocationUri, rehydrationUpdateRequest.Uri.ToString());
95+
Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
96+
}
97+
98+
[Test]
99+
[TestCase(true)]
100+
[TestCase(false)]
101+
public async Task RehydrationOperationWithCompletedOperationCanPoll(bool isAsync)
102+
{
103+
const string RequestUri = "https://www.example.com/request";
104+
105+
using var request = new MockRequest() { Method = RequestMethod.Post };
106+
using var response = new MockResponse(200);
107+
var transport = new MockTransport(response);
108+
var pipeline = new HttpPipeline(transport);
109+
110+
request.Uri.Reset(new Uri(RequestUri));
111+
112+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
113+
var token = operation.GetRehydrationToken();
114+
_ = isAsync
115+
? await Operation.RehydrateAsync(pipeline, token.Value)
116+
: Operation.Rehydrate(pipeline, token.Value);
117+
var rehydrationUpdateRequest = transport.SingleRequest;
118+
119+
Assert.AreEqual(RequestUri, rehydrationUpdateRequest.Uri.ToString());
120+
Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
121+
}
122+
123+
[Test]
124+
[TestCase(true)]
125+
[TestCase(false)]
126+
public async Task RehydrationOperationOfTCanPoll(bool isAsync)
127+
{
128+
const string RequestUri = "https://www.example.com/request";
129+
const string LocationUri = "https://www.example.com/oper-ati-oni-d00";
130+
131+
using var request = new MockRequest() { Method = RequestMethod.Post };
132+
using var response = new MockResponse(200);
133+
var transport = new MockTransport(response);
134+
var pipeline = new HttpPipeline(transport);
135+
using var responseContentStream = new MemoryStream(Encoding.UTF8.GetBytes("""
136+
{
137+
"IntValue": 1,
138+
"StringValue": "one"
139+
}
140+
"""));
141+
142+
request.Uri.Reset(new Uri(RequestUri));
143+
response.AddHeader("Location", LocationUri);
144+
response.ContentStream = responseContentStream;
145+
146+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
147+
var token = operation.GetRehydrationToken();
148+
var rehydrationOperation = isAsync
149+
? await Operation.RehydrateAsync<MockJsonModel>(pipeline, token.Value)
150+
: Operation.Rehydrate<MockJsonModel>(pipeline, token.Value);
151+
var rehydrationUpdateRequest = transport.SingleRequest;
152+
153+
Assert.AreEqual(LocationUri, rehydrationUpdateRequest.Uri.ToString());
154+
Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
155+
156+
Assert.True(rehydrationOperation.HasValue);
157+
Assert.AreEqual(rehydrationOperation.Value.IntValue, 1);
158+
Assert.AreEqual(rehydrationOperation.Value.StringValue, "one");
159+
}
160+
161+
[Test]
162+
[TestCase(true)]
163+
[TestCase(false)]
164+
public async Task RehydrationOperationOfTWithCompletedOperationCanPoll(bool isAsync)
165+
{
166+
const string RequestUri = "https://www.example.com/request";
167+
168+
using var request = new MockRequest() { Method = RequestMethod.Post };
169+
using var response = new MockResponse(200);
170+
var transport = new MockTransport(response);
171+
var pipeline = new HttpPipeline(transport);
172+
using var responseContentStream = new MemoryStream(Encoding.UTF8.GetBytes("""
173+
{
174+
"IntValue": 1,
175+
"StringValue": "one"
176+
}
177+
"""));
178+
179+
request.Uri.Reset(new Uri(RequestUri));
180+
response.ContentStream = responseContentStream;
181+
182+
var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
183+
var token = operation.GetRehydrationToken();
184+
var rehydrationOperation = isAsync
185+
? await Operation.RehydrateAsync<MockJsonModel>(pipeline, token.Value)
186+
: Operation.Rehydrate<MockJsonModel>(pipeline, token.Value);
187+
var rehydrationUpdateRequest = transport.SingleRequest;
188+
189+
Assert.AreEqual(RequestUri, rehydrationUpdateRequest.Uri.ToString());
190+
Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
191+
192+
Assert.True(rehydrationOperation.HasValue);
193+
Assert.AreEqual(rehydrationOperation.Value.IntValue, 1);
194+
Assert.AreEqual(rehydrationOperation.Value.StringValue, "one");
195+
}
196+
}
197+
}

‎sdk/formrecognizer/Azure.AI.FormRecognizer/src/AnalyzeDocumentOperation.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,5 +204,9 @@ async ValueTask<OperationState<AnalyzeResult>> IOperation<AnalyzeResult>.UpdateS
204204

205205
return OperationState<AnalyzeResult>.Pending(rawResponse);
206206
}
207+
208+
// This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
209+
RehydrationToken IOperation<AnalyzeResult>.GetRehydrationToken() =>
210+
throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
207211
}
208212
}

‎sdk/formrecognizer/Azure.AI.FormRecognizer/src/BuildDocumentClassifierOperation.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,5 +193,9 @@ async ValueTask<OperationState<DocumentClassifierDetails>> IOperation<DocumentCl
193193

194194
return OperationState<DocumentClassifierDetails>.Pending(rawResponse);
195195
}
196+
197+
// This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
198+
RehydrationToken IOperation<DocumentClassifierDetails>.GetRehydrationToken() =>
199+
throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
196200
}
197201
}

0 commit comments

Comments
 (0)