{ServiceConnector}: supporting customized keys in passwordless and fix admin user name

src/serviceconnector-passwordless/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+0.3.9
+++++++
+* Support `--customized-keys` and make some improvements.
+
 0.3.8
 ++++++
 * Make some improvements.
@@ -48,4 +52,4 @@ Release History
 
 0.1.0
 ++++++
-* Initial release.
+* Initial release.

src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py

@@ -177,6 +177,8 @@ class TargetHandler:
     user_object_id = ""
     aad_username = ""
 
+    admin_username = ""
+
     identity_name = ""
     identity_client_id = ""
     identity_object_id = ""
@@ -307,8 +309,9 @@ def set_user_admin(self, user_object_id, **kwargs):
                 logger.warning(
                     'Unable to check if current user is AAD admin. Please confirm current user as AAD admin manually.')
                 return
-        is_admin = any(ad.get('sid') == user_object_id for ad in admins)
-        if is_admin:
+        admin_info = next((ad for ad in admins if ad.get('sid') == user_object_id), None)
+        if admin_info:
+            self.admin_username = admin_info.get('login')
             return
 
         logger.warning('Set current user as DB Server AAD Administrators.')
@@ -326,6 +329,7 @@ def set_user_admin(self, user_object_id, **kwargs):
                 self.resource_group, self.server, self.subscription, mysql_identity_id))
         run_cli_cmd('az mysql flexible-server ad-admin create -g {} -s {} --subscription {} -u {} -i {} --identity {}'.format(
             self.resource_group, self.server, self.subscription, self.login_username, user_object_id, mysql_identity_id))
+        self.admin_username = self.login_username
 
     def create_aad_user(self):
         query_list = self.get_create_query()
@@ -441,7 +445,7 @@ def get_connection_string(self):
         return {
             'host': self.server + self.endpoint,
             'database': self.dbname,
-            'user': self.login_username,
+            'user': self.admin_username,
             'password': password,
             'ssl': {"fake_flag_to_enable_tls": True},
             'autocommit': True
@@ -503,12 +507,13 @@ def set_user_admin(self, user_object_id, **kwargs):
                 logger.warning(
                     'Unable to check if current user is AAD admin. Please confirm current user as AAD admin manually.')
                 return
-        is_admin = any(ad.get('sid') == user_object_id for ad in admins)
-        if not is_admin:
+        admin_info = next((ad for ad in admins if ad.get('sid') == user_object_id), None)
+        if not admin_info:
             logger.warning('Setting current user as database server AAD admin:'
                            ' user=%s object id=%s', self.login_username, user_object_id)
-            run_cli_cmd('az sql server ad-admin create -g {} --server-name {} --display-name {} --object-id {} --subscription {}'.format(
-                self.resource_group, self.server, self.login_username, user_object_id, self.subscription)).get('objectId')
+            admin_info = run_cli_cmd('az sql server ad-admin create -g {} --server-name {} --display-name {} --object-id {} --subscription {}'.format(
+                self.resource_group, self.server, self.login_username, user_object_id, self.subscription))
+        self.admin_username = admin_info.get('login', self.login_username)
 
     def create_aad_user(self):
 
@@ -699,12 +704,12 @@ def set_user_admin(self, user_object_id, **kwargs):
                 logger.warning(
                     'Unable to check if current user is AAD admin. Please confirm current user as AAD admin manually.')
                 return
-        is_admin = any(user_object_id in u.get("objectId", "") for u in admins)
-        if is_admin:
-            return
-        logger.warning('Set current user as DB Server AAD Administrators.')
-        run_cli_cmd('az postgres flexible-server ad-admin create -u {} -i {} -g {} -s {} --subscription {} -t {}'.format(
-            self.login_username, user_object_id, self.resource_group, self.db_server, self.subscription, self.login_usertype))
+        admin_info = next((ad for ad in admins if ad.get('objectId', "") == user_object_id), None)
+        if not admin_info:
+            logger.warning('Set current user as DB Server AAD Administrators.')
+            admin_info = run_cli_cmd('az postgres flexible-server ad-admin create -u {} -i {} -g {} -s {} --subscription {} -t {}'.format(
+                self.login_username, user_object_id, self.resource_group, self.db_server, self.subscription, self.login_usertype))
+        self.admin_username = admin_info.get('principalName', self.login_username)
 
     def create_aad_user(self):
         query_list = self.get_create_query()
@@ -823,7 +828,7 @@ def get_connection_string(self):
 
         # extension functions require the extension to be available, which is the case for postgres (default) database.
         conn_string = "host={} user={} dbname=postgres password={} sslmode=require".format(
-            self.host, self.login_username, password)
+            self.host, self.admin_username, password)
         return conn_string
 
     def get_create_query(self):
@@ -866,7 +871,6 @@ def set_user_admin(self, user_object_id, **kwargs):
         sub = self.subscription
         rg = self.resource_group
         server = self.db_server
-        is_admin = True
 
         # pylint: disable=not-an-iterable
         admins = run_cli_cmd(
@@ -882,12 +886,13 @@ def set_user_admin(self, user_object_id, **kwargs):
                 logger.warning(
                     'Unable to check if current user is AAD admin. Please confirm current user as AAD admin manually.')
                 return
-        is_admin = any(ad.get('sid') == user_object_id for ad in admins)
-        if not is_admin:
+        admin_info = next((ad for ad in admins if ad.get('sid') == user_object_id), None)
+        if not admin_info:
             logger.warning('Setting current user as database server AAD admin:'
                            ' user=%s object id=%s', self.login_username, user_object_id)
-            run_cli_cmd('az postgres server ad-admin create -g {} --server-name {} --display-name {} --object-id {}'
-                        ' --subscription {}'.format(rg, server, self.login_username, user_object_id, sub)).get('objectId')
+            admin_info = run_cli_cmd('az postgres server ad-admin create -g {} --server-name {} --display-name {} --object-id {}'
+                                     ' --subscription {}'.format(rg, server, self.login_username, user_object_id, sub))
+        self.admin_username = admin_info.get('login', self.login_username)
 
     def set_target_firewall(self, is_add, ip_name, start_ip=None, end_ip=None):
         sub = self.subscription
@@ -928,7 +933,7 @@ def get_connection_string(self):
 
         # extension functions require the extension to be available, which is the case for postgres (default) database.
         conn_string = "host={} user={} dbname={} password={} sslmode=require".format(
-            self.host, self.login_username + '@' + self.db_server, self.dbname, password)
+            self.host, self.admin_username + '@' + self.db_server, self.dbname, password)
         return conn_string
 
     def get_create_query(self):

src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_params.py

@@ -14,7 +14,8 @@
     add_vnet_block,
     add_connection_string_argument,
     add_secret_store_argument,
-    add_local_connection_block
+    add_local_connection_block,
+    add_customized_keys_argument
 )
 from azure.cli.command_modules.serviceconnector._validators import (
     get_default_object_id_of_current_user
@@ -63,6 +64,7 @@ def load_arguments(self, _):
             add_secret_store_argument(c)
             add_vnet_block(c, target)
             add_local_connection_block(c)
+            add_customized_keys_argument(c)
             c.argument('yes', arg_type=yes_arg_type)
 
     for source in SOURCE_RESOURCES_PARAMS:
@@ -77,4 +79,5 @@ def load_arguments(self, _):
                 add_secret_store_argument(c)
                 add_vnet_block(c, target)
                 add_connection_string_argument(c, source, target)
+                add_customized_keys_argument(c)
                 c.argument('yes', arg_type=yes_arg_type)

src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py

@@ -4,5 +4,5 @@
 # --------------------------------------------------------------------------------------------
 
 
-VERSION = '0.3.8'
+VERSION = '0.3.9'
 NAME = 'serviceconnector-passwordless'

src/serviceconnector-passwordless/azext_serviceconnector_passwordless/custom.py

@@ -15,6 +15,7 @@ def connection_create_ext(cmd, client,  # pylint: disable=too-many-locals,too-ma
                           service_endpoint=None,
                           private_endpoint=None,
                           store_in_connection_string=False,
+                          customized_keys=None,
                           new_addon=False, no_wait=False,
                           yes=False,
                           # Resource.KubernetesCluster
@@ -44,6 +45,7 @@ def connection_create_ext(cmd, client,  # pylint: disable=too-many-locals,too-ma
                                   spring, app, deployment,
                                   server, database,
                                   enable_mi_for_db_linker=get_enable_mi_for_db_linker_func(yes),
+                                  customized_keys=customized_keys,
                                   **kwargs)
 
 
@@ -56,6 +58,7 @@ def local_connection_create_ext(cmd, client,  # pylint: disable=too-many-locals,
                                 secret_auth_info=None, secret_auth_info_auto=None,
                                 user_account_auth_info=None,                      # new auth info
                                 service_principal_auth_info_secret=None,
+                                customized_keys=None,
                                 no_wait=False,
                                 yes=False,
                                 # Resource.*Postgres, Resource.*Sql*
@@ -77,4 +80,5 @@ def local_connection_create_ext(cmd, client,  # pylint: disable=too-many-locals,
                                         # Resource.*Postgres, Resource.*Sql*
                                         server, database,
                                         enable_mi_for_db_linker=get_enable_mi_for_db_linker_func(yes),
+                                        customized_keys=customized_keys,
                                         **kwargs)

src/serviceconnector-passwordless/setup.py

@@ -15,7 +15,7 @@
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
 
-VERSION = '0.3.8'
+VERSION = '0.3.9'
 try:
     from azext_serviceconnector_passwordless.config import VERSION
 except ImportError: