[AKS] support NodeTaints and AksCustomHeaders (Azure#20319)

* [AKS] support NodeTaints

* [AKS] support AksCustomHeader

src/Aks/Aks.Test/ScenarioTests/KubernetesTests.cs

@@ -93,5 +93,12 @@ public void TestNodeLabelsAndTags()
         {
             TestRunner.RunTestScript("Test-NodeLabels-Tags");
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestNodeTaints()
+        {
+            TestRunner.RunTestScript("Test-NodeTaints");
+        }
     }
 }

src/Aks/Aks.Test/ScenarioTests/KubernetesTests.ps1

@@ -487,4 +487,54 @@ function Test-NodeLabels-Tags {
     finally {
         Remove-AzResourceGroup -Name $resourceGroupName -Force
     }
+}
+
+function Test-NodeTaints {
+    # Setup
+    $resourceGroupName = Get-RandomResourceGroupName
+    $kubeClusterName = Get-RandomClusterName
+    $location = 'eastus'
+    $nodeVmSize = "Standard_D2_v2"
+
+    try {
+        New-AzResourceGroup -Name $resourceGroupName -Location $location
+
+        # create aks cluster with default nodepool
+        New-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName -NodeVmSize $nodeVmSize -NodeCount 1
+        $cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
+        Assert-AreEqual 1 $cluster.AgentPoolProfiles.Count
+        Assert-AreEqual 0 $cluster.AgentPoolProfiles[0].NodeTaints.Count
+        $pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
+        Assert-AreEqual 1 $pools.Count
+        Assert-AreEqual 0 $pools.NodeTaints.Count
+
+        # create a 2nd nodepool
+        $nodetains = @("sku=gpu:NoSchedule")
+        New-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName -Name "pool2" -Count 1 -NodeTaint $nodetains
+        $cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
+        Assert-AreEqual 2 $cluster.AgentPoolProfiles.Count
+        Assert-AreEqual 0 ($cluster.AgentPoolProfiles | where {$_.Name -eq "default"}).NodeTaints.Count
+        Assert-AreEqualArray $nodetains ($cluster.AgentPoolProfiles | where {$_.Name -eq "pool2"}).NodeTaints
+        $pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
+        Assert-AreEqual 2 $pools.Count
+        Assert-AreEqual 0 ($pools | where {$_.Name -eq "default"}).NodeTaints.Count
+        Assert-AreEqualArray $nodetains ($pools | where {$_.Name -eq "pool2"}).NodeTaints
+
+        # update the 2nd nodepool
+        $nodetains2 = @("CriticalAddonsOnly=true:NoSchedule")
+        Update-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName -Name "pool2" -NodeTaint $nodetains2
+        $cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
+        Assert-AreEqual 2 $cluster.AgentPoolProfiles.Count
+        Assert-AreEqual 0 ($cluster.AgentPoolProfiles | where {$_.Name -eq "default"}).NodeTaints.Count
+        Assert-AreEqualArray $nodetains2 ($cluster.AgentPoolProfiles | where {$_.Name -eq "pool2"}).NodeTaints
+        $pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
+        Assert-AreEqual 2 $pools.Count
+        Assert-AreEqual 0 ($pools | where {$_.Name -eq "default"}).NodeTaints.Count
+        Assert-AreEqualArray $nodetains2 ($pools | where {$_.Name -eq "pool2"}).NodeTaints
+
+        $cluster | Remove-AzAksCluster -Force
+    }
+    finally {
+        Remove-AzResourceGroup -Name $resourceGroupName -Force
+    }
 }

src/Aks/Aks/ChangeLog.md

@@ -24,6 +24,8 @@
 * Added property `-NodeImageVersion` for the output of `Get-AzAksNodePool`[#19893]
 * Added parameter `-NodePoolLabel` for `Set-AzAksCluster`, `-NodeLabel` for `New-AzAksNodePool` and `Update-AzAksNodePool`
 * Added parameter `-NodePoolTag` for `New-AzAksCluster` and `Set-AzAksCluster`, `-Tag` for `New-AzAksNodePool` and `Update-AzAksNodePool`
+* Added parameter `-NodeTaint` for `New-AzAksNodePool` and `Update-AzAksNodePool`
+* Added parameter `-AksCustomHeader` for `Set-AzAksCluster`, `New-AzAksNodePool` and `Update-AzAksNodePool`
 
 ## Version 5.0.1
 * Upgraded AutoMapper to Microsoft.Azure.PowerShell.AutoMapper 6.2.2 with fix [#18721]

src/Aks/Aks/Commands/CreateOrUpdateKubeBase.cs

@@ -39,6 +39,7 @@
 using Microsoft.Azure.Commands.Common.MSGraph.Version1_0.Applications;
 using Microsoft.Azure.Commands.Common.MSGraph.Version1_0;
 using ResourceIdentityType = Microsoft.Azure.Management.ContainerService.Models.ResourceIdentityType;
+using Microsoft.Azure.Commands.Aks.Commands;
 
 namespace Microsoft.Azure.Commands.Aks
 {
@@ -190,6 +191,9 @@ public abstract class CreateOrUpdateKubeBase : KubeCmdletBase
         [Parameter(Mandatory = false, HelpMessage = "Alternative CA cert to use for connecting to proxy servers.")]
         public string HttpProxyConfigTrustedCa { get; set; }
 
+        [Parameter(Mandatory = false, HelpMessage = "Aks custom headers used for building Kubernetes network.")]
+        public Hashtable AksCustomHeader { get; set; }
+
         protected void BeforeBuildNewCluster()
         {
             if (!string.IsNullOrEmpty(ResourceGroupName) && string.IsNullOrEmpty(Location))
@@ -269,7 +273,7 @@ protected string GetSshKey(string sshKeyOrFile)
         protected AcsServicePrincipal EnsureServicePrincipal(string spId = null, string clientSecret = null)
         {
             //If user specifies service principal, just use it directly and no need to save to disk
-            if(!string.IsNullOrEmpty(spId) && !string.IsNullOrEmpty(clientSecret))
+            if (!string.IsNullOrEmpty(spId) && !string.IsNullOrEmpty(clientSecret))
             {
                 return new AcsServicePrincipal()
                 {
@@ -289,7 +293,7 @@ protected AcsServicePrincipal EnsureServicePrincipal(string spId = null, string
                 if (clientSecret == null)
                 {
                     clientSecret = RandomBase64String(16);
-                }              
+                }
 
                 acsServicePrincipal = BuildServicePrincipal(Name, clientSecret);
                 WriteVerbose(Resources.CreatedANewServicePrincipalAndAssignedTheContributorRole);
@@ -365,7 +369,7 @@ protected void AddAcrRoleAssignment(string acrName, string acrParameterName, Acs
                 var acrObjects = RmClient.Resources.List(acrQuery);
                 acrResourceId = acrObjects.First().Id;
             }
-            catch(Exception)
+            catch (Exception)
             {
                 throw new AzPSArgumentException(
                     string.Format(Resources.CouldNotFindSpecifiedAcr, acrName),
@@ -389,12 +393,12 @@ protected void AddAcrRoleAssignment(string acrName, string acrParameterName, Acs
                     var servicePrincipal = GraphClient.FilterServicePrincipals(oDataQuery).First();
                     spObjectId = servicePrincipal.Id;
                 }
-                catch(Exception ex)
+                catch (Exception ex)
                 {
                     throw new AzPSInvalidOperationException(
                         string.Format(Resources.CouldNotFindObjectIdForServicePrincipal, acsServicePrincipal.SpId),
                         ex,
-                        string.Format(Resources.CouldNotFindObjectIdForServicePrincipal,"*"));
+                        string.Format(Resources.CouldNotFindObjectIdForServicePrincipal, "*"));
                 }
             }
             var success = RetryAction(() =>
@@ -472,7 +476,7 @@ protected bool RetryAction(Action action, string actionName = null)
         protected AcsServicePrincipal LoadServicePrincipal()
         {
             var config = LoadServicePrincipals();
-            if(config?.ContainsKey(DefaultContext.Subscription.Id) == true)
+            if (config?.ContainsKey(DefaultContext.Subscription.Id) == true)
             {
                 return config[DefaultContext.Subscription.Id];
             }
@@ -518,7 +522,7 @@ protected string DefaultDnsPrefix()
             var subPart = string.Join("", DefaultContext.Subscription.Id.Take(4));
             return $"{namePart}{subPart}";
         }
-    
+
         protected ManagedClusterLoadBalancerProfile CreateOrUpdateLoadBalancerProfile(ManagedClusterLoadBalancerProfile loadBalancerProfile)
         {
             if ((this.IsParameterBound(c => c.LoadBalancerManagedOutboundIpCount) ||
@@ -666,5 +670,18 @@ protected ManagedCluster SetIdentity(ManagedCluster cluster)
 
             return cluster;
         }
+
+        private protected ManagedCluster CreateOrUpdate(string resourceGroupName, string resourceName, ManagedCluster parameters)
+        {
+            if (this.IsParameterBound(c => c.AksCustomHeader))
+            {
+                Dictionary<string, List<string>> customHeaders = Utilities.HashtableToDictionary(AksCustomHeader);
+                return Client.ManagedClusters.CreateOrUpdateWithHttpMessagesAsync(resourceGroupName, resourceName, parameters, customHeaders).GetAwaiter().GetResult().Body;
+            }
+            else
+            {
+                return Client.ManagedClusters.CreateOrUpdate(resourceGroupName, resourceName, parameters);
+            }
+        }
     }
 }

src/Aks/Aks/Commands/NewAzureRmAks.cs

@@ -13,7 +13,6 @@
 // ----------------------------------------------------------------------------------
 
 using System;
-using System.Collections;
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Diagnostics;
@@ -28,7 +27,6 @@
 using Microsoft.Azure.Commands.Common.Exceptions;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
 using Microsoft.Azure.Commands.ResourceManager.Common.Tags;
-using Microsoft.Azure.Management.ContainerService;
 using Microsoft.Azure.Management.ContainerService.Models;
 using Microsoft.Rest;
 using Microsoft.WindowsAzure.Commands.Common;
@@ -113,10 +111,6 @@ public class NewAzureRmAks : CreateOrUpdateKubeBase
         [Parameter(Mandatory = false, HelpMessage = "Docker bridge cidr used for building Kubernetes network.")]
         public string DockerBridgeCidr { get; set; }
 
-        [Parameter(Mandatory = false, HelpMessage = "Aks custom headers used for building Kubernetes network.")]
-
-        public Hashtable AksCustomHeader { get; set; }
-
         [Parameter(Mandatory = false, HelpMessage = "The load balancer sku for the managed cluster.")]
         [PSArgumentCompleter("basic", "standard")]
         public string LoadBalancerSku { get; set; }
@@ -157,26 +151,8 @@ public override void ExecuteCmdlet()
                 var managedCluster = BuildNewCluster();
                 try
                 {
-                    ManagedCluster cluster;
-                    if (this.IsParameterBound(c => c.AksCustomHeader))
-                    {
-                        Dictionary<string, List<string>> customHeaders = new Dictionary<string, List<string>>();
-                        foreach (var key in AksCustomHeader.Keys)
-                        {
-                            List<string> values = new List<string>();
-                            foreach (var value in (object[])AksCustomHeader[key])
-                            {
-                                values.Add(value.ToString());
-                            }
-                            customHeaders.Add(key.ToString(), values);
-                        }
-
-                        cluster = Client.ManagedClusters.CreateOrUpdateWithHttpMessagesAsync(ResourceGroupName, Name, managedCluster, customHeaders).GetAwaiter().GetResult().Body;
-                    }
-                    else
-                    {
-                        cluster = Client.ManagedClusters.CreateOrUpdate(ResourceGroupName, Name, managedCluster);
-                    }
+
+                    var cluster = this.CreateOrUpdate(ResourceGroupName, Name, managedCluster);
                     var psObj = PSMapper.Instance.Map<PSKubernetesCluster>(cluster);
 
                     if (this.IsParameterBound(c => c.AcrNameToAttach))

src/Aks/Aks/Commands/NewAzureRmAksNodePool.cs

@@ -15,7 +15,6 @@
 using System;
 using System.Collections.Generic;
 using System.Management.Automation;
-
 using Microsoft.Azure.Commands.Aks.Models;
 using Microsoft.Azure.Commands.Aks.Properties;
 using Microsoft.Azure.Commands.Common.Exceptions;
@@ -113,7 +112,7 @@ public override void ExecuteCmdlet()
                     ClusterName = ClusterObject.Name;
                 }
                 var agentPool = GetAgentPool();
-                var pool = Client.AgentPools.CreateOrUpdate(ResourceGroupName, ClusterName, Name, agentPool);
+                var pool = this.CreateOrUpdate(ResourceGroupName, ClusterName, Name, agentPool);
                 var psPool = PSMapper.Instance.Map<PSNodePool>(pool);
                 WriteObject(psPool);
             };
@@ -211,6 +210,10 @@ private AgentPool GetAgentPool()
                     agentPool.Tags.Add(key.ToString(), Tag[key].ToString());
                 }
             }
+            if (this.IsParameterBound(c => c.NodeTaint))
+            {
+                agentPool.NodeTaints = NodeTaint;
+            }
 
             return agentPool;
         }

src/Aks/Aks/Commands/NewOrUpdateAgentPoolBase.cs

@@ -12,7 +12,11 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.Aks.Commands;
 using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
+using Microsoft.Azure.Management.ContainerService;
+using Microsoft.Azure.Management.ContainerService.Models;
+using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using Newtonsoft.Json;
 using System.Collections;
 using System.Collections.Generic;
@@ -47,5 +51,24 @@ public class NewOrUpdateAgentPoolBase : KubeCmdletBase
 
         [Parameter(Mandatory = false, HelpMessage = "The tags to be persisted on the agent pool virtual machine scale set.")]
         public Hashtable Tag { get; set; }
+
+        [Parameter(Mandatory = false, HelpMessage = "The node taints added to new nodes during node pool create and scale")]
+        public string[] NodeTaint { get; set; }
+
+        [Parameter(Mandatory = false, HelpMessage = "Aks custom headers")]
+        public Hashtable AksCustomHeader { get; set; }
+
+        private protected AgentPool CreateOrUpdate(string resourceGroupName, string resourceName, string agentPoolName, AgentPool parameters)
+        {
+            if (this.IsParameterBound(c => c.AksCustomHeader))
+            {
+                Dictionary<string, List<string>> customHeaders = Utilities.HashtableToDictionary(AksCustomHeader);
+                return Client.AgentPools.CreateOrUpdateWithHttpMessagesAsync(resourceGroupName, resourceName, agentPoolName, parameters, customHeaders).GetAwaiter().GetResult().Body;
+            }
+            else
+            {
+                return Client.AgentPools.CreateOrUpdate(resourceGroupName, resourceName, agentPoolName, parameters);
+            }
+        }
     }
 }

src/Aks/Aks/Commands/SetAzureRmAks.cs

@@ -16,7 +16,6 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Management.Automation;
-
 using Microsoft.Azure.Commands.Aks.Models;
 using Microsoft.Azure.Commands.Aks.Properties;
 using Microsoft.Azure.Commands.Common.Exceptions;
@@ -31,11 +30,8 @@
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
 using Microsoft.Rest.Azure.OData;
 using Microsoft.WindowsAzure.Commands.Common;
-using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
 
-using ResourceIdentityType = Microsoft.Azure.Management.ContainerService.Models.ResourceIdentityType;
-
 namespace Microsoft.Azure.Commands.Aks
 {
     [Cmdlet("Set", ResourceManager.Common.AzureRMConstants.AzureRMPrefix + "AksCluster", DefaultParameterSetName = DefaultParamSet, SupportsShouldProcess = true)]
@@ -401,7 +397,7 @@ public override void ExecuteCmdlet()
                     }
                     SetIdentity(cluster);
 
-                    var kubeCluster = Client.ManagedClusters.CreateOrUpdate(ResourceGroupName, Name, cluster);
+                    var kubeCluster = this.CreateOrUpdate(ResourceGroupName, Name, cluster);
 
                     if (this.IsParameterBound(c => c.DiskEncryptionSetID))
                     {

src/Aks/Aks/Commands/UpdateAzureRmAksNodePool.cs

@@ -180,8 +180,12 @@ public override void ExecuteCmdlet()
                             pool.Tags.Add(key.ToString(), Tag[key].ToString());
                         }
                     }
+                    if (this.IsParameterBound(c => c.NodeTaint))
+                    {
+                        pool.NodeTaints = NodeTaint;
+                    }
 
-                    var updatedPool = Client.AgentPools.CreateOrUpdate(ResourceGroupName, ClusterName, Name, pool);
+                    var updatedPool = this.CreateOrUpdate(ResourceGroupName, ClusterName, Name, pool);
                     WriteObject(PSMapper.Instance.Map<PSNodePool>(updatedPool));
                 });
             }

src/Aks/Aks/Commands/Utilities.cs

@@ -13,7 +13,8 @@
 // ----------------------------------------------------------------------------------
 
 using System;
-
+using System.Collections;
+using System.Collections.Generic;
 using Microsoft.Azure.Commands.Common.Exceptions;
 
 namespace Microsoft.Azure.Commands.Aks.Commands
@@ -32,5 +33,24 @@ public static string GetParentResourceName(string parentResource, string paramet
 
             return items.Length > 0 ? items[items.Length - 1] : null;
         }
+
+        public static Dictionary<string, List<string>> HashtableToDictionary(Hashtable table)
+        {
+            Dictionary<string, List<string>> dict = new Dictionary<string, List<string>>();
+            if (table == null)
+            {
+                return dict;
+            }
+            foreach (var key in table.Keys)
+            {
+                List<string> values = new List<string>();
+                foreach (var value in (object[])table[key])
+                {
+                    values.Add(value.ToString());
+                }
+                dict.Add(key.ToString(), values);
+            }
+            return dict;
+        }
     }
 }