New-AzureRmVM/VMSS: DataDiskSizeInGb parameter

src/ResourceManager/Common/Commands.Common.Strategies/Extensions.cs

@@ -21,6 +21,9 @@ namespace Microsoft.Azure.Commands.Common.Strategies
 {
     public static class Extensions
     {
+        public static IEnumerable<T> EmptyIfNull<T>(this IEnumerable<T> v)
+            => v ?? Enumerable.Empty<T>();
+
         public static TValue GetOrNull<TKey, TValue>(
             this IDictionary<TKey, TValue> dictionary, TKey key)
             where TValue : class

src/ResourceManager/Compute/Commands.Compute/Commands.Compute.csproj

@@ -270,7 +270,9 @@
     <Compile Include="Strategies\Client.cs" />
     <Compile Include="Strategies\ComputeRp\AvailabilitySetStrategy.cs" />
     <Compile Include="Strategies\ComputeRp\ComputeStrategy.cs" />
+    <Compile Include="Strategies\ComputeRp\DataDiskStrategy.cs" />
     <Compile Include="Strategies\ComputeRp\ImageAndOsType.cs" />
+    <Compile Include="Strategies\ComputeRp\ImageDataDiskStrategy.cs" />
     <Compile Include="Strategies\ComputeRp\ImageEx.cs" />
     <Compile Include="Strategies\ComputeRp\Images.cs" />
     <Compile Include="Strategies\ComputeRp\ApiEntityReferenceStrategy.cs" />

src/ResourceManager/Compute/Commands.Compute/Manual/VirtualMachineScaleSetCreateOrUpdateMethod.cs

@@ -114,6 +114,9 @@ public partial class NewAzureRmVmss : ComputeAutomationBaseCmdlet
         [Parameter(ParameterSetName = SimpleParameterSet, Mandatory = false)]
         public int[] NatBackendPort { get; set; }
 
+        [Parameter(ParameterSetName = SimpleParameterSet, Mandatory = false)]
+        public int[] DataDiskSizeInGb { get; set; }
+
         const int FirstPortRangeStart = 50000;
 
         sealed class Parameters : IParameters<VirtualMachineScaleSet>
@@ -218,7 +221,8 @@ public async Task<ResourceConfig<VirtualMachineScaleSet>> CreateConfigAsync()
                     instanceCount: _cmdlet.InstanceCount,
                     upgradeMode: _cmdlet.MyInvocation.BoundParameters.ContainsKey(nameof(UpgradePolicyMode))
                         ? _cmdlet.UpgradePolicyMode
-                        : (UpgradeMode?)null);
+                        : (UpgradeMode?)null,
+                    dataDisks: _cmdlet.DataDiskSizeInGb);
             }
         }
 

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/ComputeStrategy.cs

@@ -43,11 +43,13 @@ public static ResourceStrategy<TModel> Create<TModel, TOperations>(
 
         public static string GetConnectionString(
             this ImageAndOsType imageAndOsType, string fqdn, string user, string port = null)
-        {
-            var url = fqdn + (port == null ? string.Empty : ":" + port);
-            return imageAndOsType.OsType == OperatingSystemTypes.Windows
-                ? "mstsc /v:" + url
-                : "ssh " + (user == null ? string.Empty : user + "@") + url;
-        }
+            => imageAndOsType.OsType == OperatingSystemTypes.Windows
+                ? "mstsc /v:" 
+                    + fqdn 
+                    + (port == null ? string.Empty : ":" + port)
+                : "ssh " 
+                    + (user == null ? string.Empty : user + "@") 
+                    + fqdn 
+                    + (port == null ? string.Empty : " -p " + port);
     }
 }

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/DataDiskStrategy.cs

@@ -0,0 +1,60 @@
+using Microsoft.Azure.Commands.Common.Strategies;
+using Microsoft.Azure.Management.Compute.Models;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Microsoft.Azure.Commands.Compute.Strategies.ComputeRp
+{
+    static class DataDiskStrategy
+    {
+        static IList<T> CreateDataDisks<T>(
+            IEnumerable<int> imageDataDiskLuns,
+            IEnumerable<int> dataDiskSizes,
+            Func<DiskCreateOptionTypes, int, int?, T> createDataDisk)
+        {
+            if (dataDiskSizes == null)
+            {
+                return null;
+            }
+            imageDataDiskLuns = imageDataDiskLuns.EmptyIfNull();
+            var firstLun = imageDataDiskLuns
+                .Select(v => v + 1)
+                .Concat(new[] { 0 })
+                .Max();
+            return imageDataDiskLuns
+                .Select(lun => createDataDisk(DiskCreateOptionTypes.FromImage, lun, null))
+                .Concat(dataDiskSizes.Select((size, i) => createDataDisk(
+                    DiskCreateOptionTypes.Empty,
+                    i + firstLun,
+                    size)))
+                .ToList();
+        }
+
+        public static IList<DataDisk> CreateDataDisks(
+            IEnumerable<int> imageDataDiskLuns,
+            IEnumerable<int> dataDiskSizes)
+            => CreateDataDisks(
+                imageDataDiskLuns,
+                dataDiskSizes,
+                (createOption, lun, size) => new DataDisk
+                {
+                    CreateOption = createOption,
+                    Lun = lun,
+                    DiskSizeGB = size,
+                });
+
+        public static IList<VirtualMachineScaleSetDataDisk> CreateVmssDataDisks(
+            IEnumerable<int> dataDisks,
+            IEnumerable<int> dataDiskSizes)
+            => CreateDataDisks(
+                dataDisks,
+                dataDiskSizes,
+                (createOption, lun, size) => new VirtualMachineScaleSetDataDisk
+                {
+                    CreateOption = createOption,
+                    Lun = lun,
+                    DiskSizeGB = size,
+                });
+    }
+}

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/ImageAndOsType.cs

@@ -13,6 +13,7 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.Management.Compute.Models;
+using System.Collections.Generic;
 
 namespace Microsoft.Azure.Commands.Compute.Strategies.ComputeRp
 {
@@ -22,10 +23,13 @@ sealed class ImageAndOsType
 
         public ImageReference Image { get; }
 
-        public ImageAndOsType(OperatingSystemTypes osType, ImageReference image)
+        public IList<int> DataDiskLuns { get; }
+
+        public ImageAndOsType(OperatingSystemTypes osType, ImageReference image, IList<int> dataDiskLuns)
         {
             OsType = osType;
             Image = image;
+            DataDiskLuns = dataDiskLuns;
         }
     }
 }

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/ImageDataDiskStrategy.cs

@@ -0,0 +1,19 @@
+using Microsoft.Azure.Management.Compute.Models;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Microsoft.Azure.Commands.Compute.Strategies.ComputeRp
+{
+    static class ImageDataDiskStrategy
+    {
+        static IList<int> GetLuns<T>(this IEnumerable<T> imageDataDisks, Func<T, int?> getLun)
+            => imageDataDisks.Select(idd => getLun(idd) ?? 0).ToList();
+
+        public static IList<int> GetLuns(this IEnumerable<ImageDataDisk> imageDataDisks)
+            => imageDataDisks.GetLuns(idd => idd.Lun);
+
+        public static IList<int> GetLuns(this IEnumerable<DataDiskImage> imageDataDisks)
+            => imageDataDisks.GetLuns(idd => idd.Lun);
+    }
+}

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/ImageEx.cs

@@ -82,8 +82,11 @@ public static async Task<ImageAndOsType> UpdateImageAndOsTypeAsync(
                 }
                 var imageModel = await compute.VirtualMachineImages.GetAsync(
                     location, image.Publisher, image.Offer, image.Sku, image.Version);
-                return new ImageAndOsType(imageModel.OsDiskImage.OperatingSystem, image);
-            }
+                return new ImageAndOsType(
+                    imageModel.OsDiskImage.OperatingSystem,
+                    image,
+                    imageModel.DataDiskImages.GetLuns());
+            } 
             else if (imageName.Contains("/"))
             {
                 var resourceId = ResourceId.TryParse(imageName);
@@ -99,22 +102,13 @@ public static async Task<ImageAndOsType> UpdateImageAndOsTypeAsync(
                     throw new ArgumentException(Resources.ComputeMismatchSubscription);
                 }
 
-                var localImage = await compute.Images.GetAsync(
-                    resourceGroupName: resourceId.ResourceGroupName,
-                    imageName: resourceId.Name);
-
-                return new ImageAndOsType(
-                    localImage.StorageProfile.OsDisk.OsType,
-                    new ImageReference { Id = localImage.Id });
+                return await compute.GetImageAndOsTypeAsync(resourceId.ResourceGroupName, resourceId.Name);
             }
             else
             {
                 try
                 {
-                    var localImage = await compute.Images.GetAsync(resourceGroupName, imageName);
-                    return new ImageAndOsType(
-                        localImage.StorageProfile.OsDisk.OsType,
-                        new ImageReference { Id = localImage.Id });
+                    return await compute.GetImageAndOsTypeAsync(resourceGroupName, imageName);
                 }
                 catch
                 {
@@ -130,7 +124,8 @@ public static async Task<ImageAndOsType> UpdateImageAndOsTypeAsync(
                             osAndMap.Key == "Windows" 
                                 ? OperatingSystemTypes.Windows
                                 : OperatingSystemTypes.Linux,
-                            nameAndImage.Value)))
+                            nameAndImage.Value,
+                            null)))
                     .FirstOrDefault();
 
                 if (result == null)
@@ -141,5 +136,15 @@ public static async Task<ImageAndOsType> UpdateImageAndOsTypeAsync(
                 return result;
             }
         }
+
+        static async Task<ImageAndOsType> GetImageAndOsTypeAsync(
+            this ComputeManagementClient compute, string resourceGroupName, string name)
+        {
+            var localImage = await compute.Images.GetAsync(resourceGroupName, name);
+            return new ImageAndOsType(
+                localImage.StorageProfile.OsDisk.OsType,
+                new ImageReference { Id = localImage.Id },
+                localImage.StorageProfile.DataDisks.GetLuns());
+        }
     }
 }

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/VirtualMachineScaleSetStrategy.cs

@@ -47,7 +47,8 @@ internal static ResourceConfig<VirtualMachineScaleSet> CreateVirtualMachineScale
             string adminPassword,
             string vmSize,
             int instanceCount,
-            UpgradeMode? upgradeMode)
+            UpgradeMode? upgradeMode,
+            IEnumerable<int> dataDisks)
             => Strategy.CreateResourceConfig(
                 resourceGroup: resourceGroup,
                 name: name,
@@ -82,7 +83,9 @@ internal static ResourceConfig<VirtualMachineScaleSet> CreateVirtualMachineScale
                         },
                         StorageProfile = new VirtualMachineScaleSetStorageProfile
                         {
-                            ImageReference = imageAndOsType?.Image
+                            ImageReference = imageAndOsType?.Image,
+                            DataDisks = DataDiskStrategy.CreateVmssDataDisks(
+                                imageAndOsType?.DataDiskLuns, dataDisks)
                         },
                         NetworkProfile = new VirtualMachineScaleSetNetworkProfile
                         {

src/ResourceManager/Compute/Commands.Compute/Strategies/ComputeRp/VirtualMachineStrategy.cs

@@ -17,6 +17,7 @@
 using Microsoft.Azure.Management.Internal.Resources.Models;
 using Microsoft.Azure.Management.Internal.Network.Version2017_10_01.Models;
 using Microsoft.Azure.Commands.Common.Strategies;
+using System.Collections.Generic;
 
 namespace Microsoft.Azure.Commands.Compute.Strategies.ComputeRp
 {
@@ -43,7 +44,8 @@ public static ResourceConfig<VirtualMachine> CreateVirtualMachineConfig(
             string adminUsername,
             string adminPassword,
             string size,
-            ResourceConfig<AvailabilitySet> availabilitySet)
+            ResourceConfig<AvailabilitySet> availabilitySet,
+            IEnumerable<int> dataDisks)
             => Strategy.CreateResourceConfig(
                 resourceGroup: resourceGroup,
                 name: name,
@@ -70,7 +72,9 @@ public static ResourceConfig<VirtualMachine> CreateVirtualMachineConfig(
                     },
                     StorageProfile = new StorageProfile
                     {
-                        ImageReference = imageAndOsType?.Image
+                        ImageReference = imageAndOsType?.Image,
+                        DataDisks = DataDiskStrategy.CreateDataDisks(
+                            imageAndOsType?.DataDiskLuns, dataDisks)
                     },
                     AvailabilitySet = engine.GetReference(availabilitySet)
                 });
@@ -82,7 +86,8 @@ public static ResourceConfig<VirtualMachine> CreateVirtualMachineConfig(
             OperatingSystemTypes osType,
             ResourceConfig<Disk> disk,
             string size,
-            ResourceConfig<AvailabilitySet> availabilitySet)
+            ResourceConfig<AvailabilitySet> availabilitySet,
+            IEnumerable<int> dataDisks)
             => Strategy.CreateResourceConfig(
                 resourceGroup: resourceGroup,
                 name: name,
@@ -107,7 +112,8 @@ public static ResourceConfig<VirtualMachine> CreateVirtualMachineConfig(
                             CreateOption = DiskCreateOptionTypes.Attach,
                             OsType = osType,
                             ManagedDisk = engine.GetReference(disk, StorageAccountTypes.PremiumLRS),
-                        }
+                        },
+                        DataDisks = DataDiskStrategy.CreateDataDisks(null, dataDisks)
                     },
                     AvailabilitySet = engine.GetReference(availabilitySet)
                 });

src/ResourceManager/Compute/Commands.Compute/Strategies/StrategyCmdlet.cs

@@ -1,5 +1,4 @@
 using Microsoft.Azure.Commands.Common.Strategies;
-using System;
 using System.Threading;
 using System.Threading.Tasks;
 

src/ResourceManager/Compute/Commands.Compute/VirtualMachine/Operation/NewAzureVMCommand.cs

@@ -28,7 +28,6 @@
 using Microsoft.Azure.Management.Compute;
 using Microsoft.Azure.Management.Compute.Models;
 using Microsoft.Azure.Management.Internal.Network.Version2017_10_01;
-using Microsoft.Azure.Management.Internal.Network.Version2017_10_01.Models;
 using Microsoft.Azure.Management.Internal.Resources;
 using Microsoft.Azure.Management.Internal.Resources.Models;
 using Microsoft.Azure.Management.Storage;
@@ -211,6 +210,10 @@ public class NewAzureVMCommand : VirtualMachineBaseCmdlet
         [Parameter(Mandatory = false, HelpMessage = "Run cmdlet in the background")]
         public SwitchParameter AsJob { get; set; }
 
+        [Parameter(ParameterSetName = SimpleParameterSet, Mandatory = false)]
+        [Parameter(ParameterSetName = DiskFileParameterSet, Mandatory = false)]
+        public int[] DataDiskSizeInGb { get; set; }
+
         public override void ExecuteCmdlet()
         {
             switch (ParameterSetName)
@@ -293,7 +296,8 @@ public async Task<ResourceConfig<VirtualMachine>> CreateConfigAsync()
                         adminPassword:
                             new NetworkCredential(string.Empty, _cmdlet.Credential.Password).Password,
                         size: _cmdlet.Size,
-                        availabilitySet: availabilitySet);
+                        availabilitySet: availabilitySet,
+                        dataDisks: _cmdlet.DataDiskSizeInGb);
                 }
                 else
                 {
@@ -307,7 +311,8 @@ public async Task<ResourceConfig<VirtualMachine>> CreateConfigAsync()
                         osType: ImageAndOsType.OsType,
                         disk: disk,
                         size: _cmdlet.Size,
-                        availabilitySet: availabilitySet);
+                        availabilitySet: availabilitySet,
+                        dataDisks: _cmdlet.DataDiskSizeInGb);
                 }
             }
         }
@@ -341,6 +346,7 @@ async Task StrategyExecuteCmdletAsync(IAsyncCmdlet asyncCmdlet)
                 }
                 parameters.ImageAndOsType = new ImageAndOsType(
                     Linux ? OperatingSystemTypes.Linux : OperatingSystemTypes.Windows,
+                    null,
                     null);
                 var storageClient = AzureSession.Instance.ClientFactory.CreateArmClient<StorageManagementClient>(
                     DefaultProfile.DefaultContext,