{ServiceConnector}: Fix user-agent missing and getting postgres db ex…

…ception (#7342) * fix error * update * update * update * update * update user agent * lint * updated
Azure · Mar 6, 2024 · d57168d · d57168d
1 parent 0f8ff5f
commit d57168d
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 11 deletions.
diff --git a/src/serviceconnector-passwordless/HISTORY.rst b/src/serviceconnector-passwordless/HISTORY.rst
@@ -2,6 +2,10 @@
 
 Release History
 ===============
+1.0.2
+++++++
+* Bypass error of `az postgres flexible-server db show`
+
 1.0.1
 ++++++
 * Make some improvements to fix non-json output issue

diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_client_factory.py
@@ -2,22 +2,18 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+import os
 
 
 # pylint: disable=consider-using-f-string
 def cf_connection_cl(cli_ctx, *_):
     from azure.mgmt.servicelinker import ServiceLinkerManagementClient
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client, get_az_user_agent
-    from azure.core.pipeline import policies
+    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     from .config import NAME, VERSION
-
-    user_agent_policy = policies.UserAgentPolicy(
-        user_agent=get_az_user_agent())
-    user_agent_policy.add_user_agent(
-        "CliExtension/{}({})".format(NAME, VERSION))
+    os.environ['AZURE_HTTP_USER_AGENT'] = (os.environ.get('AZURE_HTTP_USER_AGENT')
+                                           or '') + " CliExtension/{}({})".format(NAME, VERSION)
     return get_mgmt_service_client(cli_ctx, ServiceLinkerManagementClient,
-                                   subscription_bound=False, api_version="2022-11-01-preview",
-                                   user_agent_policy=user_agent_policy)
+                                   subscription_bound=False, api_version="2022-11-01-preview")
 
 
 def cf_linker(cli_ctx, *_):

diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/_credential_free.py
@@ -641,6 +641,18 @@ def __init__(self, cmd, target_id, target_type, auth_info, connection_name, skip
 
     def check_db_existence(self):
         try:
+            # `az postgres flexible-server db show -d postgres` will throw exception
+            if self.dbname == "postgres":
+                server_info = run_cli_cmd(
+                    'az postgres flexible-server show -n {} -g {} --subscription {}'.format(
+                        self.db_server, self.resource_group, self.subscription))
+                if server_info is None:
+                    e = ResourceNotFoundError(
+                        "No server found for '{}'".format(self.db_server))
+                    telemetry.set_exception(e, "No-Server")
+                    raise e
+                else:
+                    return
             db_info = run_cli_cmd(
                 'az postgres flexible-server db show --server-name {} --database-name {} -g {} --subscription {}'.format(
                     self.db_server, self.dbname, self.resource_group, self.subscription))

diff --git a/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py b/src/serviceconnector-passwordless/azext_serviceconnector_passwordless/config.py
@@ -4,5 +4,5 @@
 # --------------------------------------------------------------------------------------------
 
 
-VERSION = '1.0.1'
+VERSION = '1.0.2'
 NAME = 'serviceconnector-passwordless'
diff --git a/src/serviceconnector-passwordless/setup.py b/src/serviceconnector-passwordless/setup.py
@@ -15,7 +15,7 @@
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
 
-VERSION = '1.0.1'
+VERSION = '1.0.2'
 try:
     from azext_serviceconnector_passwordless.config import VERSION
 except ImportError: