[SQL] Bugfix - Fix identity cmdlets (#21086)

* [SQL] Fix identity for SetAzSqlDatabase cmdlet

* Fix changelog

* Fix logic

---------

Co-authored-by: Vinit Dinesh Parekh <[email protected]>
Co-authored-by: Beisi Zhou <[email protected]>

Author: 3 people

src/Sql/Sql/ChangeLog.md

@@ -18,6 +18,7 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
+* Fixed identity assignment in `Set-AzSqlDatabase` cmdlet
 * Added new parameters to `New-AzSqlDatabase`, `Get-AzSqlDatabase`, `Set-AzSqlDatabase`, `New-AzSqlDatabaseCopy`, `New-AzSqlDatabaseSecondary` cmdlets
    - AssignIdentity
    - EncryptionProtector

src/Sql/Sql/Common/DatabaseIdentityAndKeysHelper.cs

@@ -30,9 +30,8 @@ public class DatabaseIdentityAndKeysHelper
         /// </summary>
         /// <param name="assignIdentityIsPresent">Flag to check if AssignIdentity flag is used in the cmdlet.</param>
         /// <param name="userAssignedIdentities">User assigned identities</param>
-        /// <param name="existingIdentity">Existing identity on database</param>
         /// <returns>Database Identity</returns>
-        public static DatabaseIdentity GetDatabaseIdentity(bool assignIdentityIsPresent, string[] userAssignedIdentities, Management.Sql.Models.DatabaseIdentity existingIdentity)
+        public static DatabaseIdentity GetDatabaseIdentity(bool assignIdentityIsPresent, string[] userAssignedIdentities)
         {
             DatabaseIdentity identityResult = null;
             if (!assignIdentityIsPresent)
@@ -47,34 +46,22 @@ public static DatabaseIdentity GetDatabaseIdentity(bool assignIdentityIsPresent,
 
             Dictionary<string, DatabaseUserIdentity> identityDict = new Dictionary<string, DatabaseUserIdentity>();
 
-            // Update an existing database identity
+            // Create identity on database
             //
-            if (existingIdentity != null && userAssignedIdentities.Any() &&
-                existingIdentity.UserAssignedIdentities != null)
+            foreach (string identity in userAssignedIdentities)
             {
-                foreach (string identity in userAssignedIdentities)
-                {
-                    existingIdentity.UserAssignedIdentities.Add(identity, new DatabaseUserIdentity());
-                }
-
-                identityResult = existingIdentity;
-            }
-            else
-            {
-                // Create identity on database
-                //
-                foreach (string identity in userAssignedIdentities)
+                if (!identityDict.ContainsKey(identity))
                 {
                     identityDict.Add(identity, new DatabaseUserIdentity());
                 }
-
-                identityResult = new DatabaseIdentity()
-                {
-                    Type = DatabaseIdentityType.UserAssigned,
-                    UserAssignedIdentities = identityDict
-                };
             }
 
+            identityResult = new DatabaseIdentity()
+            {
+                Type = DatabaseIdentityType.UserAssigned,
+                UserAssignedIdentities = identityDict
+            };
+
             return identityResult;
         }
 

src/Sql/Sql/Database Backup/Cmdlet/RestoreAzureRMSqlDatabase.cs

@@ -411,7 +411,7 @@ protected override AzureSqlDatabaseModel GetEntity()
                 Tags = TagsConversionHelper.CreateTagDictionary(Tag, validate: true),
                 ZoneRedundant = this.IsParameterBound(p => p.ZoneRedundant) ? ZoneRedundant.ToBool() : (bool?)null,
                 HighAvailabilityReplicaCount = HAReplicaCount,
-                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId, null),
+                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId),
                 Keys = DatabaseIdentityAndKeysHelper.GetDatabaseKeysDictionary(this.KeyList),
                 EncryptionProtector = this.EncryptionProtector,
                 FederatedClientId = this.FederatedClientId,

src/Sql/Sql/Database/Cmdlet/NewAzureSqlDatabase.cs

@@ -371,7 +371,7 @@ protected override AzureSqlDatabaseCreateOrUpdateModel ApplyUserInputToModel(Azu
                 MaintenanceConfigurationId = MaintenanceConfigurationId,
                 EnableLedger = this.IsParameterBound(p => p.EnableLedger) ? EnableLedger.ToBool() : (bool?)null,
                 PreferredEnclaveType = this.PreferredEnclaveType,
-                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId, null),
+                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId),
                 Keys = DatabaseIdentityAndKeysHelper.GetDatabaseKeysDictionary(this.KeyList),
                 EncryptionProtector = this.EncryptionProtector,
                 FederatedClientId = this.FederatedClientId

src/Sql/Sql/Database/Cmdlet/SetAzureSqlDatabase.cs

@@ -352,7 +352,7 @@ protected override IEnumerable<AzureSqlDatabaseModel> ApplyUserInputToModel(IEnu
                 SecondaryType = SecondaryType,
                 MaintenanceConfigurationId = MaintenanceConfigurationId,
                 PreferredEnclaveType = PreferredEnclaveType,
-                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId, model.FirstOrDefault().Identity),
+                Identity = DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId),
                 EncryptionProtector = this.EncryptionProtector ?? model.FirstOrDefault().EncryptionProtector,
                 FederatedClientId = this.FederatedClientId ?? model.FirstOrDefault().FederatedClientId,
             };

src/Sql/Sql/Replication/Cmdlet/NewAzureSqlDatabaseCopy.cs

@@ -262,7 +262,7 @@ protected override IEnumerable<AzureSqlDatabaseCopyModel> ApplyUserInputToModel(
                 LicenseType = LicenseType, // note: default license type is LicenseIncluded
                 RequestedBackupStorageRedundancy = BackupStorageRedundancy,
                 ZoneRedundant = this.IsParameterBound(p => p.ZoneRedundant) ? ZoneRedundant.ToBool() : (bool?)null,
-                Identity = Common.DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId, null),
+                Identity = Common.DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId),
                 Keys = Common.DatabaseIdentityAndKeysHelper.GetDatabaseKeysDictionary(this.KeyList),
                 EncryptionProtector = this.EncryptionProtector,
                 FederatedClientId = this.FederatedClientId

src/Sql/Sql/Replication/Cmdlet/NewAzureSqlDatabaseSecondary.cs

@@ -279,7 +279,7 @@ protected override IEnumerable<AzureReplicationLinkModel> ApplyUserInputToModel(
                 SecondaryType = SecondaryType,
                 HighAvailabilityReplicaCount = this.IsParameterBound(p => p.HighAvailabilityReplicaCount) ? HighAvailabilityReplicaCount : (int?)null,
                 ZoneRedundant = this.IsParameterBound(p => p.ZoneRedundant) ? ZoneRedundant.ToBool() : (bool?)null,
-                Identity = Common.DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId, null),
+                Identity = Common.DatabaseIdentityAndKeysHelper.GetDatabaseIdentity(this.AssignIdentity.IsPresent, this.UserAssignedIdentityId),
                 Keys = Common.DatabaseIdentityAndKeysHelper.GetDatabaseKeysDictionary(this.KeyList),
                 EncryptionProtector = this.EncryptionProtector,
                 FederatedClientId = this.FederatedClientId