[RDBMS] az postgres flexible-server update: BUG FIX, Updating geo data encryption properties (#30948)

* add to ignore

* Fix bug when updating cmk geo

* Revert "add to ignore"

This reverts commit ee1192c.

* Update validators and check that primary and backup id are not the same

* lint fixes

* small changes

Author: nasc17

src/azure-cli/azure/cli/command_modules/rdbms/_flexible_server_util.py

@@ -443,34 +443,42 @@ def _is_resource_name(resource):
 
 
 def build_identity_and_data_encryption(db_engine, byok_identity=None, backup_byok_identity=None,
-                                       byok_key=None, backup_byok_key=None):
+                                       byok_key=None, backup_byok_key=None, instance=None):
     identity, data_encryption = None, None
 
-    if byok_identity and byok_key:
-        identities = {byok_identity: {}}
+    primary_user_assigned_identity_id = byok_identity
+    primary_key_uri = byok_key
+    geo_backup_user_assigned_identity_id = backup_byok_identity
+    geo_backup_key_uri = backup_byok_key
+    if (instance is not None) and (byok_identity is None) and (backup_byok_identity is not None):
+        primary_user_assigned_identity_id = instance.data_encryption.primary_user_assigned_identity_id
+        primary_key_uri = instance.data_encryption.primary_key_uri
 
-        if backup_byok_identity:
-            identities[backup_byok_identity] = {}
+    if primary_user_assigned_identity_id and primary_key_uri:
+        identities = {primary_user_assigned_identity_id: {}}
+
+        if geo_backup_user_assigned_identity_id:
+            identities[geo_backup_user_assigned_identity_id] = {}
 
         if db_engine == 'mysql':
             identity = mysql_flexibleservers.models.Identity(user_assigned_identities=identities,
                                                              type="UserAssigned")
 
             data_encryption = mysql_flexibleservers.models.DataEncryption(
-                primary_user_assigned_identity_id=byok_identity,
-                primary_key_uri=byok_key,
-                geo_backup_user_assigned_identity_id=backup_byok_identity,
-                geo_backup_key_uri=backup_byok_key,
+                primary_user_assigned_identity_id=primary_user_assigned_identity_id,
+                primary_key_uri=primary_key_uri,
+                geo_backup_user_assigned_identity_id=geo_backup_user_assigned_identity_id,
+                geo_backup_key_uri=geo_backup_key_uri,
                 type="AzureKeyVault")
         else:
             identity = postgresql_flexibleservers.models.UserAssignedIdentity(user_assigned_identities=identities,
                                                                               type="UserAssigned")
 
             data_encryption = postgresql_flexibleservers.models.DataEncryption(
-                primary_user_assigned_identity_id=byok_identity,
-                primary_key_uri=byok_key,
-                geo_backup_user_assigned_identity_id=backup_byok_identity,
-                geo_backup_key_uri=backup_byok_key,
+                primary_user_assigned_identity_id=primary_user_assigned_identity_id,
+                primary_key_uri=primary_key_uri,
+                geo_backup_user_assigned_identity_id=geo_backup_user_assigned_identity_id,
+                geo_backup_key_uri=geo_backup_key_uri,
                 type="AzureKeyVault")
 
     return identity, data_encryption

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_custom_postgres.py

@@ -396,7 +396,10 @@ def flexible_server_update_custom_func(cmd, client, instance,
 
     identity, data_encryption = build_identity_and_data_encryption(db_engine='postgres',
                                                                    byok_identity=byok_identity,
-                                                                   byok_key=byok_key)
+                                                                   byok_key=byok_key,
+                                                                   backup_byok_identity=backup_byok_identity,
+                                                                   backup_byok_key=backup_byok_key,
+                                                                   instance=instance)
 
     auth_config = instance.auth_config
     administrator_login = instance.administrator_login if instance.administrator_login else None

src/azure-cli/azure/cli/command_modules/rdbms/validators.py

@@ -556,6 +556,11 @@ def pg_byok_validator(byok_identity, byok_key, backup_byok_identity=None, backup
         raise ArgumentUsageError("User assigned identity and keyvault key need to be provided together. "
                                  "Please provide --backup-identity and --backup-key together.")
 
+    if bool(byok_identity is not None) and bool(backup_byok_identity is not None) and \
+       byok_identity.lower() == backup_byok_identity.lower():
+        raise ArgumentUsageError("Primary user assigned identity and backup identity cannot be same. "
+                                 "Please provide different identities for --identity and --backup-identity.")
+
     if (instance is not None) and \
        not (instance.data_encryption and instance.data_encryption.type == 'AzureKeyVault') and \
        (byok_key or backup_byok_key):
@@ -570,6 +575,14 @@ def pg_byok_validator(byok_identity, byok_key, backup_byok_identity=None, backup
         if instance is None and (bool(byok_key is not None) ^ bool(backup_byok_key is not None)):
             raise ArgumentUsageError("Please provide both primary as well as geo-back user assigned identity "
                                      "and keyvault key to enable Data encryption for geo-redundant backup.")
+        if instance is not None and (bool(byok_identity is None) ^ bool(backup_byok_identity is None)):
+            primary_user_assigned_identity_id = byok_identity if byok_identity else \
+                instance.data_encryption.primary_user_assigned_identity_id
+            geo_backup_user_assigned_identity_id = backup_byok_identity if backup_byok_identity else \
+                instance.data_encryption.geo_backup_user_assigned_identity_id
+            if primary_user_assigned_identity_id.lower() == geo_backup_user_assigned_identity_id.lower():
+                raise ArgumentUsageError("Primary user assigned identity and backup identity cannot be same. "
+                                         "Please provide different identities for --identity and --backup-identity.")
 
 
 def _network_arg_validator(subnet, public_access):