Resolve several Key Vault backlog issues

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

@@ -174,6 +174,10 @@ public async ValueTask<Response<T>> WaitForCompletionAsync(TimeSpan pollingInter
         /// <param name="state">The <see cref="OperationState{T}"/> used to set <see cref="OperationInternalBase.HasCompleted"/> and other members.</param>
         public void SetState(OperationState<T> state)
         {
+            if (state.HasCompleted && state.HasSucceeded)
+            {
+                Value = state.Value!;
+            }
             ApplyStateAsync(false, state.RawResponse, state.HasCompleted, state.HasSucceeded, state.OperationFailedException).EnsureCompleted();
         }
 

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

@@ -87,6 +87,67 @@ public void RawResponseInitialization()
             }
         }
 
+        [Test]
+        public void SetStateSucceeds()
+        {
+            var operationInternal = CreateOperation(isOfT, UpdateResult.Pending);
+            if (operationInternal is OperationInternal oi)
+            {
+                oi.SetState(OperationState.Success(mockResponse));
+            }
+            else if (operationInternal is OperationInternal<int> oit)
+            {
+                oit.SetState(OperationState<int>.Success(mockResponse, 1));
+            }
+
+            Assert.IsTrue(operationInternal.HasCompleted);
+            if (operationInternal is OperationInternal<int> oit2)
+            {
+                Assert.IsTrue(oit2.HasValue);
+                Assert.AreEqual(1, oit2.Value);
+            }
+        }
+
+        [Test]
+        public void SetStateIsPending()
+        {
+            var operationInternal = CreateOperation(isOfT, UpdateResult.Pending);
+            if (operationInternal is OperationInternal oi)
+            {
+                oi.SetState(OperationState.Pending(mockResponse));
+            }
+            else if (operationInternal is OperationInternal<int> oit)
+            {
+                oit.SetState(OperationState<int>.Pending(mockResponse));
+            }
+
+            Assert.IsFalse(operationInternal.HasCompleted);
+            if (operationInternal is OperationInternal<int> oit2)
+            {
+                Assert.IsFalse(oit2.HasValue);
+            }
+        }
+
+        [Test]
+        public void SetStateFails()
+        {
+            var operationInternal = CreateOperation(isOfT, UpdateResult.Pending);
+            if (operationInternal is OperationInternal oi)
+            {
+                Assert.Throws<RequestFailedException>(() => oi.SetState(OperationState.Failure(mockResponse)));
+            }
+            else if (operationInternal is OperationInternal<int> oit)
+            {
+                Assert.Throws<RequestFailedException>(() => oit.SetState(OperationState<int>.Failure(mockResponse)));
+            }
+
+            Assert.IsTrue(operationInternal.HasCompleted);
+            if (operationInternal is OperationInternal<int> oit2)
+            {
+                Assert.IsFalse(oit2.HasValue);
+            }
+        }
+
         [Test]
         public async Task UpdateStatusWhenOperationIsPending([Values(true, false)] bool async)
         {

sdk/keyvault/Azure.Security.KeyVault.Certificates/src/Azure.Security.KeyVault.Certificates.csproj

@@ -35,6 +35,8 @@
     <Compile Include="$(AzureCoreSharedSources)ContentTypeUtilities.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScope.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)OperationHelpers.cs" LinkBase="Shared" />
+    <Compile Include="$(AzureCoreSharedSources)OperationInternal.cs" LinkBase="Shared" />
+    <Compile Include="$(AzureCoreSharedSources)OperationInternalBase.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)PageResponseEnumerator.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScopeFactory.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)TaskExtensions.cs" LinkBase="Shared" />

sdk/keyvault/Azure.Security.KeyVault.Certificates/src/CertificateProperties.cs

@@ -158,18 +158,14 @@ internal void ReadProperty(JsonProperty prop)
             }
         }
 
-        private void ParseId(Uri idToParse)
+        private void ParseId(Uri id)
         {
-            // We expect an identifier with either 3 or 4 segments: host + collection + name [+ version]
-            if (idToParse.Segments.Length != 3 && idToParse.Segments.Length != 4)
-                throw new ArgumentException(string.Format(CultureInfo.InvariantCulture, "Invalid ObjectIdentifier: {0}. Bad number of segments: {1}", idToParse, idToParse.Segments.Length));
+            KeyVaultIdentifier identifier = KeyVaultIdentifier.ParseWithCollection(id, "certificates");
 
-            if (!string.Equals(idToParse.Segments[1], "certificates" + "/", StringComparison.OrdinalIgnoreCase))
-                throw new ArgumentException(string.Format(CultureInfo.InvariantCulture, "Invalid ObjectIdentifier: {0}. segment [1] should be 'certificates/', found '{1}'", idToParse, idToParse.Segments[1]));
-
-            VaultUri = new Uri($"{idToParse.Scheme}://{idToParse.Authority}");
-            Name = idToParse.Segments[2].Trim('/');
-            Version = (idToParse.Segments.Length == 4) ? idToParse.Segments[3].TrimEnd('/') : null;
+            Id = id;
+            VaultUri = identifier.VaultUri;
+            Name = identifier.Name;
+            Version = identifier.Version;
         }
     }
 }

sdk/keyvault/Azure.Security.KeyVault.Certificates/src/DeleteCertificateOperation.cs

@@ -2,35 +2,39 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
-using Azure.Core.Pipeline;
 
 namespace Azure.Security.KeyVault.Certificates
 {
     /// <summary>
     /// A long-running operation for <see cref="CertificateClient.StartDeleteCertificate(string, CancellationToken)"/> or <see cref="CertificateClient.StartDeleteCertificateAsync(string, CancellationToken)"/>.
     /// </summary>
-    public class DeleteCertificateOperation : Operation<DeletedCertificate>
+    public class DeleteCertificateOperation : Operation<DeletedCertificate>, IOperation
     {
         private static readonly TimeSpan s_defaultPollingInterval = TimeSpan.FromSeconds(2);
 
         private readonly KeyVaultPipeline _pipeline;
+        private readonly OperationInternal _operationInternal;
         private readonly DeletedCertificate _value;
-        private Response _response;
-        private bool _completed;
 
         internal DeleteCertificateOperation(KeyVaultPipeline pipeline, Response<DeletedCertificate> response)
         {
             _pipeline = pipeline;
             _value = response.Value ?? throw new InvalidOperationException("The response does not contain a value.");
-            _response = response.GetRawResponse();
+            _operationInternal = new(_pipeline.Diagnostics, this, response.GetRawResponse(), nameof(DeleteCertificateOperation), new[]
+            {
+                new KeyValuePair<string, string>("secret", _value.Name), // Retained for backward compatibility.
+                new KeyValuePair<string, string>("certificate", _value.Name),
+            });
 
-            // The recoveryId is only returned if soft-delete is enabled.
+            // The recoveryId is only returned if soft delete is enabled.
             if (_value.RecoveryId is null)
             {
-                _completed = true;
+                // If soft delete is not enabled, deleting is immediate so set success accordingly.
+                _operationInternal.SetState(OperationState.Success(response.GetRawResponse()));
             }
         }
 
@@ -50,33 +54,20 @@ protected DeleteCertificateOperation() {}
         public override DeletedCertificate Value => _value;
 
         /// <inheritdoc/>
-        public override bool HasCompleted => _completed;
+        public override bool HasCompleted => _operationInternal.HasCompleted;
 
         /// <inheritdoc/>
         public override bool HasValue => true;
 
         /// <inheritdoc/>
-        public override Response GetRawResponse() => _response;
+        public override Response GetRawResponse() => _operationInternal.RawResponse;
 
         /// <inheritdoc/>
         public override Response UpdateStatus(CancellationToken cancellationToken = default)
         {
-            if (!_completed)
+            if (!HasCompleted)
             {
-                using DiagnosticScope scope = _pipeline.CreateScope($"{nameof(DeleteCertificateOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = _pipeline.GetResponse(RequestMethod.Get, cancellationToken, CertificateClient.DeletedCertificatesPath, _value.Name);
-                    _completed = CheckCompleted(_response);
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                return _operationInternal.UpdateStatus(cancellationToken);
             }
 
             return GetRawResponse();
@@ -85,22 +76,9 @@ public override Response UpdateStatus(CancellationToken cancellationToken = defa
         /// <inheritdoc/>
         public override async ValueTask<Response> UpdateStatusAsync(CancellationToken cancellationToken = default)
         {
-            if (!_completed)
+            if (!HasCompleted)
             {
-                using DiagnosticScope scope = _pipeline.CreateScope($"{nameof(DeleteCertificateOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, CertificateClient.DeletedCertificatesPath, _value.Name).ConfigureAwait(false);
-                    _completed = await CheckCompletedAsync(_response).ConfigureAwait(false);
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                return await _operationInternal.UpdateStatusAsync(cancellationToken).ConfigureAwait(false);
             }
 
             return GetRawResponse();
@@ -114,34 +92,27 @@ public override ValueTask<Response<DeletedCertificate>> WaitForCompletionAsync(C
         public override ValueTask<Response<DeletedCertificate>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken) =>
             this.DefaultWaitForCompletionAsync(pollingInterval, cancellationToken);
 
-        private async ValueTask<bool> CheckCompletedAsync(Response response)
+        async ValueTask<OperationState> IOperation.UpdateStateAsync(bool async, CancellationToken cancellationToken)
         {
-            switch (response.Status)
-            {
-                case 200:
-                case 403: // Access denied but proof the certificate was deleted.
-                    return true;
-
-                case 404:
-                    return false;
+            Response response = async
+                ? await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, CertificateClient.DeletedCertificatesPath, _value.Name).ConfigureAwait(false)
+                : _pipeline.GetResponse(RequestMethod.Get, cancellationToken, CertificateClient.DeletedCertificatesPath, _value.Name);
 
-                default:
-                    throw await _pipeline.Diagnostics.CreateRequestFailedExceptionAsync(response).ConfigureAwait(false);
-            }
-        }
-        private bool CheckCompleted(Response response)
-        {
             switch (response.Status)
             {
                 case 200:
                 case 403: // Access denied but proof the certificate was deleted.
-                    return true;
+                    return OperationState.Success(response);
 
                 case 404:
-                    return false;
+                    return OperationState.Pending(response);
 
                 default:
-                    throw _pipeline.Diagnostics.CreateRequestFailedException(response);
+                    RequestFailedException ex = async
+                        ? await _pipeline.Diagnostics.CreateRequestFailedExceptionAsync(response).ConfigureAwait(false)
+                        : _pipeline.Diagnostics.CreateRequestFailedException(response);
+
+                    return OperationState.Failure(response, ex);
             }
         }
     }