Adding UseOriginalStorageAccountForDiskRestore parameter to Restore cmdlet.

Author: Samuel Anudeep

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.Helpers/Conversions/RecoveryPointConversions.cs

@@ -83,6 +83,8 @@ public static List<RecoveryPointBase> GetPSAzureRecoveryPoints(
                         EncryptionEnabled = recPoint.IsSourceVMEncrypted.HasValue ?
                             recPoint.IsSourceVMEncrypted.Value : false,
                         IlrSessionActive = isInstantILRSessionActive,
+                        OriginalSAEnabled = recPoint.OriginalStorageAccountOption.HasValue ?
+                            recPoint.OriginalStorageAccountOption.Value : false,
                     };
                     result.Add(rpBase);
                 }
@@ -180,6 +182,8 @@ public static RecoveryPointBase GetPSAzureRecoveryPoints(
                     IlrSessionActive = isInstantILRSessionActive,
                     SourceResourceId = item.SourceResourceId,
                     SourceVMStorageType = recPoint.SourceVMStorageType,
+                    OriginalSAEnabled = recPoint.OriginalStorageAccountOption.HasValue ?
+                        recPoint.OriginalStorageAccountOption.Value : false,
                 };
 
                 if (vmResult.EncryptionEnabled && recPoint.KeyAndSecret != null)

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.Models/AzureVmModels/AzureVmRecoveryPoint.cs

@@ -64,6 +64,12 @@ public class AzureVmRecoveryPoint : RecoveryPointBase
         /// </summary>
         public KeyAndSecretDetails KeyAndSecretDetails { get; set; }
 
+        /// <summary>
+        /// Identifies whether this recovery point supports restoring
+        /// disks to their original storage accounts.
+        /// </summary>
+        public bool OriginalSAEnabled { get; set; }
+
         public AzureVmRecoveryPoint()
         {
 

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.Models/CmdletParamEnums.cs

@@ -46,6 +46,7 @@ public enum RestoreBackupItemParams
         StorageAccountId,
         StorageAccountLocation,
         StorageAccountType,
+        OsaOption
     }
 
     public enum PolicyParams

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.Providers/Providers/IaasVmPsBackupProvider.cs

@@ -239,8 +239,9 @@ public RestAzureNS.AzureOperationResponse TriggerRestore()
                 ProviderData[RestoreBackupItemParams.StorageAccountLocation].ToString();
             string storageAccountType =
                 ProviderData[RestoreBackupItemParams.StorageAccountType].ToString();
+            bool osaOption = (bool)ProviderData[RestoreBackupItemParams.OsaOption];
 
-            var response = ServiceClientAdapter.RestoreDisk(rp, storageAccountId, storageAccountLocation, storageAccountType);
+            var response = ServiceClientAdapter.RestoreDisk(rp, storageAccountId, storageAccountLocation, storageAccountType, osaOption);
             return response;
         }
 

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.ServiceClientAdapter/BMSAPIs/RestoreDiskAPIs.cs

@@ -36,8 +36,11 @@ public RestAzureNS.AzureOperationResponse RestoreDisk(
             AzureVmRecoveryPoint rp,
             string storageAccountId,
             string storageAccountLocation,
-            string storageAccountType)
+            string storageAccountType,
+            bool osaOption)
         {
+            var useOsa = ShouldUseOsa(rp, osaOption);
+
             string resourceGroupName = BmsAdapter.GetResourceGroupName();
             string resourceName = BmsAdapter.GetResourceName();
             string vaultLocation = BmsAdapter.GetResourceLocation();
@@ -68,6 +71,7 @@ public RestAzureNS.AzureOperationResponse RestoreDisk(
                 Region = vaultLocation,
                 StorageAccountId = storageAccountId,
                 SourceResourceId = rp.SourceResourceId,
+                OriginalStorageAccountOption = useOsa,
             };
 
             RestoreRequestResource triggerRestoreRequest = new RestoreRequestResource();
@@ -85,5 +89,23 @@ public RestAzureNS.AzureOperationResponse RestoreDisk(
 
             return response;
         }
+
+        private bool ShouldUseOsa(AzureVmRecoveryPoint rp, bool osaOption)
+        {
+            bool useOsa = false;
+            if (osaOption)
+            {
+                if (rp.OriginalSAEnabled)
+                {
+                    useOsa = true;
+                }
+                else
+                {
+                    throw new Exception("This recovery point doesn’t have the capability to restore disks to their original storage account. Re-run the restore command without the UseOriginalStorageAccountForDisks parameter.");
+                }
+            }
+
+            return useOsa;
+        }
     }
 }

src/ResourceManager/RecoveryServices.Backup/Commands.RecoveryServices.Backup.Test/ScenarioTests/IaasVm/ItemTests.ps1

@@ -213,6 +213,13 @@ function Test-AzureVMFullRestore
 
 		Set-AzureRmRecoveryServicesVaultContext -Vault $vault;
 
+		Assert-ThrowsContains { Restore-AzureRmRecoveryServicesBackupItem `
+			-RecoveryPoint $rp `
+			-StorageAccountName $saName `
+			-StorageAccountResourceGroupName $resourceGroupName `
+			-UseOriginalStorageAccountForDiskRestore } `
+			"This recovery point doesn’t have the capability to restore disks to their original storage account. Re-run the restore command without the UseOriginalStorageAccountForDisks parameter.";
+
 		$restoreJob = Restore-AzureRmRecoveryServicesBackupItem `
 			-RecoveryPoint $rp `
 			-StorageAccountName $saName `