Update with UMI scenarios

src/Sql/Sql/Auditing/Services/AuditingEndpointsCommunicator.cs

@@ -226,7 +226,7 @@ public DiagnosticSettingsResource UpdateDiagnosticSettings(DiagnosticSettingsRes
             if (server.Identity == null ||
                 server.Identity.Type != ResourceIdentityType.SystemAssigned.ToString())
             {
-                server.Identity = ResourceIdentityHelper.GetIdentityObjectFromType(true, false, null);
+                server.Identity = ResourceIdentityHelper.GetIdentityObjectFromType(true, "SystemAssigned", null, null);
                 server = GetCurrentSqlClient().Servers.CreateOrUpdate(resourceGroupName, serverName, server);
             }
 

src/Sql/Sql/Common/ResourceIdentityHelper.cs

@@ -15,28 +15,60 @@
 using Microsoft.Azure.Management.Sql.Models;
 using System.Collections.Generic;
 using System.Linq;
+using System.Management.Automation;
 using System.Runtime.CompilerServices;
 
 namespace Microsoft.Azure.Commands.Sql.Common
 {
     public enum ResourceIdentityType
     {
         SystemAssigned,
+        SystemAssignedUserAssigned,
         UserAssigned,
         None
     }
 
     public class ResourceIdentityHelper
     {
-        public static Management.Sql.Models.ResourceIdentity GetIdentityObjectFromType(bool assignIdentityIsPresent, bool userAssignedIdentityIsPresent, List<string> userAssignedIdentities)
+        public static Management.Sql.Models.ResourceIdentity GetIdentityObjectFromType(bool assignIdentityIsPresent, string resourceIdentityType, List<string> userAssignedIdentities, Management.Sql.Models.ResourceIdentity existingResourceIdentity)
         {
             Management.Sql.Models.ResourceIdentity identityResult = null;
 
-            if (assignIdentityIsPresent && userAssignedIdentityIsPresent)
+            // If the user passes in IdentityType as None, then irrespective of previous config, we set the IdentityType to be None.
+            //
+            if (resourceIdentityType != null && resourceIdentityType.Equals(ResourceIdentityType.None.ToString()))
+            {
+                identityResult = new Management.Sql.Models.ResourceIdentity()
+                {
+                    Type = ResourceIdentityType.None.ToString()
+                };
+
+                return identityResult;
+            }
+
+            if (resourceIdentityType != null && assignIdentityIsPresent && resourceIdentityType.Equals(ResourceIdentityType.SystemAssignedUserAssigned.ToString()))
             {
                 Dictionary<string, UserIdentity> umiDict = new Dictionary<string, UserIdentity>();
 
-                if (userAssignedIdentities != null && userAssignedIdentities.Any())
+                if (userAssignedIdentities == null)
+                {
+                    throw new PSArgumentNullException("The list of user assigned identity ids needs to be passed if the IdentityType is UserAssigned or SystemAssignedUserAssigned");
+                }
+
+                if (existingResourceIdentity != null && userAssignedIdentities.Any()
+                    && existingResourceIdentity.UserAssignedIdentities != null)
+                {
+                    foreach (string identity in userAssignedIdentities)
+                    {
+                        existingResourceIdentity.UserAssignedIdentities.Add(identity, new UserIdentity());
+                    }
+
+                    identityResult = new Management.Sql.Models.ResourceIdentity()
+                    {
+                        Type = ResourceIdentityType.SystemAssignedUserAssigned.ToString()
+                    };
+                }
+                else if (userAssignedIdentities.Any())
                 {
                     foreach (string identity in userAssignedIdentities)
                     {
@@ -45,34 +77,70 @@ public static Management.Sql.Models.ResourceIdentity GetIdentityObjectFromType(b
 
                     identityResult = new Management.Sql.Models.ResourceIdentity()
                     {
-                        Type = ResourceIdentityType.UserAssigned.ToString(),
+                        Type = ResourceIdentityType.SystemAssignedUserAssigned.ToString(),
                         UserAssignedIdentities = umiDict
                     };
                 }
-                else
+            }
+            else if (resourceIdentityType != null && assignIdentityIsPresent && resourceIdentityType.Equals(ResourceIdentityType.UserAssigned.ToString()))
+            {
+                Dictionary<string, UserIdentity> umiDict = new Dictionary<string, UserIdentity>();
+
+                if (userAssignedIdentities == null)
                 {
+                    throw new PSArgumentNullException("The list of user assigned identity ids needs to be passed if the IdentityType is UserAssigned or SystemAssignedUserAssigned");
+                }
+
+                if (existingResourceIdentity != null && userAssignedIdentities.Any()
+                    && existingResourceIdentity.UserAssignedIdentities != null)
+                {
+                    foreach (string identity in userAssignedIdentities)
+                    {
+                        existingResourceIdentity.UserAssignedIdentities.Add(identity, new UserIdentity());
+                    }
+
                     identityResult = new Management.Sql.Models.ResourceIdentity()
                     {
-                        Type = ResourceIdentityType.SystemAssigned.ToString()
+                        Type = ResourceIdentityType.UserAssigned.ToString()
+                    };
+                }
+                else if (userAssignedIdentities.Any())
+                {
+                    foreach (string identity in userAssignedIdentities)
+                    {
+                        umiDict.Add(identity, new UserIdentity());
+                    }
+
+                    identityResult = new Management.Sql.Models.ResourceIdentity()
+                    {
+                        Type = ResourceIdentityType.UserAssigned.ToString(),
+                        UserAssignedIdentities = umiDict
                     };
-                }    
+                }
             }
             else if (assignIdentityIsPresent)
             {
-                identityResult = new Management.Sql.Models.ResourceIdentity()
+                if (existingResourceIdentity != null)
                 {
-                    Type = ResourceIdentityType.SystemAssigned.ToString()
-                };
+                    identityResult = existingResourceIdentity;
+                    identityResult.Type = ResourceIdentityType.SystemAssigned.ToString();
+                }
+                else
+                {
+                    identityResult = new Management.Sql.Models.ResourceIdentity()
+                    {
+                        Type = ResourceIdentityType.SystemAssigned.ToString()
+                    };
+                }       
             }
-            else if (!assignIdentityIsPresent && !userAssignedIdentityIsPresent)
+
+            if (!assignIdentityIsPresent && existingResourceIdentity != null && existingResourceIdentity.PrincipalId != null)
             {
-                identityResult = new Management.Sql.Models.ResourceIdentity()
-                {
-                    Type = ResourceIdentityType.None.ToString()
-                };
+                identityResult = existingResourceIdentity;
             }
 
             return identityResult;
+
         }
     }
 }

src/Sql/Sql/ManagedInstance/Cmdlet/NewAzureSqlManagedInstance.cs

@@ -342,11 +342,15 @@ public class NewAzureSqlManagedInstance : ManagedInstanceCmdletBase
         /// </summary>
         [Parameter(Mandatory = false,
             HelpMessage = "List of user assigned identities")]
-        public List<string> UserAssignedIdentity { get; set; }
+        public List<string> UserAssignedIdentityId { get; set; }
 
+        // <summary>
+        /// Type of identity to be assigned to the server..
+        /// </summary>
         [Parameter(Mandatory = false,
-            HelpMessage = "Generate and assign an Azure Active Directory User Assigned Identity for this server for use with key management services like Azure KeyVault.")]
-        public SwitchParameter AssignUserAssignIdentity { get; set; }
+            HelpMessage = "Type of Identity to be used. Possible values are SystemAsssigned, UserAssigned, SystemAssignedUserAssigned and None.")]
+        [PSArgumentCompleter("SystemAssigned", "UserAssigned", "SystemAssignedUserAssigned", "None")]
+        public string IdentityType { get; set; }
 
         /// <summary>
         /// Gets or sets whether or not to run this cmdlet in the background as a job
@@ -521,7 +525,7 @@ public override void ExecuteCmdlet()
                 AdministratorPassword = (this.AdministratorCredential != null) ? this.AdministratorCredential.Password : null,
                 AdministratorLogin = (this.AdministratorCredential != null) ? this.AdministratorCredential.UserName : null,
                 Tags = TagsConversionHelper.CreateTagDictionary(Tag, validate: true),
-                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity, this.AssignUserAssignIdentity, UserAssignedIdentity),
+                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity.IsPresent, this.IdentityType ?? null, UserAssignedIdentityId, null),
                 LicenseType = this.LicenseType,
                 // `-StorageSizeInGB 0` as a parameter to this cmdlet means "use default".
                 // For non-MI database, we can just pass in 0 and the server will treat 0 as default.

src/Sql/Sql/ManagedInstance/Cmdlet/SetAzureSqlManagedInstance.cs

@@ -223,11 +223,15 @@ public class SetAzureSqlManagedInstance : ManagedInstanceCmdletBase
         /// </summary>
         [Parameter(Mandatory = false,
             HelpMessage = "List of user assigned identities")]
-        public List<string> UserAssignedIdentity { get; set; }
+        public List<string> UserAssignedIdentityId { get; set; }
 
+        // <summary>
+        /// List of user assigned identities.
+        /// </summary>
         [Parameter(Mandatory = false,
-            HelpMessage = "Generate and assign an Azure Active Directory User Assigned Identity for this server for use with key management services like Azure KeyVault.")]
-        public SwitchParameter AssignUserAssignIdentity { get; set; }
+            HelpMessage = "Type of Identity to be used. Possible values are SystemAsssigned, UserAssigned, SystemAssignedUserAssigned and None.")]
+        [PSArgumentCompleter("SystemAssigned", "UserAssigned", "SystemAssignedUserAssigned", "None")]
+        public string IdentityType { get; set; }
 
         /// <summary>
         /// Gets or sets whether or not to run this cmdlet in the background as a job
@@ -312,7 +316,7 @@ protected override IEnumerable<AzureSqlManagedInstanceModel> ApplyUserInputToMod
                 PublicDataEndpointEnabled = this.PublicDataEndpointEnabled,
                 ProxyOverride = this.ProxyOverride,
                 Tags = TagsConversionHelper.CreateTagDictionary(Tag, validate: true),
-                Identity = model.FirstOrDefault().Identity ?? ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity, this.AssignUserAssignIdentity, UserAssignedIdentity),
+                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity.IsPresent, this.IdentityType ?? null, UserAssignedIdentityId, GetEntity().FirstOrDefault().Identity),
                 InstancePoolName = this.InstancePoolName,
                 MinimalTlsVersion = this.MinimalTlsVersion,
                 MaintenanceConfigurationId = this.MaintenanceConfigurationId,

src/Sql/Sql/Server/Cmdlet/NewAzureSqlServer.cs

@@ -108,11 +108,15 @@ public class NewAzureSqlServer : AzureSqlServerCmdletBase
         /// </summary>
         [Parameter(Mandatory = false,
             HelpMessage = "List of user assigned identities")]
-        public List<string> UserAssignedIdentity { get; set; }
+        public List<string> UserAssignedIdentityId { get; set; }
 
+        // <summary>
+        /// Type of identity to be assigned to the server..
+        /// </summary>
         [Parameter(Mandatory = false,
-            HelpMessage = "Generate and assign an Azure Active Directory User Assigned Identity for this server for use with key management services like Azure KeyVault.")]
-        public SwitchParameter AssignUserAssignIdentity { get; set; }
+            HelpMessage = "Type of Identity to be used. Possible values are SystemAsssigned, UserAssigned, SystemAssignedUserAssigned and None.")]
+        [PSArgumentCompleter("SystemAssigned", "UserAssigned", "SystemAssignedUserAssigned", "None")]
+        public string IdentityType { get; set; }
 
         /// <summary>
         /// Gets or sets whether or not to run this cmdlet in the background as a job
@@ -209,7 +213,7 @@ public override void ExecuteCmdlet()
                 SqlAdministratorPassword = (this.SqlAdministratorCredentials != null) ? this.SqlAdministratorCredentials.Password : null,
                 SqlAdministratorLogin = (this.SqlAdministratorCredentials != null) ? this.SqlAdministratorCredentials.UserName : null,
                 Tags = TagsConversionHelper.CreateTagDictionary(Tags, validate: true),
-                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity, this.AssignUserAssignIdentity, UserAssignedIdentity),
+                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity.IsPresent, this.IdentityType ?? null, UserAssignedIdentityId, null),
                 MinimalTlsVersion = this.MinimalTlsVersion,
                 PublicNetworkAccess = this.PublicNetworkAccess,
                 PrimaryUserAssignedIdentityId = this.PrimaryUserAssignedIdentityId,

src/Sql/Sql/Server/Cmdlet/SetAzureSqlServer.cs

@@ -102,11 +102,15 @@ public class SetAzureSqlServer : AzureSqlServerCmdletBase
         /// </summary>
         [Parameter(Mandatory = false,
             HelpMessage = "List of user assigned identities")]
-        public List<string> UserAssignedIdentity { get; set; }
+        public List<string> UserAssignedIdentityId { get; set; }
 
+        // <summary>
+        /// Type of identity to be assigned to the server..
+        /// </summary>
         [Parameter(Mandatory = false,
-            HelpMessage = "Generate and assign an Azure Active Directory User Assigned Identity for this server for use with key management services like Azure KeyVault.")]
-        public SwitchParameter AssignUserAssignIdentity { get; set; }
+            HelpMessage = "Type of Identity to be used. Possible values are SystemAsssigned, UserAssigned, SystemAssignedUserAssigned and None.")]
+        [PSArgumentCompleter("SystemAssigned", "UserAssigned", "SystemAssignedUserAssigned", "None")]
+        public string IdentityType { get; set; }
 
         /// <summary>
         /// Defines whether it is ok to skip the requesting of rule removal confirmation
@@ -145,10 +149,10 @@ public class SetAzureSqlServer : AzureSqlServerCmdletBase
                 Tags = TagsConversionHelper.ReadOrFetchTags(this, model.FirstOrDefault().Tags),
                 ServerVersion = this.ServerVersion,
                 Location = model.FirstOrDefault().Location,
-                Identity = model.FirstOrDefault().Identity ?? ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity, this.AssignUserAssignIdentity, UserAssignedIdentity),
+                Identity = ResourceIdentityHelper.GetIdentityObjectFromType(this.AssignIdentity.IsPresent, this.IdentityType ?? null, UserAssignedIdentityId, GetEntity().FirstOrDefault().Identity),
                 PublicNetworkAccess = this.PublicNetworkAccess,
                 MinimalTlsVersion = this.MinimalTlsVersion,
-                PrimaryUserAssignedIdentityId = this.PrimaryUserAssignedIdentityId,
+                PrimaryUserAssignedIdentityId = model.FirstOrDefault().PrimaryUserAssignedIdentityId ?? this.PrimaryUserAssignedIdentityId,
                 KeyId = this.KeyId
             });
             return updateData;

src/Sql/Sql/TransparentDataEncryption/Cmdlet/SetAzureSqlServerTransparentDataEncryptionProtector.cs

@@ -55,7 +55,7 @@ public class SetAzureSqlServerTransparentDataEncryptionProtector : AzureSqlServe
         ValueFromPipelineByPropertyName = true,
             HelpMessage = "The Key Auto Rotation status")]
         [ValidateNotNullOrEmpty]
-        public SwitchParameter AutoRotationEnabled { get; set; }
+        public bool? AutoRotationEnabled { get; set; }
 
         /// <summary>
         /// Defines whether it is ok to skip the requesting of setting Transparent Data Encryption protector confirmation