Support Lro interim response

samples/Azure.Management.Storage/Generated/LongRunningOperation/StorageAccountRestoreBlobRangesOperation.cs

@@ -0,0 +1,101 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+// <auto-generated/>
+
+#nullable disable
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure;
+using Azure.Core;
+using Azure.Core.Pipeline;
+using Azure.Management.Storage.Models;
+using Azure.ResourceManager;
+
+namespace Azure.Management.Storage
+{
+    /// <summary> A class representing the specific long-running operation StorageAccountRestoreBlobRangesOperation. </summary>
+    public class StorageAccountRestoreBlobRangesOperation : ArmOperation<BlobRestoreStatus>
+    {
+        private readonly StorageArmOperation<BlobRestoreStatus> _operation;
+
+        private readonly IOperationSource<BlobRestoreStatus> _operationSource;
+
+        private readonly AsyncLockWithValue<BlobRestoreStatus> _stateLock;
+
+        private readonly Response _interimResponse;
+
+        /// <summary> Initializes a new instance of StorageAccountRestoreBlobRangesOperation for mocking. </summary>
+        protected StorageAccountRestoreBlobRangesOperation()
+        {
+        }
+
+        internal StorageAccountRestoreBlobRangesOperation(IOperationSource<BlobRestoreStatus> source, ClientDiagnostics clientDiagnostics, HttpPipeline pipeline, Request request, Response response, OperationFinalStateVia finalStateVia)
+        {
+            _operation = new StorageArmOperation<BlobRestoreStatus>(source, clientDiagnostics, pipeline, request, response, finalStateVia);
+            _operationSource = source;
+            _stateLock = new AsyncLockWithValue<BlobRestoreStatus>();
+            _interimResponse = response;
+        }
+
+        /// <inheritdoc />
+#pragma warning disable CA1822
+        public override string Id => throw new NotImplementedException();
+#pragma warning restore CA1822
+
+        /// <inheritdoc />
+        public override BlobRestoreStatus Value => _operation.Value;
+
+        /// <inheritdoc />
+        public override bool HasValue => _operation.HasValue;
+
+        /// <inheritdoc />
+        public override bool HasCompleted => _operation.HasCompleted;
+
+        /// <inheritdoc />
+        public override Response GetRawResponse() => _operation.GetRawResponse();
+
+        /// <inheritdoc />
+        public override Response UpdateStatus(CancellationToken cancellationToken = default) => _operation.UpdateStatus(cancellationToken);
+
+        /// <inheritdoc />
+        public override ValueTask<Response> UpdateStatusAsync(CancellationToken cancellationToken = default) => _operation.UpdateStatusAsync(cancellationToken);
+
+        /// <inheritdoc />
+        public override Response<BlobRestoreStatus> WaitForCompletion(CancellationToken cancellationToken = default) => _operation.WaitForCompletion(cancellationToken);
+
+        /// <inheritdoc />
+        public override Response<BlobRestoreStatus> WaitForCompletion(TimeSpan pollingInterval, CancellationToken cancellationToken = default) => _operation.WaitForCompletion(pollingInterval, cancellationToken);
+
+        /// <inheritdoc />
+        public override ValueTask<Response<BlobRestoreStatus>> WaitForCompletionAsync(CancellationToken cancellationToken = default) => _operation.WaitForCompletionAsync(cancellationToken);
+
+        /// <inheritdoc />
+        public override ValueTask<Response<BlobRestoreStatus>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken = default) => _operation.WaitForCompletionAsync(pollingInterval, cancellationToken);
+
+        /// <summary> Gets interim status of the long-running operation. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        /// <returns> The interim status of the long-running operation. </returns>
+        public virtual async ValueTask<BlobRestoreStatus> GetCurrentStatusAsync(CancellationToken cancellationToken = default) => await GetCurrentState(true, cancellationToken).ConfigureAwait(false);
+
+        /// <summary> Gets interim status of the long-running operation. </summary>
+        /// <param name="cancellationToken"> The cancellation token to use. </param>
+        /// <returns> The interim status of the long-running operation. </returns>
+        public virtual BlobRestoreStatus GetCurrentStatus(CancellationToken cancellationToken = default) => GetCurrentState(false, cancellationToken).EnsureCompleted();
+
+        private async ValueTask<BlobRestoreStatus> GetCurrentState(bool async, CancellationToken cancellationToken)
+        {
+            using var asyncLock = await _stateLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
+            if (asyncLock.HasValue)
+            {
+                return asyncLock.Value;
+            }
+            var val = async ? await _operationSource.CreateResultAsync(_interimResponse, cancellationToken).ConfigureAwait(false)
+                    : _operationSource.CreateResult(_interimResponse, cancellationToken);
+            asyncLock.SetValue(val);
+            return val;
+        }
+    }
+}

samples/Azure.Management.Storage/readme.md

@@ -57,4 +57,7 @@ directive:
   - from: swagger-document
     where: $.paths["/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Storage/storageAccounts/{accountName}/fileServices/default/shares"].get.parameters[4].type
     transform: return "integer"
+  - from: swagger-document
+    where: $.paths["/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Storage/storageAccounts/{accountName}/restoreBlobRanges"].post
+    transform: $["x-ms-long-running-operation-options"]["enable-interim-state"] = true
 ```

src/AutoRest.CSharp/Common/Input/CodeModelPartials.cs

@@ -35,6 +35,20 @@ public OperationFinalStateVia LongRunningFinalStateVia
             }
         }
 
+        // This is a new extension introduced by generator to control whether interim state returns are supported in lro.
+        public bool IsInterimLongRunningStateEnabled
+        {
+            get
+            {
+                var isInterimStatusEnabled = Extensions.GetValue<IDictionary<object, object>>("x-ms-long-running-operation-options")?.GetValue<string>("enable-interim-state");
+                return isInterimStatusEnabled switch
+                {
+                    "true" => true,
+                    _ => false,
+                };
+            }
+        }
+
         public string? Accessibility => Extensions.GetValue<string>("x-accessibility");
 
         public ServiceResponse LongRunningInitialResponse

src/AutoRest.CSharp/Mgmt/AutoRest/MgmtOutputLibrary.cs

@@ -130,6 +130,8 @@ private static void ApplyGlobalConfigurations()
         public Dictionary<CSharpType, OperationSource> CSharpTypeToOperationSource { get; } = new Dictionary<CSharpType, OperationSource>();
         public IEnumerable<OperationSource> OperationSources => CSharpTypeToOperationSource.Values;
 
+        public ICollection<LongRunningInterimOperation> InterimOperations { get; } = new List<LongRunningInterimOperation>();
+
         private IEnumerable<Schema> UpdateBodyParameters()
         {
             Dictionary<Schema, int> usageCounts = new Dictionary<Schema, int>();

src/AutoRest.CSharp/Mgmt/AutoRest/MgmtTarget.cs

@@ -146,6 +146,13 @@ public static async Task ExecuteAsync(GeneratedCodeWorkspace project, CodeModel
             lroWriter.Write();
             AddGeneratedFile(project, lroWriter.Filename, lroWriter.ToString());
 
+            foreach (var interimOperation in MgmtContext.Library.InterimOperations.Distinct(LongRunningInterimOperation.LongRunningInterimOperationComparer))
+            {
+                var writer = new MgmtLongRunningInterimOperationWriter(interimOperation);
+                writer.Write();
+                AddGeneratedFile(project, $"LongRunningOperation/{interimOperation.TypeName}.cs", writer.ToString());
+            }
+
             foreach (var operationSource in MgmtContext.Library.OperationSources)
             {
                 var writer = new OperationSourceWriter(operationSource);

src/AutoRest.CSharp/Mgmt/Generation/MgmtClientBaseWriter.cs

@@ -835,10 +835,17 @@ protected virtual void WriteLROMethodBranch(MgmtRestOperation operation, IEnumer
 
         protected virtual void WriteLROResponse(string diagnosticsVariableName, string pipelineVariableName, MgmtRestOperation operation, IEnumerable<ParameterMapping> parameterMapping, bool async)
         {
-            _writer.Append($"var operation = new {LibraryArmOperation}");
-            if (operation.ReturnType.IsGenericType)
+            if (operation.InterimOperation is not null)
             {
-                _writer.Append($"<{operation.MgmtReturnType}>");
+                _writer.Append($"var operation = new {operation.InterimOperation.TypeName}");
+            }
+            else
+            {
+                _writer.Append($"var operation = new {LibraryArmOperation}");
+                if (operation.ReturnType.IsGenericType)
+                {
+                    _writer.Append($"<{operation.MgmtReturnType}>");
+                }
             }
             _writer.Append($"(");
             if (operation.IsFakeLongRunningOperation)