Migrate Restore-AzStorageBlobRange and ManagementPolicy related cmdlets

src/Storage/Storage.Management/Models/PSBlobRestore.cs

@@ -57,7 +57,7 @@ public PSBlobRestoreRange(Track2Models.BlobRestoreRange range)
                 foreach (PSBlobRestoreRange range in ranges)
                 {
                     re.Add(
-                        new Track2Models.BlobRestoreRange(range.EndRange, range.StartRange));
+                        new Track2Models.BlobRestoreRange(range.StartRange, range.EndRange));
                 }
             }
             return re;
@@ -112,6 +112,14 @@ public PSBlobRestoreStatus(Track2Models.BlobRestoreStatus status)
                 this.Parameters = status.Parameters is null ? null : new PSBlobRestoreParameters(status.Parameters);
             }
         }
+
+        public PSBlobRestoreStatus(string status, string failureReason, string restoreId, PSBlobRestoreParameters parameters)
+        {
+            Status = status;
+            FailureReason = failureReason;
+            RestoreId = restoreId;
+            Parameters = parameters;
+        }
     }
 
     /// <summary>

src/Storage/Storage.Management/Storage.Management.csproj

@@ -16,11 +16,10 @@
   <ItemGroup>
     <Compile Remove="StorageAccount\AddAzureStorageAccountNetworkRule.cs" />
     <Compile Remove="StorageAccount\GetAzureStorageAccountNetworkRuleSet.cs" />
-    <Compile Remove="StorageAccount\InvokeAzureStorageAccountFailover.cs" />
     <Compile Remove="StorageAccount\InvokeAzureStorageAccountHierarchicalNamespaceUpgrade.cs" />
     <Compile Remove="StorageAccount\RemoveAzureStorageAccountNetworkRule.cs" />
-    <Compile Remove="StorageAccount\RestoreAzStorageBlobRange.cs" />
     <Compile Remove="StorageAccount\UpdateAzureStorageAccountNetworkRuleSet.cs" />
+    <Compile Remove="StorageAccount\InvokeAzureStorageAccountFailover.cs" />
   </ItemGroup>
   <ItemGroup>
     <!-- <PackageReference Include="Azure.Core" Version="1.24.0" /> -->

src/Storage/Storage.Management/StorageAccount/AddAzureStorageAccountManagementPolicyAction.cs

@@ -14,8 +14,6 @@
 
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
 using System.Globalization;
 using System.Management.Automation;

src/Storage/Storage.Management/StorageAccount/GetAzureStorageAccountManagementPolicy.cs

@@ -12,11 +12,10 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Azure.ResourceManager.Storage;
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
 using System.Management.Automation;
 
@@ -95,11 +94,11 @@ public override void ExecuteCmdlet()
                     break;
             }
 
-            ManagementPolicy managementPolicy = this.StorageClient.ManagementPolicies.Get(
-                 this.ResourceGroupName,
-                 this.StorageAccountName);
+            ManagementPolicyResource policy = this.StorageClientTrack2.GetManagementPolicyResource(
+                this.ResourceGroupName, this.StorageAccountName, "default").Get();
 
-            WriteObject(new PSManagementPolicy(managementPolicy, this.ResourceGroupName, this.StorageAccountName), true);
+            var result = new PSManagementPolicy(policy, this.ResourceGroupName, this.StorageAccountName);
+            WriteObject(result, true);
         }
     }
 }

src/Storage/Storage.Management/StorageAccount/NewAzureStorageAccountManagementPolicyFilter.cs

@@ -14,8 +14,6 @@
 
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using System.Management.Automation;
 
 namespace Microsoft.Azure.Commands.Management.Storage

src/Storage/Storage.Management/StorageAccount/NewAzureStorageAccountManagementPolicyRule.cs

@@ -14,8 +14,6 @@
 
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
 using System.Management.Automation;
 

src/Storage/Storage.Management/StorageAccount/RemoveAzureStorageAccountManagementPolicy.cs

@@ -12,11 +12,10 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Azure;
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using System.Management.Automation;
 
 namespace Microsoft.Azure.Commands.Management.Storage
@@ -116,9 +115,8 @@ public override void ExecuteCmdlet()
                         break;
                 }
 
-                this.StorageClient.ManagementPolicies.Delete(
-                    this.ResourceGroupName,
-                    this.StorageAccountName);
+                this.StorageClientTrack2.GetManagementPolicyResource(this.ResourceGroupName, this.StorageAccountName, "default")
+                    .Delete(WaitUntil.Completed);
 
                 if (PassThru.IsPresent)
                 {

src/Storage/Storage.Management/StorageAccount/RestoreAzStorageBlobRange.cs

@@ -12,15 +12,18 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Azure;
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
 using Microsoft.Rest.Azure;
+using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using System;
+using System.Collections.Generic;
 using System.Management.Automation;
 using System.Threading.Tasks;
+using Track2 = Azure.ResourceManager.Storage;
+using Track2Models = Azure.ResourceManager.Storage.Models;
 
 namespace Microsoft.Azure.Commands.Management.Storage
 {
@@ -79,7 +82,7 @@ public class RestoreAzureStorageBlobRangeCommand : StorageAccountBaseCmdlet
 
         [Parameter(Mandatory = true, HelpMessage = "The Time to Restore Blob.")]
         [ValidateNotNull]
-        public DateTime TimeToRestore { get; set; }
+        public DateTimeOffset TimeToRestore { get; set; }
 
         [Parameter(Mandatory = false, HelpMessage = "The blob range to Restore.")]
         [ValidateNotNull]
@@ -118,50 +121,116 @@ public override void ExecuteCmdlet()
 
             if (ShouldProcess(this.StorageAccountName, "Restore Blob Range"))
             {
-                if (waitForComplete)
-                {                    
-                    Task<AzureOperationResponse<BlobRestoreStatus>> beginTask = this.StorageClient.StorageAccounts.BeginRestoreBlobRangesWithHttpMessagesAsync(
-                     this.ResourceGroupName,
-                     this.StorageAccountName,
-                     this.TimeToRestore,
-                     PSBlobRestoreRange.ParseBlobRestoreRanges(this.BlobRestoreRange));
-
-                    beginTask.Wait();
-
-                    AzureOperationResponse<BlobRestoreStatus> response = beginTask.Result;
+                Track2.StorageAccountResource account = this.StorageClientTrack2.GetStorageAccount(this.ResourceGroupName, this.StorageAccountName);
+                var restoreLro = account.RestoreBlobRanges(
+                    WaitUntil.Started,
+                    new Track2Models.BlobRestoreContent(
+                        this.TimeToRestore.ToUniversalTime(),
+                        PSBlobRestoreRange.ParseBlobRestoreRanges(this.BlobRestoreRange)));
+
+                // This is a temporary workaround of SDK issue https://github.com/Azure/azure-sdk-for-net/issues/29060
+                // The workaround is to get the raw response and parse it into the output desired
+                // The Blob restore status should be got from SDK directly once the issue is fixed
+                Dictionary<string, object> temp = restoreLro.GetRawResponse().Content.ToObjectFromJson() as Dictionary<string, object>;
 
-                    WriteWarning(string.Format("Restore blob ranges with Id '{0}' started. Restore blob ranges time to complete is dependent on the size of the restore.", response.Body is null ? "" : response.Body.RestoreId));
+                if (waitForComplete)
+                {
+                    if (temp == null)
+                    {
+                        throw new InvalidJobStateException("Could not fetch the Blob restore response.");
+                    }
+                    PSBlobRestoreStatus blobRestoreStatus = ParseRestoreRawResponse(temp);
+                    if (blobRestoreStatus.RestoreId != null)
+                    {
+                        WriteWarning(string.Format("Restore blob ranges with Id '{0}' started. Restore blob ranges time to complete is dependent on the size of the restore.", blobRestoreStatus.RestoreId));
+                    }
+                    else
+                    {
+                        WriteWarning(string.Format("Could not fetch the Restore Id."));
+                    }
 
-                    Task<AzureOperationResponse<BlobRestoreStatus>> waitTask = ((StorageManagementClient)this.StorageClient).GetPostOrDeleteOperationResultAsync(response, null, new System.Threading.CancellationToken());
                     try
                     {
-                        waitTask.Wait();
+                        var result = restoreLro.WaitForCompletion().Value;
+                        WriteObject(new PSBlobRestoreStatus(result));
                     }
                     catch (System.AggregateException ex) when (ex.InnerException is CloudException)
                     {
-                        throw new InvalidJobStateException(string.Format("Blob ranges restore failed with information: '{0}'.", ((CloudException)ex.InnerException).Response.Content));
+                        throw new InvalidJobStateException(string.Format("Blob ranges restore failed with information: '{0}'.", ex.ToString()));
                     }
-
-                    AzureOperationResponse<BlobRestoreStatus> result = waitTask.Result;
-
-                    WriteObject(new PSBlobRestoreStatus(result.Body));
-
                 }
                 else
                 {
-                    BlobRestoreStatus status = this.StorageClient.StorageAccounts.BeginRestoreBlobRanges(
-                     this.ResourceGroupName,
-                     this.StorageAccountName,
-                     this.TimeToRestore,
-                     PSBlobRestoreRange.ParseBlobRestoreRanges(this.BlobRestoreRange));
-
-                    WriteObject(new PSBlobRestoreStatus(status));
-                    if (status != null && status.Status == BlobRestoreProgressStatus.Failed)
+                    if (temp == null)
+                    {
+                        throw new InvalidJobStateException("Could not fetch the Blob restore response.");
+                    }
+
+                    PSBlobRestoreStatus blobRestoreStatus = ParseRestoreRawResponse(temp);
+
+                    if (blobRestoreStatus.Status != null)
+                    {
+                        if (blobRestoreStatus.Status == Track2Models.BlobRestoreProgressStatus.Failed.ToString())
+                        {
+                            throw new InvalidJobStateException("Blob ranges restore failed.");
+                        }
+                    } 
+                    else
                     {
-                        throw new InvalidJobStateException("Blob ranges restore failed.");
+                        WriteWarning(string.Format("Could not fetch the status."));
                     }
+                    WriteObject(blobRestoreStatus);
                 }
             }
         }
+
+        private PSBlobRestoreStatus ParseRestoreRawResponse(Dictionary<string, object> response)
+        {
+            response.TryGetValue("restoreId", out object restoreId);
+            response.TryGetValue("status", out object jobStatus);
+            response.TryGetValue("parameters", out object parameters);
+
+            PSBlobRestoreParameters blobRestoreParameters;
+            Dictionary<string, object> paramMap = parameters as Dictionary<string, object>;
+
+            if (paramMap == null)
+            {
+                blobRestoreParameters = null;
+            }
+            else
+            {
+                blobRestoreParameters = new PSBlobRestoreParameters();
+                paramMap.TryGetValue("timetoRestore", out object timeToRestore);
+                DateTimeOffset.TryParse(timeToRestore?.ToString(), out DateTimeOffset parseDate);
+                blobRestoreParameters.TimeToRestore = parseDate;
+
+                paramMap.TryGetValue("blobRanges", out object ranges);
+                List<PSBlobRestoreRange> blobRestoreRanges = new List<PSBlobRestoreRange>();
+
+                object[] rangesList = ranges as object[];
+                foreach (object range in rangesList)
+                {
+                    Dictionary<string, object> rangeMap = range as Dictionary<string, object>;
+
+                    rangeMap.TryGetValue("startRange", out object startRange);
+                    rangeMap.TryGetValue("endRange", out object endRange);
+
+                    PSBlobRestoreRange blobRestoreRange = new PSBlobRestoreRange
+                    {
+                        StartRange = startRange?.ToString(),
+                        EndRange = endRange?.ToString()
+                    };
+
+                    blobRestoreRanges.Add(blobRestoreRange);
+                }
+                blobRestoreParameters.BlobRanges = blobRestoreRanges.ToArray();
+            }
+
+            return new PSBlobRestoreStatus(
+                status: jobStatus?.ToString(),
+                failureReason: null,
+                restoreId: restoreId?.ToString(),
+                parameters: blobRestoreParameters);
+        }
     }
 }

src/Storage/Storage.Management/StorageAccount/SetAzureStorageAccountManagementPolicy.cs

@@ -12,16 +12,16 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Azure.ResourceManager.Storage;
 using Microsoft.Azure.Commands.Management.Storage.Models;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
-using Microsoft.Azure.Management.Storage;
-using Microsoft.Azure.Management.Storage.Models;
-using System.Management.Automation;
-using Newtonsoft.Json.Linq;
-using System.Globalization;
-using System.Collections.Generic;
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
 using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
+using Newtonsoft.Json.Linq;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Management.Automation;
+using Track2 = Azure.ResourceManager.Storage;
 
 namespace Microsoft.Azure.Commands.Management.Storage
 {
@@ -62,7 +62,7 @@ public class SetAzureStorageAccountManagementPolicyCommand : StorageAccountBaseC
             Position = 0,
             Mandatory = true,
             HelpMessage = "Resource Group Name.",
-           ParameterSetName = AccountNamePolicyRuleParameterSet)]       
+           ParameterSetName = AccountNamePolicyRuleParameterSet)]
         [Parameter(
             Position = 0,
             Mandatory = true,
@@ -155,7 +155,7 @@ public override void ExecuteCmdlet()
             base.ExecuteCmdlet();
             if (ShouldProcess(this.StorageAccountName, "Set Storage Account Management Policy"))
             {
-                if ((this.ParameterSetName == AccountObjectPolicyRuleParameterSet) 
+                if ((this.ParameterSetName == AccountObjectPolicyRuleParameterSet)
                     || (this.ParameterSetName == AccountObjectPolicyObjectParameterSet))
                 {
                     this.ResourceGroupName = StorageAccount.ResourceGroupName;
@@ -168,35 +168,43 @@ public override void ExecuteCmdlet()
                     this.ResourceGroupName = accountResource.ResourceGroupName;
                     this.StorageAccountName = accountResource.ResourceName;
                 }
-                ManagementPolicy managementPolicy;
+
+                ManagementPolicyResource managementPolicyResource;
+                Track2.ManagementPolicyData data;
 
                 switch (this.ParameterSetName)
                 {
                     case AccountObjectPolicyRuleParameterSet:
                     case AccountNamePolicyRuleParameterSet:
                     case AccountResourceIdPolicyRuleParameterSet:
-                        managementPolicy = this.StorageClient.ManagementPolicies.CreateOrUpdate(
-                            this.ResourceGroupName,
-                            this.StorageAccountName,
-                            new ManagementPolicySchema(
-                                //this.version,
-                                PSManagementPolicyRule.ParseManagementPolicyRules(this.Rule)));
+
+                        data = new Track2.ManagementPolicyData
+                        {
+                            Rules = PSManagementPolicyRule.ParseManagementPolicyRules(this.Rule)
+                        };
+
+                        managementPolicyResource = this.StorageClientTrack2
+                            .GetManagementPolicyResource(this.ResourceGroupName, this.StorageAccountName, "default")
+                            .CreateOrUpdate(global::Azure.WaitUntil.Completed, data).Value;
                         break;
                     case AccountObjectPolicyObjectParameterSet:
                     case AccountNamePolicyObjectParameterSet:
                     case AccountResourceIdPolicyObjectParameterSet:
-                        managementPolicy = this.StorageClient.ManagementPolicies.CreateOrUpdate(
-                            this.ResourceGroupName,
-                            this.StorageAccountName,
-                            new ManagementPolicySchema(
-                                //this.Policy.Version,
-                                PSManagementPolicyRule.ParseManagementPolicyRules(this.Policy.Rules)));
+
+                        data = new Track2.ManagementPolicyData()
+                        {
+                            Rules = PSManagementPolicyRule.ParseManagementPolicyRules(this.Policy.Rules)
+                        };
+
+                        managementPolicyResource = this.StorageClientTrack2
+                            .GetManagementPolicyResource(this.ResourceGroupName, this.StorageAccountName, "default")
+                            .CreateOrUpdate(global::Azure.WaitUntil.Completed, data).Value;
                         break;
                     default:
                         throw new PSArgumentException(string.Format(CultureInfo.InvariantCulture, "Invalid ParameterSet: {0}", this.ParameterSetName));
                 }
 
-                WriteObject(new PSManagementPolicy(managementPolicy, this.ResourceGroupName, this.StorageAccountName), true);
+                WriteObject(new PSManagementPolicy(managementPolicyResource, this.ResourceGroupName, this.StorageAccountName), true);
             }
         }
     }