From 197426af0001952dd691a29b0876cf20c35e405b Mon Sep 17 00:00:00 2001
From: Heath Stewart <heaths@microsoft.com>
Date: Tue, 30 Nov 2021 17:21:40 -0800
Subject: [PATCH] Use operation helpers

Partially fixes #22912
---
 .../Azure.Security.KeyVault.Secrets.csproj    |  4 +-
 .../src/DeleteSecretOperation.cs              | 76 +++++-------------
 .../src/RecoverDeletedSecretOperation.cs      | 77 +++++--------------
 3 files changed, 45 insertions(+), 112 deletions(-)

diff --git a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/Azure.Security.KeyVault.Secrets.csproj b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/Azure.Security.KeyVault.Secrets.csproj
index 0fba7ca7bd074..f04b10d59a74c 100644
--- a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/Azure.Security.KeyVault.Secrets.csproj
+++ b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/Azure.Security.KeyVault.Secrets.csproj
@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <Description>This is the Microsoft Azure Key Vault Secrets client library</Description>
@@ -37,6 +37,8 @@
     <Compile Include="$(AzureCoreSharedSources)PageResponseEnumerator.cs" LinkBase="Shared" />
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScopeFactory.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)TaskExtensions.cs" LinkBase="Shared" />
   </ItemGroup>
 
diff --git a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/DeleteSecretOperation.cs b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/DeleteSecretOperation.cs
index 98e42d5e9ae2b..a457a74d9578f 100644
--- a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/DeleteSecretOperation.cs
+++ b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/DeleteSecretOperation.cs
@@ -5,27 +5,26 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
-using Azure.Core.Pipeline;
 
 namespace Azure.Security.KeyVault.Secrets
 {
     /// <summary>
     /// A long-running operation for <see cref="SecretClient.StartDeleteSecret(string, CancellationToken)"/> or <see cref="SecretClient.StartDeleteSecretAsync(string, CancellationToken)"/>.
     /// </summary>
-    public class DeleteSecretOperation : Operation<DeletedSecret>
+    public class DeleteSecretOperation : Operation<DeletedSecret>, IOperation
     {
         private static readonly TimeSpan s_defaultPollingInterval = TimeSpan.FromSeconds(2);
 
         private readonly KeyVaultPipeline _pipeline;
+        private readonly OperationInternal _operationInternal;
         private readonly DeletedSecret _value;
-        private Response _response;
-        private bool _completed;
+        private readonly bool _completed;
 
         internal DeleteSecretOperation(KeyVaultPipeline pipeline, Response<DeletedSecret> 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(DeleteSecretOperation));
 
             // The recoveryId is only returned if soft-delete is enabled.
             if (_value.RecoveryId is null)
@@ -50,33 +49,20 @@ protected DeleteSecretOperation() {}
         public override DeletedSecret Value => _value;
 
         /// <inheritdoc/>
-        public override bool HasCompleted => _completed;
+        public override bool HasCompleted => _completed || _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(DeleteSecretOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = _pipeline.GetResponse(RequestMethod.Get, cancellationToken, SecretClient.DeletedSecretsPath, _value.Name);
-                    _completed = CheckCompleted(_response);
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                return _operationInternal.UpdateStatus(cancellationToken);
             }
 
             return GetRawResponse();
@@ -85,22 +71,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(DeleteSecretOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, SecretClient.DeletedSecretsPath, _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 +87,27 @@ public override ValueTask<Response<DeletedSecret>> WaitForCompletionAsync(Cancel
         public override ValueTask<Response<DeletedSecret>> 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 secret was deleted.
-                    return true;
-
-                case 404:
-                    return false;
+            Response response = async
+                ? await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, SecretClient.DeletedSecretsPath, _value.Name).ConfigureAwait(false)
+                : _pipeline.GetResponse(RequestMethod.Get, cancellationToken, SecretClient.DeletedSecretsPath, _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 secret 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);
             }
         }
     }
diff --git a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/RecoverDeletedSecretOperation.cs b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/RecoverDeletedSecretOperation.cs
index 26d061a1ca6a6..798bf07ed1d8a 100644
--- a/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/RecoverDeletedSecretOperation.cs
+++ b/sdk/keyvault/Azure.Security.KeyVault.Secrets/src/RecoverDeletedSecretOperation.cs
@@ -5,27 +5,25 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
-using Azure.Core.Pipeline;
 
 namespace Azure.Security.KeyVault.Secrets
 {
     /// <summary>
     /// A long-running operation for <see cref="SecretClient.StartRecoverDeletedSecret(string, CancellationToken)"/> or <see cref="SecretClient.StartRecoverDeletedSecretAsync(string, CancellationToken)"/>.
     /// </summary>
-    public class RecoverDeletedSecretOperation : Operation<SecretProperties>
+    public class RecoverDeletedSecretOperation : Operation<SecretProperties>, IOperation
     {
         private static readonly TimeSpan s_defaultPollingInterval = TimeSpan.FromSeconds(2);
 
         private readonly KeyVaultPipeline _pipeline;
+        private readonly OperationInternal _operationInternal;
         private readonly SecretProperties _value;
-        private Response _response;
-        private bool _completed;
 
         internal RecoverDeletedSecretOperation(KeyVaultPipeline pipeline, Response<SecretProperties> 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(RecoverDeletedSecretOperation));
         }
 
         /// <summary> Initializes a new instance of <see cref="RecoverDeletedSecretOperation" /> for mocking. </summary>
@@ -44,33 +42,20 @@ protected RecoverDeletedSecretOperation() {}
         public override SecretProperties 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(RecoverDeletedSecretOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = _pipeline.GetResponse(RequestMethod.Get, cancellationToken, SecretClient.SecretsPath, _value.Name, "/", _value.Version);
-                    _completed = CheckCompleted(_response);
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                return _operationInternal.UpdateStatus(cancellationToken);
             }
 
             return GetRawResponse();
@@ -79,22 +64,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(RecoverDeletedSecretOperation)}.{nameof(UpdateStatus)}");
-                scope.AddAttribute("secret", _value.Name);
-                scope.Start();
-
-                try
-                {
-                    _response = await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, SecretClient.SecretsPath, _value.Name, "/", _value.Version).ConfigureAwait(false);
-                    _completed = await CheckCompletedAsync(_response).ConfigureAwait(false);
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                return await _operationInternal.UpdateStatusAsync(cancellationToken).ConfigureAwait(false);
             }
 
             return GetRawResponse();
@@ -108,34 +80,27 @@ public override ValueTask<Response<SecretProperties>> WaitForCompletionAsync(Can
         public override ValueTask<Response<SecretProperties>> 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 secret was recovered.
-                    return true;
-
-                case 404:
-                    return false;
+            Response response = async
+                ? await _pipeline.GetResponseAsync(RequestMethod.Get, cancellationToken, SecretClient.SecretsPath, _value.Name, "/", _value.Version).ConfigureAwait(false)
+                : _pipeline.GetResponse(RequestMethod.Get, cancellationToken, SecretClient.SecretsPath, _value.Name, "/", _value.Version);
 
-                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 secret was recovered.
-                    return true;
+                case 403: // Access denied but proof the secret was deleted.
+                    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);
             }
         }
     }